#!/usr/bin/env python
from __future__ import annotations

import argparse
import ast
from dataclasses import dataclass
from datetime import datetime, timezone
import json
import os
from pathlib import Path
from typing import Any


META_ROOT = Path(__file__).resolve().parents[2]
ROOT = Path(os.environ.get("GOVOPLAN_CORE_ROOT", META_ROOT.parent / "govoplan-core")).resolve()
DEFAULT_BASELINE_FILE = ROOT / "docs" / "migration-release-baselines.json"
MIGRATION_TRACK_RELEASE = "release"
MIGRATION_TRACK_DEV = "dev"
MIGRATION_TRACKS = (MIGRATION_TRACK_RELEASE, MIGRATION_TRACK_DEV)


@dataclass(frozen=True, slots=True)
class Migration:
    owner: str
    path: Path
    revision: str
    down_revisions: tuple[str, ...]
    depends_on: tuple[str, ...]
    branch_labels: tuple[str, ...]


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Audit GovOPlaN Alembic migrations against the release baseline policy.",
    )
    parser.add_argument(
        "--workspace-root",
        type=Path,
        default=ROOT.parent,
        help="Directory containing govoplan-* checkouts. Defaults to the parent of govoplan-core.",
    )
    parser.add_argument(
        "--baseline-file",
        type=Path,
        default=DEFAULT_BASELINE_FILE,
        help="JSON file recording released migration heads.",
    )
    parser.add_argument(
        "--track",
        choices=MIGRATION_TRACKS,
        default=MIGRATION_TRACK_RELEASE,
        help="Migration track to audit. The release track is compared to release baselines; the dev track only validates the detailed graph.",
    )
    parser.add_argument(
        "--strict",
        action="store_true",
        help="Fail when current heads are not recorded in the latest release baseline.",
    )
    parser.add_argument(
        "--strict-if-baseline",
        action="store_true",
        help="Run in strict mode only after at least one release baseline is recorded.",
    )
    parser.add_argument(
        "--record-release",
        metavar="VERSION",
        help="Record the current reviewed migration heads as a release baseline.",
    )
    parser.add_argument(
        "--replace-release",
        action="store_true",
        help="Replace an existing release entry when used with --record-release.",
    )
    parser.add_argument(
        "--squash-plan",
        action="store_true",
        help="Print the reviewed/manual squash checklist for the current graph.",
    )
    parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON.")
    args = parser.parse_args()

    migrations = discover_migrations(args.workspace_root, track=args.track)
    baseline = load_baseline_file(args.baseline_file)
    report = build_report(
        migrations,
        baseline=baseline,
        baseline_file=args.baseline_file,
        workspace_root=args.workspace_root,
        track=args.track,
    )

    if args.record_release:
        if args.track != MIGRATION_TRACK_RELEASE:
            raise SystemExit("--record-release is only valid for --track release")
        if report["graph_errors"]:
            print_text_report(report)
            return 1
        baseline = record_release_baseline(
            baseline,
            report,
            release=args.record_release,
            replace=args.replace_release,
        )
        write_baseline_file(args.baseline_file, baseline)
        report = build_report(
            migrations,
            baseline=baseline,
            baseline_file=args.baseline_file,
            workspace_root=args.workspace_root,
            track=args.track,
        )

    if args.squash_plan:
        print_squash_plan(report)
    elif args.json:
        print(json.dumps(report, indent=2, sort_keys=True))
    else:
        print_text_report(report)

    has_graph_errors = bool(report["graph_errors"])
    has_strict_errors = bool(report["strict_errors"])
    strict_required = args.strict or (args.strict_if_baseline and report["release_count"] > 0)
    if has_graph_errors or (strict_required and has_strict_errors):
        return 1
    return 0


def discover_migrations(workspace_root: Path, *, track: str = MIGRATION_TRACK_RELEASE) -> list[Migration]:
    migrations: list[Migration] = []
    for versions_dir in _migration_version_dirs(workspace_root, track=track):
        owner = owner_for_versions_dir(versions_dir)
        for path in sorted(versions_dir.glob("*.py")):
            if path.name == "__init__.py":
                continue
            migration = parse_migration_file(owner, path)
            if migration is not None:
                migrations.append(migration)
    return migrations


def _migration_version_dirs(workspace_root: Path, *, track: str = MIGRATION_TRACK_RELEASE) -> list[Path]:
    version_dir_name = "dev_versions" if track == MIGRATION_TRACK_DEV else "versions"
    candidates = [ROOT / "alembic" / version_dir_name]
    for repo in sorted(workspace_root.glob("govoplan-*")):
        candidates.extend(sorted(repo.glob(f"src/govoplan_*/backend/migrations/{version_dir_name}")))
    seen: set[Path] = set()
    result: list[Path] = []
    for candidate in candidates:
        resolved = candidate.resolve()
        if resolved in seen or not candidate.is_dir():
            continue
        seen.add(resolved)
        result.append(candidate)
    return result


def owner_for_versions_dir(versions_dir: Path) -> str:
    core_locations = {
        (ROOT / "alembic" / "versions").resolve(),
        (ROOT / "alembic" / "dev_versions").resolve(),
    }
    if versions_dir.resolve() in core_locations:
        return "govoplan-core"
    for parent in versions_dir.parents:
        if parent.name.startswith("govoplan-"):
            return parent.name
    return versions_dir.parent.name


def parse_migration_file(owner: str, path: Path) -> Migration | None:
    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
    values: dict[str, Any] = {}
    for statement in tree.body:
        if isinstance(statement, ast.Assign):
            for target in statement.targets:
                if isinstance(target, ast.Name) and target.id in {"revision", "down_revision", "depends_on", "branch_labels"}:
                    values[target.id] = ast.literal_eval(statement.value)
        elif (
            isinstance(statement, ast.AnnAssign)
            and isinstance(statement.target, ast.Name)
            and statement.target.id in {"revision", "down_revision", "depends_on", "branch_labels"}
            and statement.value is not None
        ):
            values[statement.target.id] = ast.literal_eval(statement.value)
    revision = values.get("revision")
    if not isinstance(revision, str):
        return None
    return Migration(
        owner=owner,
        path=path,
        revision=revision,
        down_revisions=_normalize_revision_tuple(values.get("down_revision")),
        depends_on=_normalize_revision_tuple(values.get("depends_on")),
        branch_labels=_normalize_string_tuple(values.get("branch_labels")),
    )


def _normalize_revision_tuple(value: Any) -> tuple[str, ...]:
    if value is None:
        return ()
    if isinstance(value, str):
        return (value,)
    if isinstance(value, (list, tuple, set)):
        return tuple(str(item) for item in value if item is not None)
    return (str(value),)


def _normalize_string_tuple(value: Any) -> tuple[str, ...]:
    if value is None:
        return ()
    if isinstance(value, str):
        return (value,)
    if isinstance(value, (list, tuple, set)):
        return tuple(str(item) for item in value)
    return (str(value),)


def load_baseline_file(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {"version": 1, "releases": [], "_missing": True}
    payload = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise SystemExit(f"{path} must contain a JSON object")
    releases = payload.get("releases")
    if not isinstance(releases, list):
        raise SystemExit(f"{path} must contain a releases list")
    return payload


def build_report(
    migrations: list[Migration],
    *,
    baseline: dict[str, Any],
    baseline_file: Path,
    workspace_root: Path,
    track: str = MIGRATION_TRACK_RELEASE,
) -> dict[str, Any]:
    revisions: dict[str, Migration] = {}
    graph_errors: list[str] = []
    for migration in migrations:
        existing = revisions.get(migration.revision)
        if existing is not None:
            graph_errors.append(
                f"duplicate revision {migration.revision}: {existing.path} and {migration.path}"
            )
        revisions[migration.revision] = migration

    referenced = {
        revision
        for migration in migrations
        for revision in (*migration.down_revisions, *migration.depends_on)
    }
    for migration in migrations:
        for down_revision in migration.down_revisions:
            if down_revision not in revisions:
                graph_errors.append(
                    f"{migration.revision} references missing down_revision {down_revision} in {migration.path}"
                )
        for dependency in migration.depends_on:
            if dependency not in revisions:
                graph_errors.append(
                    f"{migration.revision} references missing depends_on {dependency} in {migration.path}"
                )

    current_heads = sorted(revision for revision in revisions if revision not in referenced)
    current_head_entries = [
        {
            "owner": revisions[revision].owner,
            "revision": revision,
        }
        for revision in current_heads
    ]
    owner_rows = []
    for owner in sorted({migration.owner for migration in migrations}):
        owner_migrations = [migration for migration in migrations if migration.owner == owner]
        owner_revisions = {migration.revision for migration in owner_migrations}
        owner_referenced = {
            revision
            for migration in owner_migrations
            for revision in (*migration.down_revisions, *migration.depends_on)
            if revision in owner_revisions
        }
        owner_heads = sorted(owner_revisions - owner_referenced)
        owner_rows.append(
            {
                "owner": owner,
                "migrations": len(owner_migrations),
                "heads": owner_heads,
            }
        )

    releases = baseline.get("releases") or []
    latest_release = releases[-1] if releases else None
    latest_heads = _latest_release_heads(latest_release)
    unrecorded_heads = sorted(set(current_heads) - latest_heads) if latest_release else current_heads

    strict_errors: list[str] = []
    if track == MIGRATION_TRACK_RELEASE:
        if baseline.get("_missing"):
            strict_errors.append(f"baseline file is missing: {baseline_file}")
        if not releases:
            strict_errors.append("no released migration baseline has been recorded yet")
        if unrecorded_heads:
            strict_errors.append(
                "current migration heads are not recorded in the latest release baseline: "
                + ", ".join(unrecorded_heads)
            )
    return {
        "workspace_root": str(workspace_root),
        "track": track,
        "baseline_file": str(baseline_file),
        "baseline_file_missing": bool(baseline.get("_missing")),
        "release_count": len(releases),
        "latest_release": latest_release,
        "owners": owner_rows,
        "current_heads": current_heads,
        "current_head_entries": current_head_entries,
        "graph_errors": graph_errors,
        "strict_errors": strict_errors,
    }


def _latest_release_heads(latest_release: Any) -> set[str]:
    if not isinstance(latest_release, dict):
        return set()
    heads = latest_release.get("heads") or []
    graph_heads = latest_release.get("graph_heads") or []
    result: set[str] = set()
    for head_list in (heads, graph_heads):
        if not isinstance(head_list, list):
            continue
        for entry in head_list:
            if isinstance(entry, str):
                result.add(entry)
            elif isinstance(entry, dict) and isinstance(entry.get("revision"), str):
                result.add(entry["revision"])
    owner_heads = latest_release.get("owner_heads") or []
    if isinstance(owner_heads, list):
        for entry in owner_heads:
            if not isinstance(entry, dict):
                continue
            revisions = entry.get("revisions") or []
            if isinstance(revisions, str):
                result.add(revisions)
            elif isinstance(revisions, list):
                result.update(revision for revision in revisions if isinstance(revision, str))
    return result


def record_release_baseline(
    baseline: dict[str, Any],
    report: dict[str, Any],
    *,
    release: str,
    replace: bool,
) -> dict[str, Any]:
    payload = {
        key: value
        for key, value in baseline.items()
        if not key.startswith("_")
    }
    payload.setdefault("version", 1)
    releases = list(payload.get("releases") or [])
    existing_index = next(
        (index for index, item in enumerate(releases) if isinstance(item, dict) and item.get("release") == release),
        None,
    )
    if existing_index is not None and not replace:
        raise SystemExit(f"release {release} already exists in baseline file; pass --replace-release to update it")

    entry = {
        "release": release,
        "recorded_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z"),
        "track": MIGRATION_TRACK_RELEASE,
        "squash_policy": "reviewed-manual",
        "heads": report["current_head_entries"],
        "owner_heads": [
            {
                "owner": owner["owner"],
                "revisions": owner["heads"],
            }
            for owner in report["owners"]
        ],
    }
    if existing_index is None:
        releases.append(entry)
    else:
        releases[existing_index] = entry
    payload["releases"] = releases
    return payload


def write_baseline_file(path: Path, baseline: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(baseline, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def print_text_report(report: dict[str, Any]) -> None:
    print("Migration release audit")
    print(f"  workspace: {report['workspace_root']}")
    print(f"  track: {report['track']}")
    print(f"  baseline file: {report['baseline_file']}")
    print(f"  recorded releases: {report['release_count']}")
    print()
    print("Owners:")
    for owner in report["owners"]:
        heads = ", ".join(owner["heads"]) if owner["heads"] else "-"
        print(f"  {owner['owner']}: {owner['migrations']} migration(s), head(s): {heads}")
    print()
    print("Current heads:")
    for entry in report["current_head_entries"]:
        print(f"  {entry['owner']}: {entry['revision']}")
    if not report["current_head_entries"]:
        print("  -")

    if report["graph_errors"]:
        print()
        print("Graph errors:")
        for error in report["graph_errors"]:
            print(f"  - {error}")

    if report["strict_errors"]:
        print()
        print("Release baseline warnings:")
        for error in report["strict_errors"]:
            print(f"  - {error}")


def print_squash_plan(report: dict[str, Any]) -> None:
    print("Manual migration squash plan")
    print()
    print("Policy:")
    print("  - Do not rewrite revision IDs that have shipped to real installations.")
    print("  - Keep detailed development migrations on the dev track.")
    print("  - Add reviewed release-track baselines or release-to-release step-up migrations.")
    print("  - Review generated release shortcuts like normal schema code.")
    print("  - Run PostgreSQL migration smoke checks after any squash.")
    print()
    print("Current owner heads:")
    for owner in report["owners"]:
        heads = ", ".join(owner["heads"]) if owner["heads"] else "-"
        print(f"  - {owner['owner']}: {heads}")
    print()
    print("Current Alembic graph heads:")
    for entry in report["current_head_entries"]:
        print(f"  - {entry['owner']}: {entry['revision']}")
    if not report["current_head_entries"]:
        print("  -")
    print()
    print("Release steps:")
    print("  1. Decide which dev-track revisions are unreleased and should be folded into the next release shortcut.")
    print("  2. Add reviewed release-track baseline or step-up migrations without deleting the dev-track chain.")
    print("  3. Run tools/release/release-migration-audit.py, tools/release/release-migration-audit.py --track dev, and PostgreSQL release checks.")
    print("  4. Record the reviewed heads with tools/release/release-migration-audit.py --record-release <x.y.z>.")
    print("  5. Cut the release with tools/release/push-release-tag.sh; audit is strict after the first baseline.")


if __name__ == "__main__":
    raise SystemExit(main())