import csv
import os
from django.core.management.base import BaseCommand

class Command(BaseCommand):
    help = 'Flags staff as heads based on whether they appear in the Manager column'

    def add_arguments(self, parser):
        parser.add_argument('input_csv', type=str, help='Path to original CSV')
        parser.add_argument('output_csv', type=str, help='Path to save the new CSV')

    def handle(self, *args, **options):
        input_path = options['input_csv']
        output_path = options['output_csv']

        if not os.path.exists(input_path):
            self.stdout.write(self.style.ERROR(f"File {input_path} not found."))
            return

        rows = []
        manager_ids_set = set()

        # Pass 1: Collect all IDs from the Manager column
        with open(input_path, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            fieldnames = reader.fieldnames
            
            for row in reader:
                rows.append(row)
                manager_val = row.get('Manager', '')
                
                if manager_val and '-' in manager_val:
                    # Split by '-' and take the first part (the ID)
                    m_id = manager_val.split('-')[0].strip()
                    if m_id:
                        manager_ids_set.add(m_id)

        # Pass 2: Flag rows as 'is_head' if their Staff ID is in the set
        for row in rows:
            staff_id = row.get('Staff ID', '').strip()
            if staff_id in manager_ids_set:
                row['is_head'] = 1
            else:
                row['is_head'] = 0

        # Pass 3: Write the new CSV
        new_fieldnames = fieldnames + ['is_head']
        with open(output_path, 'w', encoding='utf-8', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=new_fieldnames)
            writer.writeheader()
            writer.writerows(rows)

        self.stdout.write(self.style.SUCCESS(f"Processed {len(rows)} rows. Found {len(manager_ids_set)} unique heads."))
        self.stdout.write(self.style.SUCCESS(f"Saved to: {output_path}"))