#!/usr/bin/env python3
"""Check for outdated dependencies in pyproject.toml and optionally apply upgrades.

Usage:
    python check_outdated_deps.py check                            # Output JSON matrix of upgrade jobs
    python check_outdated_deps.py table                            # Human-readable outdated table
    python check_outdated_deps.py apply <pkg> <old_ver> <new_ver>  # Update a single dep in-place
    python check_outdated_deps.py apply-group <group_name>         # Bump a coupled group as a set
    python check_outdated_deps.py resolve                          # Resolution gate (uv); exit 1 on conflict
"""

import fnmatch
import json
import re
import subprocess
import sys
import tomllib
import urllib.request
from pathlib import Path

from packaging.specifiers import SpecifierSet
from packaging.version import Version

PYPROJECT_PATH = Path(__file__).resolve().parents[2] / "pyproject.toml"
PYPI_URL = "https://pypi.org/pypi/{}/json"
# Python version used to resolve dependencies. Matches the "Test CCA against
# local wheel" job so the resolution gate reproduces exactly what CI checks.
RESOLVE_PYTHON_VERSION = "3.12"


def normalize_name(name: str) -> str:
    """Normalize package name per PEP 503."""
    return re.sub(r"[-_.]+", "-", name).lower()


def parse_compatible_release(dep_string: str) -> tuple[str, str, str | None] | None:
    """Parse a dependency string with a ~= specifier.

    Returns (package_name, pinned_version, extras) or None if not a compatible release pin.
    Extras like [hf_xet] or [openai] are separated out.
    """
    # Match: name[extras]~=version
    m = re.match(r"^([A-Za-z0-9_.-]+)(\[[^\]]+\])?\s*~=\s*([^\s,;]+)$", dep_string.strip())
    if not m:
        return None
    name, extras, version = m.group(1), m.group(2), m.group(3)
    return name, version, extras


def get_latest_stable_version(package_name: str, max_major: int | None = None) -> str | None:
    """Query PyPI for the latest stable (non-prerelease, non-yanked) version.

    If max_major is given, only versions with that major number are considered.
    """
    url = PYPI_URL.format(package_name)
    try:
        req = urllib.request.Request(url, headers={"Accept": "application/json"})
        with urllib.request.urlopen(req, timeout=15) as resp:
            data = json.loads(resp.read())
    except Exception as e:
        print(f"  Warning: failed to query PyPI for {package_name}: {e}", file=sys.stderr)
        return None

    releases = data.get("releases", {})
    stable_versions = []
    for ver_str, files in releases.items():
        try:
            v = Version(ver_str)
        except Exception:
            continue
        if v.is_prerelease or v.is_devrelease:
            continue
        # Skip yanked releases (all files yanked = yanked version)
        if files and all(f.get("yanked", False) for f in files):
            continue
        # Skip versions with no files at all
        if not files:
            continue
        if max_major is not None and v.major != max_major:
            continue
        stable_versions.append(v)

    if not stable_versions:
        return None
    return str(max(stable_versions))


def check_outdated() -> list[dict]:
    """Return list of outdated compatible-release deps."""
    with open(PYPROJECT_PATH, "rb") as f:
        pyproject = tomllib.load(f)

    deps = pyproject.get("project", {}).get("dependencies", [])
    pin_major = load_pin_major()
    outdated = []
    seen = set()

    for dep in deps:
        parsed = parse_compatible_release(dep)
        if parsed is None:
            continue
        name, current_version, extras = parsed
        # Skip duplicates (e.g. boto3 and boto3[crt] share the same base package and version)
        key = (normalize_name(name), current_version)
        if key in seen:
            continue
        seen.add(key)
        print(f"Checking {name} (current: ~={current_version})...", file=sys.stderr)
        # For pin-major packages, only look at versions within the same major.
        max_major = Version(current_version).major if normalize_name(name) in pin_major else None
        if max_major is not None:
            print(f"  (pin-major: restricting to {max_major}.x)", file=sys.stderr)
        latest = get_latest_stable_version(name, max_major=max_major)
        if latest is None:
            continue
        latest_ver = Version(latest)
        # Only flag as outdated if the latest version is OUTSIDE the ~= range.
        # e.g. ~=0.7.14 allows 0.7.16, so that's not outdated.
        # But ~=0.7.14 does NOT allow 0.8.0, so that IS outdated.
        specifier = SpecifierSet(f"~={current_version}")
        if latest_ver > Version(current_version) and latest_ver not in specifier:
            outdated.append({
                "package": name,
                "current_version": current_version,
                "latest_version": latest,
                "extras": extras,
            })
            print(f"  -> outdated: {current_version} -> {latest} (outside ~={current_version})", file=sys.stderr)
        else:
            print(f"  -> up to date (latest {latest} within ~={current_version})", file=sys.stderr)

    return outdated


def load_pin_major() -> set[str]:
    """Load the pin-major package list from [tool.upgrade-deps].pin-major."""
    with open(PYPROJECT_PATH, "rb") as f:
        pyproject = tomllib.load(f)
    names = pyproject.get("tool", {}).get("upgrade-deps", {}).get("pin-major", [])
    return {normalize_name(n) for n in names}


def load_groups() -> list[dict]:
    """Load coupled-dependency group definitions from [tool.upgrade-deps].groups."""
    with open(PYPROJECT_PATH, "rb") as f:
        pyproject = tomllib.load(f)
    return pyproject.get("tool", {}).get("upgrade-deps", {}).get("groups", [])


def group_for(package: str, groups: list[dict]) -> str | None:
    """Return the name of the group a package belongs to, or None.

    Members are glob patterns (e.g. "langchain*") matched against the
    PEP 503-normalized package name.
    """
    nname = normalize_name(package)
    for group in groups:
        for pattern in group.get("members", []):
            if fnmatch.fnmatch(nname, normalize_name(pattern)):
                return group.get("name")
    return None


def pinned_dependencies() -> list[tuple[str, str, str | None]]:
    """Return all ~= pinned deps from pyproject as (name, version, extras)."""
    with open(PYPROJECT_PATH, "rb") as f:
        pyproject = tomllib.load(f)
    pinned = []
    for dep in pyproject.get("project", {}).get("dependencies", []):
        parsed = parse_compatible_release(dep)
        if parsed is not None:
            pinned.append(parsed)
    return pinned


def resolve_versions(relax_names: set[str]) -> dict[str, str]:
    """Resolve the full dependency set with the named members relaxed to >=current.

    Returns a mapping of normalized package name -> resolved version. The members
    are loosened so the resolver is free to pick the latest mutually-compatible
    versions across the whole coupled set. Requires `uv` on PATH.
    """
    with open(PYPROJECT_PATH, "rb") as f:
        pyproject = tomllib.load(f)
    lines = []
    for dep in pyproject.get("project", {}).get("dependencies", []):
        parsed = parse_compatible_release(dep)
        if parsed is not None and normalize_name(parsed[0]) in relax_names:
            name, version, extras = parsed
            lines.append(f"{name}{extras or ''}>={version}")
        else:
            lines.append(dep)

    result = subprocess.run(
        ["uv", "pip", "compile", "-", "--python-version", RESOLVE_PYTHON_VERSION, "--no-header"],
        input="\n".join(lines),
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        print(result.stderr, file=sys.stderr)
        print("Error: failed to resolve relaxed dependency set", file=sys.stderr)
        sys.exit(1)

    resolved = {}
    for line in result.stdout.splitlines():
        m = re.match(r"^([A-Za-z0-9_.-]+)\s*==\s*([^\s;]+)", line.strip())
        if m:
            resolved[normalize_name(m.group(1))] = m.group(2)
    return resolved


def apply_group(group_name: str) -> None:
    """Bump every member of a coupled group to the latest mutually-compatible set.

    Membership is declared in config; the target versions are chosen by the
    resolver (not hard-coded). Only members whose resolved version falls OUTSIDE
    their current ~= range are re-pinned, keeping the diff minimal.
    """
    group = next((g for g in load_groups() if g.get("name") == group_name), None)
    if group is None:
        print(f"Error: no group named '{group_name}' in [tool.upgrade-deps].groups", file=sys.stderr)
        sys.exit(1)

    patterns = group.get("members", [])
    members = [
        (name, version, extras)
        for name, version, extras in pinned_dependencies()
        if any(fnmatch.fnmatch(normalize_name(name), normalize_name(p)) for p in patterns)
    ]
    if not members:
        print(f"Error: group '{group_name}' matched no ~= pinned dependencies", file=sys.stderr)
        sys.exit(1)

    relax_names = {normalize_name(name) for name, _, _ in members}
    resolved = resolve_versions(relax_names)

    changed = []
    for name, current, _extras in members:
        target = resolved.get(normalize_name(name))
        if target is None:
            continue
        # Only re-pin members that genuinely need to move (resolved version is
        # outside the current compatible-release range). Members that resolve
        # within their existing range are left untouched.
        if Version(target) > Version(current) and Version(target) not in SpecifierSet(f"~={current}"):
            apply_upgrade(name, current, target)
            changed.append((name, current, target))

    if not changed:
        print(f"Group '{group_name}': no members needed bumping.", file=sys.stderr)
        return
    print(f"Group '{group_name}' bumped:", file=sys.stderr)
    for name, current, target in changed:
        print(f"  {name}: ~={current} -> ~={target}", file=sys.stderr)


def resolve_check() -> None:
    """Resolution gate: confirm the current pyproject resolves. Exit 1 on conflict.

    Reproduces the resolution the downstream "Test CCA against local wheel" job
    performs, so an unsatisfiable upgrade is caught before a PR is opened.
    """
    result = subprocess.run(
        ["uv", "pip", "compile", str(PYPROJECT_PATH), "--python-version", RESOLVE_PYTHON_VERSION, "--no-header"],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        print(result.stderr, file=sys.stderr)
        print("Resolution gate FAILED: dependency set is unsatisfiable.", file=sys.stderr)
        sys.exit(1)
    print("Resolution gate passed: dependency set resolves cleanly.", file=sys.stderr)


def build_matrix() -> list[dict]:
    """Build the GitHub Actions job matrix from outdated deps.

    Outdated packages that belong to a coupled group are folded into a single
    grouped job; everything else becomes an individual job. Grouped jobs are
    emitted first so they survive the workflow's per-run PR cap.
    """
    groups = load_groups()
    outdated = check_outdated()

    singles, buckets = [], {}
    for dep in outdated:
        gname = group_for(dep["package"], groups)
        if gname:
            buckets.setdefault(gname, []).append(dep)
        else:
            singles.append(dep)

    entries = []
    for gname, deps in buckets.items():
        lead = next((d for d in deps if normalize_name(d["package"]) == normalize_name(gname)), None)
        if lead is None:
            lead = sorted(deps, key=lambda d: d["package"])[0]
        members = ", ".join(sorted(d["package"] for d in deps))
        entries.append({
            "kind": "group",
            "group": gname,
            "package": lead["package"],
            "current_version": lead["current_version"],
            "latest_version": lead["latest_version"],
            "extras": lead["extras"],
            "members": members,
            "branch": f"upgrade/{gname}-group-{lead['latest_version']}",
            "title": (
                f"Upgrade {gname} group to latest compatible "
                f"({lead['package']} {lead['current_version']} → {lead['latest_version']})"
            ),
        })

    for dep in singles:
        entries.append({
            "kind": "single",
            "group": "",
            "package": dep["package"],
            "current_version": dep["current_version"],
            "latest_version": dep["latest_version"],
            "extras": dep["extras"],
            "members": dep["package"],
            "branch": f"upgrade/{dep['package']}-{dep['latest_version']}",
            "title": f"Upgrade {dep['package']} from {dep['current_version']} to {dep['latest_version']}",
        })

    return entries


def apply_upgrade(package: str, old_version: str, new_version: str) -> None:
    """Update a single dependency version in pyproject.toml, preserving formatting."""
    content = PYPROJECT_PATH.read_text()

    # Build a pattern that matches the dep line for this package (with optional extras)
    # We need to match the original casing in the file, so use case-insensitive match
    # Allow flexible name matching (hyphens, underscores, dots are interchangeable per PEP 503).
    # Normalize first, then escape, so re.escape doesn't break the character class.
    flexible_name = re.sub(r"[-_.]+", "[-_.]", package)
    escaped_old = re.escape(old_version)

    pattern = re.compile(
        rf"^(\s*\"{flexible_name}(?:\[[^\]]*\])?)\s*~=\s*{escaped_old}(\")",
        re.MULTILINE | re.IGNORECASE,
    )

    new_content, count = pattern.subn(rf"\g<1>~={new_version}\2", content)
    if count == 0:
        print(f"Error: could not find {package}~={old_version} in {PYPROJECT_PATH}", file=sys.stderr)
        sys.exit(1)

    PYPROJECT_PATH.write_text(new_content)
    print(f"Updated {package} from ~={old_version} to ~={new_version}", file=sys.stderr)


def bump_patch_version() -> str:
    """Increment the patch version in pyproject.toml. Returns the new version string."""
    content = PYPROJECT_PATH.read_text()
    pattern = re.compile(r'^(version\s*=\s*")(\d+\.\d+\.)(\d+)(")', re.MULTILINE)
    match = pattern.search(content)
    if not match:
        print("Error: could not find version in pyproject.toml", file=sys.stderr)
        sys.exit(1)
    old_patch = int(match.group(3))
    new_patch = old_patch + 1
    new_version = f"{match.group(2)}{new_patch}"
    new_content = pattern.sub(rf"\g<1>{new_version}\4", content)
    PYPROJECT_PATH.write_text(new_content)
    print(f"Bumped version to {new_version}", file=sys.stderr)
    return new_version


def main() -> None:
    if len(sys.argv) < 2:
        print(__doc__, file=sys.stderr)
        sys.exit(1)

    command = sys.argv[1]

    if command == "check":
        print(json.dumps(build_matrix()))

    elif command == "table":
        outdated = check_outdated()
        if not outdated:
            print("\nAll dependencies are up to date.")
        else:
            # Calculate column widths
            pkg_w = max(len("Package"), max(len(d["package"]) for d in outdated))
            cur_w = max(len("Current"), max(len(d["current_version"]) for d in outdated))
            lat_w = max(len("Latest"), max(len(d["latest_version"]) for d in outdated))

            header = f"  {'Package':<{pkg_w}}   {'Current':>{cur_w}}   {'Latest':>{lat_w}}"
            sep = f"  {'─' * pkg_w}   {'─' * cur_w}   {'─' * lat_w}"

            print(f"\n  {len(outdated)} outdated {'dependency' if len(outdated) == 1 else 'dependencies'}:\n")
            print(header)
            print(sep)
            for d in outdated:
                print(f"  {d['package']:<{pkg_w}}   {d['current_version']:>{cur_w}}   {d['latest_version']:>{lat_w}}")
            print()

    elif command == "apply":
        if len(sys.argv) != 5:
            print("Usage: check_outdated_deps.py apply <package> <old_version> <new_version>", file=sys.stderr)
            sys.exit(1)
        apply_upgrade(sys.argv[2], sys.argv[3], sys.argv[4])

    elif command == "apply-group":
        if len(sys.argv) != 3:
            print("Usage: check_outdated_deps.py apply-group <group_name>", file=sys.stderr)
            sys.exit(1)
        apply_group(sys.argv[2])

    elif command == "resolve":
        resolve_check()

    elif command == "bump-version":
        new_version = bump_patch_version()
        print(new_version)

    else:
        print(f"Unknown command: {command}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
