#!/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, NoReturn


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:
    if path.stat().st_size > 2 * 1024 * 1024:
        raise ValueError(f"{path}: migration source exceeds the 2 MiB audit bound")
    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
    values = _imported_release_metadata(owner, path, tree)
    wrapped = _wrapped_release_metadata(owner, path, tree)
    if values and wrapped:
        raise ValueError(f"{path}: mixed migration metadata wrapper styles are ambiguous")
    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] = _migration_assignment_value(
                        target.id,
                        statement.value,
                        wrapped=wrapped,
                    )
        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] = _migration_assignment_value(
                statement.target.id,
                statement.value,
                wrapped=wrapped,
            )
    revision = values.get("revision")
    if not isinstance(revision, str):
        if "revision" in values:
            raise ValueError(f"{path}: migration revision must be an explicit string")
        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 _ast_binding_count(tree: ast.Module, name: str) -> int:
    count = 0
    for node in ast.walk(tree):
        if isinstance(node, ast.Name) and node.id == name and isinstance(node.ctx, (ast.Store, ast.Del)):
            count += 1
        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.name == name:
            count += 1
        elif isinstance(node, (ast.Import, ast.ImportFrom)):
            count += sum(
                (alias.asname or (alias.name.split(".")[0] if isinstance(node, ast.Import) else alias.name)) == name
                for alias in node.names
            )
    return count


def _release_wrapper_peer(owner: str, path: Path, filename: str) -> Migration:
    versions = path.parent.parent / "versions"
    peer = versions / filename
    if (
        path.parent.name != "dev_versions"
        or Path(filename).name != filename or not filename.endswith(".py")
        or versions.is_symlink() or peer.is_symlink() or not peer.is_file()
        or peer.resolve().parent != versions.resolve()
        or peer.stat().st_size > 2 * 1024 * 1024
    ):
        raise ValueError(f"{path}: unsupported or ambiguous development migration wrapper")
    migration = parse_migration_file(owner, peer)
    if migration is None:
        raise ValueError(f"{path}: release wrapper target has no migration metadata")
    return migration


def _imported_release_metadata(owner: str, path: Path, tree: ast.Module) -> dict[str, Any]:
    """Recognize explicit metadata re-exports, never star/dynamic imports."""
    names = {"revision", "down_revision", "depends_on", "branch_labels"}
    prefix = f"{owner.replace('-', '_')}.backend.migrations.versions."
    values: dict[str, Any] = {}
    targets: set[str] = set()
    for statement in tree.body:
        if not isinstance(statement, ast.ImportFrom):
            continue
        selected = [alias for alias in statement.names if alias.name in names or (alias.asname or alias.name) in names]
        if not selected:
            if (statement.module or "").startswith(prefix) and any(alias.name == "*" for alias in statement.names):
                raise ValueError(f"{path}: unsupported wildcard migration metadata")
            continue
        stem = (statement.module or "").removeprefix(prefix)
        if (
            statement.level or not (statement.module or "").startswith(prefix)
            or not stem or any(character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_" for character in stem)
        ):
            raise ValueError(f"{path}: unsupported or ambiguous imported migration metadata")
        targets.add(stem)
        if len(targets) != 1:
            raise ValueError(f"{path}: ambiguous imported migration metadata sources")
        migration = _release_wrapper_peer(owner, path, stem + ".py")
        projection = {
            "revision": migration.revision, "down_revision": migration.down_revisions,
            "depends_on": migration.depends_on, "branch_labels": migration.branch_labels,
        }
        for alias in selected:
            name = alias.asname or alias.name
            if name != alias.name or name in values or _ast_binding_count(tree, name) != 1:
                raise ValueError(f"{path}: ambiguous imported migration metadata assignment")
            values[name] = projection[name]
    return values


def _wrapped_release_metadata(owner: str, path: Path, tree: ast.Module) -> dict[str, Migration]:
    """Resolve only known literal sibling wrappers, without importing any code."""
    metadata_names = {"revision", "down_revision", "depends_on", "branch_labels"}
    aliases: set[str] = set()
    bindings: dict[str, list[ast.expr]] = {}
    imported: dict[str, list[str]] = {}
    for statement in tree.body:
        if isinstance(statement, ast.ImportFrom) and not statement.level:
            for alias in statement.names:
                imported.setdefault(alias.asname or alias.name, []).append(f"{statement.module}.{alias.name}")
        targets: list[ast.expr] = []
        value: ast.expr | None = None
        if isinstance(statement, ast.Assign):
            targets, value = statement.targets, statement.value
        elif isinstance(statement, ast.AnnAssign) and statement.value is not None:
            targets, value = [statement.target], statement.value
        for target in targets:
            if isinstance(target, ast.Name) and value is not None:
                bindings.setdefault(target.id, []).append(value)
                if target.id in metadata_names and isinstance(value, ast.Attribute) and isinstance(value.value, ast.Name):
                    aliases.add(value.value.id)
    if not aliases:
        return {}
    if path.parent.name != "dev_versions":
        raise ValueError(f"{path}: non-literal release migration metadata is unsupported")

    def reject() -> NoReturn:
        raise ValueError(f"{path}: unsupported or ambiguous development migration wrapper")

    def binding_count(name: str) -> int:
        return _ast_binding_count(tree, name)

    def assigned(name: str) -> ast.expr:
        values = bindings.get(name, ())
        if len(values) != 1 or binding_count(name) != 1 or name in imported:
            reject()
        return values[0]

    def imported_as(node: ast.expr, qualified: str) -> bool:
        return (
            isinstance(node, ast.Name)
            and imported.get(node.id) == [qualified]
            and binding_count(node.id) == 1
        )

    def call(node: ast.expr, qualified: str, arguments: int) -> bool:
        return isinstance(node, ast.Call) and imported_as(node.func, qualified) and len(node.args) == arguments and not node.keywords

    def file_wrapper_target(node: ast.expr) -> str:
        if binding_count("__file__"):
            reject()
        if isinstance(node, ast.Name):
            node = assigned(node.id)
        if not (
            isinstance(node, ast.BinOp) and isinstance(node.op, ast.Div)
            and isinstance(node.right, ast.Constant) and isinstance(node.right.value, str)
            and isinstance(node.left, ast.BinOp) and isinstance(node.left.op, ast.Div)
            and isinstance(node.left.right, ast.Constant) and node.left.right.value == "versions"
        ):
            reject()
        root = node.left.left
        for name in imported:
            if imported_as(ast.Name(id=name), "pathlib.Path"):
                expected = ast.parse(f"{name}(__file__).resolve().parents[1]", mode="eval").body
                if ast.dump(root) == ast.dump(expected):
                    return node.right.value
        reject()

    result: dict[str, Migration] = {}
    resolved: set[Path] = set()
    for alias in sorted(aliases):
        value = assigned(alias)
        if call(value, "importlib.import_module", 1):
            argument = value.args[0]
            if not isinstance(argument, ast.Constant) or not isinstance(argument.value, str):
                reject()
            prefix = f"{owner.replace('-', '_')}.backend.migrations.versions."
            if not argument.value.startswith(prefix):
                reject()
            stem = argument.value.removeprefix(prefix)
            if not stem or any(character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_" for character in stem):
                reject()
            filename = stem + ".py"
        elif call(value, "importlib.util.module_from_spec", 1):
            argument = value.args[0]
            if not isinstance(argument, ast.Name):
                reject()
            specification = assigned(argument.id)
            if not call(specification, "importlib.util.spec_from_file_location", 2):
                reject()
            if not isinstance(specification.args[0], ast.Constant) or not isinstance(specification.args[0].value, str):
                reject()
            filename = file_wrapper_target(specification.args[1])
        else:
            reject()
        migration = _release_wrapper_peer(owner, path, filename)
        peer = migration.path
        resolved.add(peer.resolve())
        if len(resolved) > 1:
            reject()
        result[alias] = migration
    return result


def _migration_assignment_value(
    name: str,
    value: ast.expr,
    *,
    wrapped: dict[str, Migration],
) -> Any:
    try:
        return ast.literal_eval(value)
    except (ValueError, TypeError):
        if (
            not isinstance(value, ast.Attribute)
            or not isinstance(value.value, ast.Name)
            or value.value.id not in wrapped
            or value.attr != name
        ):
            raise ValueError(f"Unsupported or ambiguous migration metadata: {name}")
    migration = wrapped[value.value.id]
    return {
        "revision": migration.revision,
        "down_revision": migration.down_revisions,
        "depends_on": migration.depends_on,
        "branch_labels": migration.branch_labels,
    }[name]


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())