#!/usr/bin/env python3
"""Upsert derived QUAIDAL controls from YAML into compliance.derived_controls.

Reads:
  control-pipeline/data/quaidal/master_controls.yaml
  control-pipeline/data/quaidal/atomic_controls.yaml
  control-pipeline/data/quaidal/mitigations.yaml
  control-pipeline/data/quaidal/metrics.yaml

Writes: compliance.derived_controls (idempotent UPSERT by derived_id)

Usage:
    # Mac Mini direct:
    python3 control-pipeline/scripts/apply_quaidal_to_db.py

    # Via SSH (locally, against macmini DB):
    DB_HOST=macmini python3 control-pipeline/scripts/apply_quaidal_to_db.py
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

try:
    import psycopg
    import yaml
except ImportError as e:
    print(f"ERROR: missing dependency {e.name}. Install with: pip install psycopg[binary] pyyaml", file=sys.stderr)
    sys.exit(2)

REPO_ROOT = Path(__file__).resolve().parents[2]
DATA_DIR = REPO_ROOT / "control-pipeline" / "data" / "quaidal"

KIND_FILES = {
    "criterion":     "master_controls.yaml",
    "building_block": "atomic_controls.yaml",
    "measure":       "mitigations.yaml",
    "metric":        "metrics.yaml",
}

UPSERT_SQL = """
INSERT INTO compliance.derived_controls (
    derived_id, kind, canonical_name, description, regulation_anchor,
    related_quaidal_ids, external_refs,
    source_framework, source_section, source_url, source_commit_sha,
    source_title_original, source_license_note,
    plagiarism_score_at_generation, generated_by_model, yaml_path
) VALUES (
    %(derived_id)s, %(kind)s, %(canonical_name)s, %(description)s, %(regulation_anchor)s,
    %(related_quaidal_ids)s::jsonb, %(external_refs)s::jsonb,
    %(source_framework)s, %(source_section)s, %(source_url)s, %(source_commit_sha)s,
    %(source_title_original)s, %(source_license_note)s,
    %(plagiarism_score)s, %(generated_by_model)s, %(yaml_path)s
)
ON CONFLICT (derived_id) DO UPDATE SET
    kind = EXCLUDED.kind,
    canonical_name = EXCLUDED.canonical_name,
    description = EXCLUDED.description,
    regulation_anchor = EXCLUDED.regulation_anchor,
    related_quaidal_ids = EXCLUDED.related_quaidal_ids,
    external_refs = EXCLUDED.external_refs,
    source_framework = EXCLUDED.source_framework,
    source_section = EXCLUDED.source_section,
    source_url = EXCLUDED.source_url,
    source_commit_sha = EXCLUDED.source_commit_sha,
    source_title_original = EXCLUDED.source_title_original,
    source_license_note = EXCLUDED.source_license_note,
    plagiarism_score_at_generation = EXCLUDED.plagiarism_score_at_generation,
    generated_by_model = EXCLUDED.generated_by_model,
    yaml_path = EXCLUDED.yaml_path
"""


def load_yaml_records(yaml_path: Path) -> tuple[list[dict], str | None, str | None]:
    if not yaml_path.exists():
        return [], None, None
    data = yaml.safe_load(yaml_path.read_text(encoding="utf-8"))
    return data.get("controls", []), data.get("commit_sha"), data.get("generated_by_model")


def to_row(ctrl: dict, yaml_path: Path, default_model: str | None, default_commit: str | None) -> dict:
    source = ctrl.get("source") or {}
    return {
        "derived_id":           ctrl["id"],
        "kind":                 ctrl["kind"],
        "canonical_name":       ctrl["canonical_name"],
        "description":          ctrl["description"],
        "regulation_anchor":    ctrl.get("regulation_anchor"),
        "related_quaidal_ids":  json.dumps(ctrl.get("related_quaidal_ids", []), ensure_ascii=False),
        "external_refs":        json.dumps(ctrl.get("external_refs", []), ensure_ascii=False),
        "source_framework":     source.get("framework", "BSI QUAIDAL"),
        "source_section":       source.get("section", ""),
        "source_url":           source.get("url"),
        "source_commit_sha":    source.get("commit_sha") or default_commit,
        "source_title_original": source.get("title_original_de"),
        "source_license_note":  source.get("license_note"),
        "plagiarism_score":     ctrl.get("plagiarism_score_at_generation"),
        "generated_by_model":   default_model,
        "yaml_path":            str(yaml_path.relative_to(REPO_ROOT)),
    }


def build_dsn(args: argparse.Namespace) -> str:
    if args.dsn:
        return args.dsn
    return (
        f"host={args.db_host} port={args.db_port} "
        f"dbname={args.db_name} user={args.db_user} password={args.db_password}"
    )


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--dsn", help="Full DSN; overrides individual flags")
    ap.add_argument("--db-host", default=os.environ.get("DB_HOST", "localhost"))
    ap.add_argument("--db-port", default=os.environ.get("DB_PORT", "5432"))
    ap.add_argument("--db-name", default=os.environ.get("DB_NAME", "breakpilot_db"))
    ap.add_argument("--db-user", default=os.environ.get("DB_USER", "breakpilot"))
    ap.add_argument("--db-password", default=os.environ.get("DB_PASSWORD", "breakpilot"))
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    total = 0
    rows: list[dict] = []
    for kind, fname in KIND_FILES.items():
        path = DATA_DIR / fname
        records, commit, model = load_yaml_records(path)
        for rec in records:
            rows.append(to_row(rec, path, model, commit))
        if records:
            print(f"  {fname}: {len(records)} entries", file=sys.stderr)
        total += len(records)

    if not rows:
        print("ERROR: no YAML records found; run derive_quaidal_mcs.py first", file=sys.stderr)
        return 2

    print(f"Total rows: {total}", file=sys.stderr)
    if args.dry_run:
        print("Dry run — sample row:", file=sys.stderr)
        print(json.dumps({k: (v[:200] if isinstance(v, str) else v) for k, v in rows[0].items()}, indent=2, ensure_ascii=False))
        return 0

    dsn = build_dsn(args)
    print(f"Connecting to {args.db_host}:{args.db_port}/{args.db_name}", file=sys.stderr)
    inserted = updated = 0
    with psycopg.connect(dsn) as conn:
        with conn.cursor() as cur:
            for row in rows:
                cur.execute(
                    "SELECT 1 FROM compliance.derived_controls WHERE derived_id = %s",
                    (row["derived_id"],),
                )
                existed = cur.fetchone() is not None
                cur.execute(UPSERT_SQL, row)
                if existed:
                    updated += 1
                else:
                    inserted += 1
        conn.commit()
    print(f"Inserted: {inserted}, Updated: {updated}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())
