#!/usr/bin/env python3
"""
G-pre2 INCREMENTAL: Add new atomic controls to Master Controls without rebuild.

Unlike gpre2_master_controls.py which DELETEs and rebuilds the entire
master_controls table, this script is non-destructive:
  - Existing master_controls stay untouched (same UUIDs, same MC-IDs)
  - For each object_group that gained new atomic controls:
    * If MC exists: append new members + update total_controls/phase_counts
    * If MC missing AND group now has >= min_phases: create new MC + all members

Usage:
  python3 /app/scripts/gpre2_master_controls_incremental.py --since 2026-05-18T02:53:00+00:00
  python3 /app/scripts/gpre2_master_controls_incremental.py --since 2026-05-18T02:53:00+00:00 --dry-run
  python3 /app/scripts/gpre2_master_controls_incremental.py --since 2026-05-18T02:53:00+00:00 --min-phases 2
"""

import argparse
import json
import logging
import os
from collections import defaultdict
from datetime import datetime

from sqlalchemy import create_engine, text

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("gpre2_incremental")

DB_URL = os.getenv("DATABASE_URL", "postgresql://breakpilot:breakpilot123@postgres:5432/breakpilot_db")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--since", required=True, help="ISO datetime — only consider atomics created at/after this")
    parser.add_argument("--min-phases", type=int, default=2, help="Min distinct phases to form a new MC (default 2)")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    since_dt = datetime.fromisoformat(args.since.replace("Z", "+00:00"))
    logger.info("Incremental run since %s, min_phases=%d, dry_run=%s",
                since_dt.isoformat(), args.min_phases, args.dry_run)

    engine = create_engine(DB_URL, connect_args={"options": "-c search_path=compliance,public"})

    # Step 1: object → group_id reverse index
    object_to_group = {}
    with engine.connect() as c:
        groups = c.execute(text("SELECT group_id, canonical_name, members FROM object_groups")).fetchall()
    for gid, canonical, members_json in groups:
        members = json.loads(members_json) if isinstance(members_json, str) else members_json
        for member in members:
            object_to_group[member] = (gid, canonical)
    logger.info("Reverse index: %d objects → %d groups", len(object_to_group), len(groups))

    # Step 2: Load ALL atomics with merge_group_hint (we need full picture)
    with engine.connect() as c:
        all_rows = c.execute(text("""
            SELECT id, control_id,
                   generation_metadata->>'merge_group_hint' AS hint,
                   title,
                   created_at
            FROM canonical_controls
            WHERE generation_metadata->>'merge_group_hint' IS NOT NULL
              AND generation_metadata->>'merge_group_hint' != ''
              AND release_state NOT IN ('deprecated', 'rejected', 'duplicate')
        """)).fetchall()
    logger.info("Loaded %d atomic controls total", len(all_rows))

    # Step 3: Build group_phases (gid → phase → [(uuid, control_id, action, title, is_new)])
    from services.control_dedup import normalize_object
    group_phases: dict[int, dict[str, list]] = defaultdict(lambda: defaultdict(list))
    group_names: dict[int, str] = {}
    new_atomic_count = 0
    new_groups_touched: set[int] = set()
    unmatched = 0

    for uuid, control_id, hint, title, created_at in all_rows:
        parts = hint.split(":", 2)
        if len(parts) < 2:
            continue
        action = parts[0]
        obj = parts[1]
        phase = parts[2] if len(parts) > 2 else "implementation"
        normed = normalize_object(obj)
        if normed in object_to_group:
            gid, canonical = object_to_group[normed]
        elif obj in object_to_group:
            gid, canonical = object_to_group[obj]
        else:
            unmatched += 1
            continue
        is_new = created_at >= since_dt
        group_phases[gid][phase].append((str(uuid), control_id, action, title, is_new))
        group_names[gid] = canonical
        if is_new:
            new_atomic_count += 1
            new_groups_touched.add(gid)

    logger.info("Total: %d new atomics across %d object_groups (%d unmatched)",
                new_atomic_count, len(new_groups_touched), unmatched)

    if not new_groups_touched:
        logger.info("Nothing to do — no new atomics matched to any object_group.")
        return

    # Step 4: For each touched object_group, decide action
    stats = {
        "groups_examined": len(new_groups_touched),
        "mcs_existing_updated": 0,
        "mcs_new_created": 0,
        "members_inserted": 0,
        "members_skipped_existing": 0,
        "groups_skipped_below_min_phases": 0,
        "groups_skipped_no_member_change": 0,
    }

    # Load existing master_controls index: master_control_id → uuid
    with engine.connect() as c:
        mc_index = {row[1]: (str(row[0]), row[2]) for row in c.execute(text(
            "SELECT id, master_control_id, total_controls FROM master_controls"
        )).fetchall()}
    logger.info("Existing master_controls: %d", len(mc_index))

    # Load existing members for touched MCs (avoid duplicate inserts)
    touched_mc_ids = ["MC-%d" % gid for gid in new_groups_touched]
    existing_members: dict[str, set[str]] = defaultdict(set)
    with engine.connect() as c:
        for mc_id_str in touched_mc_ids:
            mc_uuid_info = mc_index.get(mc_id_str)
            if not mc_uuid_info:
                continue
            mc_uuid = mc_uuid_info[0]
            for row in c.execute(text(
                "SELECT control_uuid FROM master_control_members WHERE master_control_uuid = CAST(:u AS uuid)"
            ), {"u": mc_uuid}).fetchall():
                existing_members[mc_id_str].add(str(row[0]))

    # Build INSERT/UPDATE plans
    inserts_new_mcs = []
    inserts_members = []
    updates_mcs = []

    PHASE_ORDER = {
        "scope": 0, "definition": 1, "governance": 1, "design": 2,
        "implementation": 3, "configuration": 3, "operation": 4, "training": 4,
        "monitoring": 5, "testing": 6, "review": 7, "assessment": 8,
        "remediation": 8, "validation": 9, "reporting": 10, "evidence": 11,
    }

    for gid in new_groups_touched:
        mc_id_str = "MC-%d" % gid
        phases = group_phases[gid]
        canonical = group_names[gid]
        all_phases = sorted(phases.keys(), key=lambda p: PHASE_ORDER.get(p, 99))
        phase_counts = {p: len(ctrls) for p, ctrls in phases.items()}
        total = sum(phase_counts.values())

        existing_mc = mc_index.get(mc_id_str)

        if existing_mc:
            # MC exists — append only NEW atomics that aren't already members
            mc_uuid = existing_mc[0]
            existing_set = existing_members[mc_id_str]
            added_for_this_mc = 0
            for phase, controls in phases.items():
                for ctrl_uuid, ctrl_id, action, title, is_new in controls:
                    if ctrl_uuid in existing_set:
                        stats["members_skipped_existing"] += 1
                        continue
                    inserts_members.append({
                        "mc_uuid": mc_uuid, "control_uuid": ctrl_uuid,
                        "phase": phase, "action": action,
                    })
                    stats["members_inserted"] += 1
                    added_for_this_mc += 1
            if added_for_this_mc > 0:
                updates_mcs.append({
                    "mc_uuid": mc_uuid,
                    "phases_covered": json.dumps(all_phases),
                    "phase_control_count": json.dumps(phase_counts),
                    "total_controls": total,
                })
                stats["mcs_existing_updated"] += 1
            else:
                stats["groups_skipped_no_member_change"] += 1
        else:
            # MC missing — create only if group now meets min_phases threshold
            if len(phases) < args.min_phases:
                stats["groups_skipped_below_min_phases"] += 1
                continue
            inserts_new_mcs.append({
                "master_control_id": mc_id_str,
                "object_group_id": gid,
                "canonical_name": canonical,
                "phases_covered": json.dumps(all_phases),
                "phase_control_count": json.dumps(phase_counts),
                "total_controls": total,
                "_members": [
                    {"control_uuid": c[0], "phase": p, "action": c[2]}
                    for p, ctrls in phases.items() for c in ctrls
                ],
            })
            stats["mcs_new_created"] += 1

    logger.info("Plan summary: %s", stats)

    if args.dry_run:
        logger.info("DRY RUN — no writes")
        # Show first few examples
        if inserts_new_mcs:
            logger.info("Sample NEW MCs (up to 5):")
            for mc in inserts_new_mcs[:5]:
                logger.info("  %s: %s — total=%d, phases=%s",
                            mc["master_control_id"], mc["canonical_name"],
                            mc["total_controls"], mc["phases_covered"])
        if updates_mcs:
            logger.info("Updates to existing MCs: %d", len(updates_mcs))
        return

    # Step 5: WRITE — strictly INSERT/UPDATE, no DELETE
    with engine.begin() as c:
        c.execute(text("SET search_path TO compliance, public"))

        # 5a: Insert new MCs + their members
        for mc in inserts_new_mcs:
            new_uuid_row = c.execute(text("""
                INSERT INTO master_controls
                    (master_control_id, object_group_id, canonical_name,
                     phases_covered, phase_control_count, total_controls)
                VALUES (:master_control_id, :object_group_id, :canonical_name,
                        CAST(:phases_covered AS jsonb), CAST(:phase_control_count AS jsonb),
                        :total_controls)
                RETURNING id
            """), {k: v for k, v in mc.items() if k != "_members"}).fetchone()
            new_mc_uuid = str(new_uuid_row[0])
            for mem in mc["_members"]:
                c.execute(text("""
                    INSERT INTO master_control_members
                        (master_control_uuid, control_uuid, phase, action)
                    VALUES (CAST(:mc_uuid AS uuid), CAST(:control_uuid AS uuid), :phase, :action)
                """), {"mc_uuid": new_mc_uuid, **mem})

        # 5b: Append new members to existing MCs
        for mem in inserts_members:
            c.execute(text("""
                INSERT INTO master_control_members
                    (master_control_uuid, control_uuid, phase, action)
                VALUES (CAST(:mc_uuid AS uuid), CAST(:control_uuid AS uuid), :phase, :action)
            """), mem)

        # 5c: Update phase counts / totals on touched existing MCs
        for upd in updates_mcs:
            c.execute(text("""
                UPDATE master_controls
                SET phases_covered = CAST(:phases_covered AS jsonb),
                    phase_control_count = CAST(:phase_control_count AS jsonb),
                    total_controls = :total_controls
                WHERE id = CAST(:mc_uuid AS uuid)
            """), upd)

    logger.info("DONE — wrote %d new MCs, updated %d existing MCs, %d members inserted",
                stats["mcs_new_created"], stats["mcs_existing_updated"], stats["members_inserted"])


if __name__ == "__main__":
    main()
