#!/usr/bin/env python3
"""Synchronize a GovOPlaN Python environment when package metadata changes."""

from __future__ import annotations

import argparse
from dataclasses import dataclass
import hashlib
import json
import os
from pathlib import Path
import shlex
import subprocess
import sys
from typing import Any

META_ROOT = Path(__file__).resolve().parents[2]
STAMP_VERSION = 1


@dataclass(frozen=True, slots=True)
class RequirementEntry:
    line: str
    install_args: tuple[str, ...]
    is_local: bool
    key: str
    target_path: str | None = None
    pyproject: str | None = None

    def as_dict(self) -> dict[str, object]:
        payload: dict[str, object] = {
            "line": self.line,
            "install_args": list(self.install_args),
            "is_local": self.is_local,
            "key": self.key,
        }
        if self.target_path:
            payload["target_path"] = self.target_path
        if self.pyproject:
            payload["pyproject"] = self.pyproject
        return payload


@dataclass(frozen=True, slots=True)
class InstallPlan:
    mode: str
    reason: str
    commands: tuple[tuple[str, ...], ...]
    warnings: tuple[str, ...] = ()


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--requirements",
        type=Path,
        default=META_ROOT / "requirements-dev.txt",
        help="Requirements file to install. Defaults to requirements-dev.txt.",
    )
    parser.add_argument(
        "--python",
        default=sys.executable,
        help="Python executable to use for pip. Defaults to the current interpreter.",
    )
    parser.add_argument(
        "--stamp",
        type=Path,
        default=None,
        help="Stamp file used to remember the last synchronized metadata digest.",
    )
    parser.add_argument("--force", action="store_true", help="Run pip even when the environment stamp is current.")
    parser.add_argument("--check", action="store_true", help="Exit 1 if synchronization would be required.")
    parser.add_argument("--dry-run", action="store_true", help="Print what would run without changing the environment.")
    parser.add_argument("--upgrade-pip", action="store_true", help="Upgrade pip before installing requirements.")
    args = parser.parse_args()

    requirements = args.requirements.expanduser().resolve()
    if not requirements.exists():
        print(f"Requirements file not found: {requirements}", file=sys.stderr)
        return 2

    python = normalize_python_command(args.python)
    stamp = (args.stamp or default_stamp_path(python, requirements)).expanduser().resolve()
    local_requirements = local_requirement_entries(requirements)
    fingerprint = build_fingerprint(requirements=requirements, python=python, local_requirements=local_requirements)
    previous = load_stamp(stamp)
    current = previous.get("digest") == fingerprint["digest"]

    if current and not args.force:
        payload = stamp_payload(
            requirements=requirements,
            python=python,
            fingerprint=fingerprint,
            entries=parse_requirement_entries(requirements),
        )
        if stamp_needs_refresh(previous, payload):
            if args.dry_run:
                print(f"Python environment is current for {requirements.name}. Stamp metadata would be refreshed.")
                print("No pip command required; the environment stamp would be refreshed.")
                return 0
            if args.check:
                print(f"Python environment is current for {requirements.name}.")
                return 0
            write_stamp(stamp, payload)
            print(f"Python environment is current for {requirements.name}. Stamp metadata refreshed.")
            return 0
        print(f"Python environment is current for {requirements.name}.")
        return 0

    print(f"Python environment metadata is stale for {requirements.name}.")
    print(f"Tracked inputs: {len(fingerprint['inputs'])}")

    if args.check:
        return 1

    plan = build_install_plan(
        previous=previous,
        fingerprint=fingerprint,
        requirements=requirements,
        python=python,
        local_requirements=local_requirements,
        force=args.force,
    )

    print(f"{plan.mode}: {plan.reason}")
    for warning in plan.warnings:
        print(f"warning: {warning}", file=sys.stderr)

    commands: list[tuple[str, ...]] = []
    if args.upgrade_pip:
        commands.append((python, "-m", "pip", "install", "--upgrade", "pip"))
    commands.extend(plan.commands)

    if args.dry_run:
        for command in commands:
            print("+ " + format_command(command))
        if not commands:
            print("No pip command required; the environment stamp would be refreshed.")
        return 0

    for command in commands:
        subprocess.run(command, check=True)

    write_stamp(
        stamp,
        stamp_payload(
            requirements=requirements,
            python=python,
            fingerprint=fingerprint,
            entries=parse_requirement_entries(requirements),
        ),
    )
    print(f"Python environment synchronized. Stamp written to {stamp}.")
    return 0


def default_stamp_path(python: str, requirements: Path) -> Path:
    executable = Path(python)
    if executable.name.startswith("python") and executable.parent.name == "bin":
        return executable.parent.parent / f".govoplan-sync-{requirements.stem}.json"
    return META_ROOT / "runtime" / f".govoplan-sync-{requirements.stem}.json"


def normalize_python_command(value: str) -> str:
    if os.sep not in value:
        return value
    path = Path(value).expanduser()
    if not path.is_absolute():
        path = Path.cwd() / path
    return str(path)


def stamp_payload(
    *,
    requirements: Path,
    python: str,
    fingerprint: dict[str, Any],
    entries: tuple[RequirementEntry, ...],
) -> dict[str, object]:
    return {
        "version": STAMP_VERSION,
        "digest": fingerprint["digest"],
        "requirements": str(requirements),
        "python": python,
        "inputs": fingerprint["inputs"],
        "requirements_entries": [entry.as_dict() for entry in entries],
    }


def stamp_needs_refresh(previous: dict[str, Any], payload: dict[str, object]) -> bool:
    return any(previous.get(key) != value for key, value in payload.items())


def build_fingerprint(*, requirements: Path, python: str, local_requirements: tuple[RequirementEntry, ...]) -> dict[str, Any]:
    hasher = hashlib.sha256()
    inputs: list[dict[str, str]] = []

    add_text(hasher, "stamp-version", str(STAMP_VERSION))
    add_text(hasher, "python", python)
    add_file(hasher, inputs, requirements)

    for project in local_pyprojects(local_requirements):
        add_file(hasher, inputs, project)

    return {"digest": hasher.hexdigest(), "inputs": inputs}


def build_install_plan(
    *,
    previous: dict[str, Any],
    fingerprint: dict[str, Any],
    requirements: Path,
    python: str,
    local_requirements: tuple[RequirementEntry, ...],
    force: bool,
) -> InstallPlan:
    full_command = (python, "-m", "pip", "install", "-r", str(requirements))
    if force:
        return InstallPlan("Full Python environment sync", "--force was requested.", (full_command,))
    if previous.get("version") != STAMP_VERSION:
        return InstallPlan("Full Python environment sync", "the environment stamp is missing or uses an older format.", (full_command,))
    if previous.get("python") != python:
        return InstallPlan("Full Python environment sync", "the Python executable changed.", (full_command,))

    previous_inputs = inputs_by_path(previous.get("inputs"))
    current_inputs = inputs_by_path(fingerprint.get("inputs"))
    if not previous_inputs:
        return InstallPlan("Full Python environment sync", "there is no previous input fingerprint.", (full_command,))

    requirements_key = str(requirements)
    changed_paths = {
        path
        for path, digest in current_inputs.items()
        if previous_inputs.get(path) != digest
    }
    removed_paths = set(previous_inputs) - set(current_inputs)
    stale_requirements = requirements_key in changed_paths or requirements_key in removed_paths

    commands: list[tuple[str, ...]] = []
    warnings: list[str] = []

    if stale_requirements:
        delta = requirement_delta_plan(previous.get("requirements_entries"), parse_requirement_entries(requirements))
        if delta is None:
            return InstallPlan(
                "Full Python environment sync",
                "requirements changed in a way that needs the full resolver.",
                (full_command,),
            )
        for install_args in delta.commands:
            commands.append((python, "-m", "pip", "install", *install_args))
        warnings.extend(delta.warnings)
        changed_paths.discard(requirements_key)
        removed_paths.discard(requirements_key)

    local_by_pyproject = {
        entry.pyproject: entry
        for entry in local_requirements
        if entry.pyproject
    }
    for path in sorted(changed_paths):
        entry = local_by_pyproject.get(path)
        if entry is None:
            return InstallPlan(
                "Full Python environment sync",
                f"tracked input changed outside a known local project: {path}",
                (full_command,),
            )
        commands.append((python, "-m", "pip", "install", *entry.install_args))

    if removed_paths:
        return InstallPlan(
            "Full Python environment sync",
            "tracked input files disappeared; using the full requirements installer.",
            (full_command,),
        )

    deduped_commands = tuple(dict.fromkeys(commands))
    if deduped_commands:
        return InstallPlan(
            "Selective Python environment sync",
            f"installing {len(deduped_commands)} stale local requirement(s).",
            deduped_commands,
            tuple(warnings),
        )
    return InstallPlan(
        "Stamp-only Python environment sync",
        "requirements metadata changed without a pip-installable requirement change.",
        (),
        tuple(warnings),
    )


def requirement_delta_plan(previous_entries: object, current_entries: tuple[RequirementEntry, ...]) -> InstallPlan | None:
    parsed_previous = previous_requirement_entries(previous_entries)
    if parsed_previous is None:
        return None

    previous_by_key = {entry.key: entry for entry in parsed_previous}
    current_by_key = {entry.key: entry for entry in current_entries}
    changed_current = [
        entry
        for key, entry in current_by_key.items()
        if previous_by_key.get(key) is None or previous_by_key[key].install_args != entry.install_args or previous_by_key[key].line != entry.line
    ]
    removed_keys = set(previous_by_key) - set(current_by_key)
    if any(not entry.is_local or not entry.pyproject for entry in changed_current):
        return None

    warnings = []
    if removed_keys:
        warnings.append("Removed requirements are not uninstalled automatically; pip install -r had the same limitation.")
    commands = tuple((entry.install_args) for entry in changed_current)
    return InstallPlan(
        "Requirement delta",
        "requirements file changed only in local editable requirements.",
        commands,
        tuple(warnings),
    )


def previous_requirement_entries(value: object) -> tuple[RequirementEntry, ...] | None:
    if not isinstance(value, list):
        return None
    entries: list[RequirementEntry] = []
    for item in value:
        if not isinstance(item, dict):
            return None
        raw_install_args = item.get("install_args")
        if not isinstance(raw_install_args, list) or not all(isinstance(arg, str) for arg in raw_install_args):
            return None
        line = item.get("line")
        key = item.get("key")
        if not isinstance(line, str) or not isinstance(key, str):
            return None
        entries.append(
            RequirementEntry(
                line=line,
                install_args=tuple(raw_install_args),
                is_local=bool(item.get("is_local")),
                key=key,
                target_path=item.get("target_path") if isinstance(item.get("target_path"), str) else None,
                pyproject=item.get("pyproject") if isinstance(item.get("pyproject"), str) else None,
            )
        )
    return tuple(entries)


def inputs_by_path(value: object) -> dict[str, str]:
    if not isinstance(value, list):
        return {}
    inputs: dict[str, str] = {}
    for item in value:
        if not isinstance(item, dict):
            continue
        path = item.get("path")
        digest = item.get("sha256")
        if isinstance(path, str) and isinstance(digest, str):
            inputs[path] = digest
    return inputs


def local_requirement_entries(requirements: Path) -> tuple[RequirementEntry, ...]:
    return tuple(entry for entry in parse_requirement_entries(requirements) if entry.is_local and entry.pyproject)


def parse_requirement_entries(requirements: Path) -> tuple[RequirementEntry, ...]:
    entries: list[RequirementEntry] = []
    for raw_line in requirements.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or line.startswith(("-r ", "--requirement")):
            continue
        editable = editable_target(line)
        target = editable or local_path_target(line)
        if target is None:
            entries.append(RequirementEntry(line=line, install_args=(line,), is_local=False, key=f"requirement:{line}"))
            continue
        path = resolve_requirement_path(requirements.parent, target)
        if path is None:
            entries.append(RequirementEntry(line=line, install_args=(line,), is_local=False, key=f"requirement:{line}"))
            continue
        pyproject = path / "pyproject.toml"
        resolved_pyproject = str(pyproject.resolve()) if pyproject.exists() else None
        install_target = absolute_local_target(requirements.parent, target)
        install_args = ("-e", install_target) if editable is not None else (install_target,)
        key = f"local:{resolved_pyproject or path.resolve()}"
        entries.append(
            RequirementEntry(
                line=line,
                install_args=install_args,
                is_local=True,
                key=key,
                target_path=str(path.resolve()),
                pyproject=resolved_pyproject,
            )
        )
    return tuple(entries)


def local_pyprojects(local_requirements: tuple[RequirementEntry, ...]) -> tuple[Path, ...]:
    projects = {Path(entry.pyproject) for entry in local_requirements if entry.pyproject}
    return tuple(sorted(projects))


def editable_target(line: str) -> str | None:
    if line.startswith("-e "):
        return line[3:].strip()
    if line.startswith("--editable "):
        return line.removeprefix("--editable ").strip()
    if line.startswith("--editable="):
        return line.removeprefix("--editable=").strip()
    return None


def local_path_target(line: str) -> str | None:
    if line.startswith(("./", "../", "/")):
        return line
    return None


def resolve_requirement_path(base: Path, target: str) -> Path | None:
    cleaned, _suffix = split_local_target(target)
    if "://" in cleaned or cleaned.startswith("git+"):
        return None
    if not cleaned:
        return None
    return (base / cleaned).resolve()


def absolute_local_target(base: Path, target: str) -> str:
    path_part, suffix = split_local_target(target)
    return f"{(base / path_part).resolve()}{suffix}"


def split_local_target(target: str) -> tuple[str, str]:
    cleaned = target.split("#", 1)[0].strip()
    fragment = target[len(cleaned) :].strip()
    extras = ""
    if "[" in cleaned:
        cleaned, extra_part = cleaned.split("[", 1)
        extras = f"[{extra_part}"
    return cleaned.strip(), f"{extras}{fragment}"


def add_file(hasher: hashlib._Hash, inputs: list[dict[str, str]], path: Path) -> None:
    data = path.read_bytes()
    digest = hashlib.sha256(data).hexdigest()
    add_text(hasher, str(path), digest)
    inputs.append({"path": str(path), "sha256": digest})


def add_text(hasher: hashlib._Hash, key: str, value: str) -> None:
    hasher.update(key.encode("utf-8"))
    hasher.update(b"\0")
    hasher.update(value.encode("utf-8"))
    hasher.update(b"\0")


def load_stamp(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return {}
    return payload if isinstance(payload, dict) else {}


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


def format_command(command: tuple[str, ...]) -> str:
    return " ".join(shlex.quote(item) for item in command)


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