"""Candidate validation, migration, clone, and rollback rehearsal helpers."""

from __future__ import annotations
import json
import os
import shutil
import stat
import tempfile
from pathlib import Path
from typing import Final, cast
from uuid import uuid4
from scripts.nutricoach_v150_live_upgrade_common import (
    JsonValue,
    LiveTarget,
    UpgradeDenied,
    canonical,
    load_object,
    object_at,
    sha256_bytes,
    sha256_file,
)
from scripts.nutricoach_v150_live_upgrade_state import tree_snapshot

MANIFEST_SHA: Final = "78b2d1790b67cf0d9730cd9a4b0d2d0d3de5bf5f5836e31a70f5be3a62dd7786"
CANDIDATE: Final = "066a794d44861d2cd0fe8c1ea14c0050e00219b38a2d2026b389772783056dad"
QUALIFICATION_SHA: Final = (
    "f8bfee22b543ffd96873dec3b857c9e45deadaeb18a83328028fb15ebc6497a3"
)
FINAL_QUALIFICATION_SHA: Final = (
    "82eb4e3821de7b26f3e79efc1695014c893cfa67478c38a415e9efb2d91000db"
)
OLD_DENIAL_SHA: Final = (
    "56a6773241559e4f9d4e0a75fb5bc05eb6b7efe7876592f730cf74ac24a9836d"
)
CAPACITY: Final = 5


def write_receipt(path: Path, value: JsonValue) -> None:
    _ = path.write_bytes(canonical(value) + b"\n")
    path.chmod(0o600)


def validate_candidate(manifest_path: Path) -> dict[str, JsonValue]:
    manifest = load_object(manifest_path)
    capabilities = object_at(manifest.get("capabilities"), "capabilities")
    multi = object_at(
        capabilities.get("nutricoach_multi_customer_v1"), "multi_customer"
    )
    inbox = object_at(capabilities.get("nutricoach_channel_inbox_v1"), "channel_inbox")
    expected_off = {"authorized": False, "compiled": True, "configured": False}
    if (
        sha256_file(manifest_path) != MANIFEST_SHA
        or manifest.get("candidate_digest") != CANDIDATE
        or manifest.get("status") != "QUALIFIED_PENDING_LIVE_AUTHORIZATION"
        or any(inbox.get(key) != value for key, value in expected_off.items())
        or any(multi.get(key) != value for key, value in expected_off.items())
        or multi.get("authorized_capacity") != 0
        or multi.get("post_migration_capacity") != CAPACITY
    ):
        raise UpgradeDenied("candidate_contract")
    qualification = manifest_path.with_name("qualification.json")
    final = manifest_path.parents[1] / "final-qualification.json"
    old_denial = (
        manifest_path.parents[1] / "task-3-supersession/old-manifest-denial.json"
    )
    if (
        sha256_file(qualification) != QUALIFICATION_SHA
        or sha256_file(final) != FINAL_QUALIFICATION_SHA
        or sha256_file(old_denial) != OLD_DENIAL_SHA
    ):
        raise UpgradeDenied("qualification_contract")
    return manifest


def migration(before: bytes) -> tuple[bytes, dict[str, JsonValue]]:
    try:
        document = cast(dict[str, JsonValue], json.loads(before))
    except (UnicodeError, json.JSONDecodeError) as exc:
        raise UpgradeDenied("registry_invalid") from exc
    if document.get("registry_mode") != "ordinary_v1":
        raise UpgradeDenied("registry_mode")
    if document.get("admission_policy") is not None:
        raise UpgradeDenied("migration_already_applied")
    customers = document.get("customers")
    if not isinstance(customers, list):
        raise UpgradeDenied("registry_customers")
    enabled = sum(
        1
        for customer in customers
        if isinstance(customer, dict) and customer.get("enabled") is True
    )
    if enabled > CAPACITY:
        raise UpgradeDenied("capacity")
    document["admission_policy"] = {
        "schema": "nutricoach-multi-customer-v1",
        "candidate_digest": CANDIDATE,
        "max_enabled_customers": CAPACITY,
    }
    after = (json.dumps(document, ensure_ascii=False, indent=2) + "\n").encode()
    before_sha = sha256_bytes(before)
    after_sha = sha256_bytes(after)
    proposal_input: JsonValue = {
        "after_sha256": after_sha,
        "before_sha256": before_sha,
        "candidate_digest": CANDIDATE,
        "max_enabled_customers": CAPACITY,
        "capability": "nutricoach_multi_customer_v1",
    }
    proposal = sha256_bytes(canonical(proposal_input))
    return after, {
        **cast(dict[str, JsonValue], proposal_input),
        "enabled_customers": enabled,
        "proposal_digest": proposal,
        "approval_phrase": f"AUTHORIZE NUTRICOACH MULTI CUSTOMER V1 {proposal}",
        "persisted": False,
        "status": "CAPACITY_5_DRY_RUN",
    }


def _atomic_clone_write(path: Path, payload: bytes, mode: int) -> None:
    temporary = path.with_name(f".{path.name}.{uuid4().hex}.tmp")
    _ = temporary.write_bytes(payload)
    temporary.chmod(mode)
    os.replace(temporary, path)


def cleanup_clone(root: Path) -> None:
    for path in sorted(root.rglob("*"), reverse=True):
        if path.is_dir() and not path.is_symlink():
            path.chmod(stat.S_IMODE(path.stat().st_mode) | stat.S_IRWXU)
    root.chmod(0o700)
    shutil.rmtree(root)


def rollback_rehearsal(
    target: LiveTarget, after: bytes, parent: Path
) -> dict[str, JsonValue]:
    rehearsal = Path(tempfile.mkdtemp(prefix=".nutricoach-v150-rehearsal-", dir=parent))
    clone = rehearsal / "profile"
    registry_relative = Path("customers/registry.json")
    receipt: dict[str, JsonValue] | None = None
    try:
        _ = shutil.copytree(
            target.profile_root, clone, copy_function=os.link, symlinks=True
        )
        before = (clone / registry_relative).read_bytes()
        mode = stat.S_IMODE((clone / registry_relative).stat().st_mode)
        initial = tree_snapshot(clone, classify_volatile=True)
        _atomic_clone_write(clone / registry_relative, after, mode)
        migrated = tree_snapshot(clone, classify_volatile=True)
        _atomic_clone_write(clone / registry_relative, before, mode)
        restored = tree_snapshot(clone, classify_volatile=True)
        if (
            initial["stable_digest"] != restored["stable_digest"]
            or initial["stable_digest"] == migrated["stable_digest"]
        ):
            raise UpgradeDenied("rollback_drift")
        receipt = {
            "initial_digest": initial["stable_digest"],
            "migrated_digest": migrated["stable_digest"],
            "restored_digest": restored["stable_digest"],
            "registry_restored_sha256": sha256_bytes(before),
            "status": "BYTE_EXACT_ROLLBACK_REHEARSED",
        }
    finally:
        cleanup_clone(rehearsal)
    if rehearsal.exists():
        raise UpgradeDenied("rehearsal_cleanup")
    receipt["cleanup_complete"] = True
    return receipt


def changed_paths(
    before: dict[str, JsonValue], after: dict[str, JsonValue]
) -> list[JsonValue]:
    def hashes(snapshot: dict[str, JsonValue]) -> dict[tuple[str, str], JsonValue]:
        result: dict[tuple[str, str], JsonValue] = {}
        for raw in cast(list[JsonValue], snapshot["volatile"]):
            row = cast(dict[str, JsonValue], raw)
            key = (str(row.get("profile", "")), str(row["path"]))
            result[key] = row.get("sha256", row.get("target"))
        return result

    before_hashes, after_hashes = hashes(before), hashes(after)
    return [
        "/".join(key)
        for key in sorted(set(before_hashes) | set(after_hashes))
        if before_hashes.get(key) != after_hashes.get(key)
    ]