"""Upgrade a ModelForge deployment, or refuse and say why.

    python scripts/upgrade.py --plan          # preflight only, changes nothing
    python scripts/upgrade.py --apply         # preflight, then migrate

The preflight is the point. An upgrade that starts and then discovers it has no recoverable backup,
or that the schema in front of it is one this release does not understand, has already taken the
deployment down to learn something it could have known first.

It refuses by default when no verified backup exists. That refusal is overridable, because an
operator who has taken a backup by other means should not be blocked by a tool that cannot see it —
but the override is explicit and recorded, never a default.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any

from alembic import command
from alembic.config import Config
from sqlalchemy import create_engine, inspect, text

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "backend" / "src"))

from modelforge_api.domain.release import (
    MINIMUM_POSTGRES_MAJOR,
    MINIMUM_UPGRADE_SOURCE,
    PRODUCT_NAME,
    PRODUCT_VERSION,
    SUPPORTED_AGENT_PROTOCOL_VERSIONS,
    SUPPORTED_UPGRADE_SOURCE_SCHEMA_REVISIONS,
    TARGET_SCHEMA_REVISION,
    Compatibility,
    agent_protocol_compatibility,
    schema_compatibility,
    upgrade_required,
)

#: A backup older than this is reported as stale. It does not block on its own — the operator sees
#: the age and decides — but an upgrade behind a two-day-old recovery point is worth saying out loud.
BACKUP_FRESHNESS = timedelta(hours=26)


@dataclass
class Finding:
    name: str
    ok: bool
    detail: str
    blocking: bool = True


@dataclass
class UpgradePlan:
    findings: list[Finding] = field(default_factory=list)

    def add(self, name: str, ok: bool, detail: str, *, blocking: bool = True) -> None:
        self.findings.append(Finding(name, ok, detail, blocking))

    @property
    def blockers(self) -> list[Finding]:
        return [item for item in self.findings if item.blocking and not item.ok]

    def as_dict(self) -> dict[str, Any]:
        return {
            "product": PRODUCT_NAME,
            "target_version": PRODUCT_VERSION,
            "target_schema": TARGET_SCHEMA_REVISION,
            "minimum_upgrade_source": MINIMUM_UPGRADE_SOURCE,
            "findings": [
                {
                    "name": item.name,
                    "ok": item.ok,
                    "detail": item.detail,
                    "blocking": item.blocking,
                }
                for item in self.findings
            ],
            "blockers": [item.name for item in self.blockers],
        }


def build_plan(database_url: str, *, require_backup: bool) -> UpgradePlan:
    plan = UpgradePlan()
    engine = create_engine(database_url, pool_pre_ping=True)
    try:
        with engine.connect() as connection:
            server_version = int(
                connection.exec_driver_sql("show server_version_num").scalar_one()
            )
            has_schema = inspect(connection).has_table("alembic_version")
            revision = (
                connection.execute(
                    text("select version_num from alembic_version")
                ).scalar_one_or_none()
                if has_schema
                else None
            )
            major = server_version // 10000
            plan.add(
                "postgresql version",
                major >= MINIMUM_POSTGRES_MAJOR,
                f"server major {major} (minimum {MINIMUM_POSTGRES_MAJOR})",
            )

            if revision == TARGET_SCHEMA_REVISION:
                plan.add(
                    "current schema",
                    True,
                    f"{revision} — already at the target; no migration will run",
                )
            elif revision in SUPPORTED_UPGRADE_SOURCE_SCHEMA_REVISIONS:
                plan.add(
                    "current schema",
                    True,
                    f"supported upgrade source {revision} -> {TARGET_SCHEMA_REVISION}",
                )
            else:
                compatibility = schema_compatibility(revision)
                plan.add(
                    "current schema",
                    False,
                    f"{revision or '(empty)'} is {compatibility}; "
                    f"{upgrade_required(compatibility)}",
                )

            # A database can carry alembic_version without carrying the rest of the schema — a
            # half-finished migration, or a database that simply is not ModelForge's. Probing it as
            # though the tables exist turns a preflight into a stack trace, which is the opposite of
            # what a preflight is for.
            inspector = inspect(connection)
            expected_tables = (
                "backup_sets",
                "serving_jobs",
                "compute_nodes",
                "migration_cutover_operations",
            )
            missing_tables = [
                name for name in expected_tables if not inspector.has_table(name)
            ]
            if has_schema and missing_tables:
                plan.add(
                    "schema completeness",
                    False,
                    f"alembic_version is present but {len(missing_tables)} expected table(s) are "
                    f"missing ({', '.join(missing_tables)}); this database is either partially "
                    "migrated or not a ModelForge control plane",
                )

            if has_schema and not missing_tables:
                newest = connection.execute(
                    text(
                        "select backup_id, created_at from backup_sets "
                        "where state = 'VERIFIED' order by created_at desc limit 1"
                    )
                ).first()
                if newest is None:
                    plan.add(
                        "verified backup",
                        not require_backup,
                        "no verified backup exists; an upgrade without a recovery point cannot be "
                        "undone if the schema turns out to be irreversible",
                        blocking=require_backup,
                    )
                else:
                    backup_id, created_at = newest
                    age = datetime.now(UTC) - created_at
                    plan.add(
                        "verified backup",
                        True,
                        f"{backup_id}, {age.total_seconds() / 3600:.1f} h old",
                    )
                    plan.add(
                        "backup freshness",
                        age <= BACKUP_FRESHNESS,
                        f"{age.total_seconds() / 3600:.1f} h old "
                        f"(advisory threshold {BACKUP_FRESHNESS.total_seconds() / 3600:.0f} h)",
                        blocking=False,
                    )

                stale_work = connection.execute(
                    text(
                        "select count(*) from serving_jobs "
                        "where status in ('queued','leased','running')"
                    )
                ).scalar_one()
                plan.add(
                    "no work in flight",
                    int(stale_work) == 0,
                    f"{stale_work} serving job(s) queued, leased or running",
                    blocking=False,
                )

                protocols = connection.execute(
                    text(
                        "select distinct agent_protocol_version from compute_nodes "
                        "where enabled and agent_protocol_version is not null"
                    )
                ).scalars().all()
                incompatible = [
                    version
                    for version in protocols
                    if agent_protocol_compatibility(int(version)) is not Compatibility.COMPATIBLE
                ]
                plan.add(
                    "agent protocol",
                    not incompatible,
                    f"enabled nodes speak {sorted(protocols) or ['(none reported)']}; "
                    f"this release supports {list(SUPPORTED_AGENT_PROTOCOL_VERSIONS)}"
                    + (f"; incompatible: {incompatible}" if incompatible else ""),
                )

                cutovers = connection.execute(
                    text(
                        "select count(*) from migration_cutover_operations "
                        "where stage not in ('COMMITTED','ROLLED_BACK','ABORTED')"
                    )
                ).scalar_one()
                plan.add(
                    "no migration cutover mid-flight",
                    int(cutovers) == 0,
                    f"{cutovers} cutover(s) in an intermediate stage; these are never "
                    "auto-resolved and an upgrade will not resolve them either",
                    blocking=False,
                )
    finally:
        engine.dispose()

    plan.add(
        "rollback",
        True,
        "the schema changes during this upgrade; application rollback requires restoring the "
        "verified pre-upgrade backup. The 0022 -> 0021 DDL downgrade is rehearsal-only because "
        "it removes terminal node tombstones and decommission-operation evidence",
        blocking=False,
    )
    return plan


def apply_migrations(database_url: str) -> tuple[str | None, str, float]:
    started = time.perf_counter()
    engine = create_engine(database_url)
    with engine.connect() as connection:
        before = (
            connection.execute(text("select version_num from alembic_version")).scalar_one_or_none()
            if inspect(connection).has_table("alembic_version")
            else None
        )
    config = Config(str(ROOT / "backend" / "alembic.ini"))
    config.set_main_option("script_location", str(ROOT / "backend" / "alembic"))
    config.set_main_option("sqlalchemy.url", database_url)
    command.upgrade(config, "head")
    with engine.connect() as connection:
        after = connection.execute(text("select version_num from alembic_version")).scalar_one()
    engine.dispose()
    return before, str(after), time.perf_counter() - started


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--database-url", required=True)
    parser.add_argument("--plan", action="store_true", help="preflight only")
    parser.add_argument("--apply", action="store_true", help="preflight, then migrate")
    parser.add_argument(
        "--allow-without-backup",
        action="store_true",
        help="proceed with no verified backup; explicit and recorded, never a default",
    )
    parser.add_argument("--report", default=None)
    args = parser.parse_args(argv)
    if not (args.plan or args.apply):
        parser.error("choose --plan or --apply")

    print(f"{PRODUCT_NAME} upgrade to {PRODUCT_VERSION}", flush=True)
    plan = build_plan(args.database_url, require_backup=not args.allow_without_backup)
    for finding in plan.findings:
        mark = "OK  " if finding.ok else ("BLOCK" if finding.blocking else "WARN")
        print(f"  {mark:5} {finding.name:32} {finding.detail}", flush=True)

    report: dict[str, Any] = plan.as_dict()
    report["applied"] = False
    if plan.blockers:
        print(f"\nrefusing to upgrade: {len(plan.blockers)} blocker(s)", flush=True)
    elif args.apply:
        before, after, seconds = apply_migrations(args.database_url)
        report["applied"] = True
        report["migration"] = {
            "from": before,
            "to": after,
            "seconds": round(seconds, 3),
            "changed": before != after,
        }
        verb = "already current" if before == after else "migrated"
        print(f"\n  schema {verb}: {before} -> {after} in {seconds:.2f}s", flush=True)

    if args.report:
        Path(args.report).write_text(json.dumps(report, indent=2), encoding="utf-8")
        print(f"  report written to {args.report}", flush=True)
    return 1 if plan.blockers else 0


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