#!/usr/bin/env python3
"""
Merges fresh Personio employee data with an existing headcount snapshot.

Usage:
    python3 build_snapshot.py --existing data/headcount_snapshot.json --new-data /tmp/fresh_employees.json --output data/headcount_snapshot.json

The --new-data file should be a JSON array of employee objects from the Personio MCP workflow.
Fields like first_name, last_name, and email are automatically stripped (PII).

Optionally pass --team-mapping to also rebuild the team mapping file.
"""

import json
import argparse
from datetime import datetime, timezone
from collections import defaultdict


def strip_pii(employee: dict) -> dict:
    """Remove personally identifiable information, keep only operational fields."""
    return {
        "personio_id": employee.get("personio_id") or employee.get("id"),
        "status": employee.get("status"),
        "team_id": employee.get("team_id", ""),
        "department_id": employee.get("department_id", ""),
        "office_id": employee.get("office_id", ""),
        "start_date": employee.get("start_date"),
        "end_date": employee.get("end_date"),
    }


def merge_employees(existing: list, new_data: list) -> list:
    """
    Merge new employee data into existing snapshot.
    - Existing records are updated if the same personio_id appears in new_data
    - New IDs are added
    - IDs only in existing are preserved (historical records)
    """
    index = {e["personio_id"]: e for e in existing}

    updated = 0
    added = 0
    for emp in new_data:
        clean = strip_pii(emp)
        pid = clean["personio_id"]
        if pid in index:
            index[pid] = clean
            updated += 1
        else:
            index[pid] = clean
            added += 1

    print(f"Merge result: {updated} updated, {added} added, {len(index)} total")
    return list(index.values())


def build_summary(employees: list) -> dict:
    """Count employees by status."""
    statuses = {}
    for e in employees:
        s = e.get("status", "UNKNOWN")
        statuses[s] = statuses.get(s, 0) + 1
    return statuses


def rebuild_team_mapping(employees: list, existing_mapping_path: str = None) -> dict:
    """Rebuild team mapping counts from employee data, using CSV as source of truth for names."""
    import csv
    import os

    team_names = {}
    # Try to load names from CSV (single source of truth)
    if existing_mapping_path:
        csv_path = os.path.join(os.path.dirname(existing_mapping_path), "team-mapping.csv")
        try:
            with open(csv_path) as f:
                for row in csv.DictReader(f):
                    team_names[row['ID'].strip()] = row['Name'].strip()
            print(f"Loaded {len(team_names)} team names from CSV")
        except FileNotFoundError:
            pass

    # Fallback: load from existing JSON if CSV not available
    if not team_names and existing_mapping_path:
        try:
            with open(existing_mapping_path) as f:
                existing = json.load(f)
            for t in existing.get("teams", []):
                team_names[t["team_id"]] = t["team_name"]
        except FileNotFoundError:
            pass

    # Count per team
    team_stats = defaultdict(lambda: {"active": 0, "total": 0})
    for e in employees:
        tid = e.get("team_id", "")
        if tid:
            team_stats[tid]["total"] += 1
            if e.get("status") == "ACTIVE":
                team_stats[tid]["active"] += 1

    teams_list = []
    for tid, stats in sorted(team_stats.items(), key=lambda x: team_names.get(x[0], x[0])):
        teams_list.append({
            "team_id": tid,
            "team_name": team_names.get(tid, f"Unknown ({tid})"),
            "active_headcount": stats["active"],
            "total_records": stats["total"],
        })

    return {"total_teams": len(teams_list), "teams": teams_list}


def main():
    parser = argparse.ArgumentParser(description="Build/update headcount snapshot")
    parser.add_argument("--existing", help="Path to existing snapshot JSON (optional)")
    parser.add_argument("--new-data", required=True, help="Path to fresh employee data JSON array")
    parser.add_argument("--output", required=True, help="Output path for updated snapshot")
    parser.add_argument("--team-mapping", help="Path to team_mapping.json to update (optional)")
    args = parser.parse_args()

    # Load new data
    with open(args.new_data) as f:
        new_employees = json.load(f)
    print(f"Loaded {len(new_employees)} new employee records")

    # Load existing snapshot if available
    existing_employees = []
    if args.existing:
        try:
            with open(args.existing) as f:
                existing = json.load(f)
            existing_employees = existing.get("employees", [])
            print(f"Loaded {len(existing_employees)} existing records from snapshot")
        except FileNotFoundError:
            print("No existing snapshot found — creating fresh")

    # Merge
    merged = merge_employees(existing_employees, new_employees)

    # Build snapshot
    summary = build_summary(merged)
    snapshot = {
        "last_refreshed": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "total_records": len(merged),
        "summary": summary,
        "employees": merged,
    }

    with open(args.output, "w") as f:
        json.dump(snapshot, f, indent=2)

    print(f"\nSnapshot saved to {args.output}")
    print(f"Total records: {len(merged)}")
    print(f"Summary: {json.dumps(summary, indent=2)}")
    active = summary.get("ACTIVE", 0)
    print(f"\n=> Current active headcount: {active}")

    # Optionally rebuild team mapping
    if args.team_mapping:
        team_map = rebuild_team_mapping(merged, args.team_mapping)
        with open(args.team_mapping, "w") as f:
            json.dump(team_map, f, indent=2)
        print(f"\nTeam mapping updated: {team_map['total_teams']} teams")


if __name__ == "__main__":
    main()