#!/usr/bin/env python3
"""
Update the crate versions across the workspace.
"""

from __future__ import annotations

import re
import sys
from pathlib import Path
from typing import Iterable, List, Tuple


def usage(script_name: str) -> None:
    print(f"Usage: {script_name} <new-version>")
    print(f"Example: {script_name} 0.2.0")
    sys.exit(1)


def find_cargo_tomls(root_dir: Path) -> Tuple[List[Path], List[str]]:
    """Return every Cargo.toml we should touch and collect crate names."""
    cargo_tomls: List[Path] = [root_dir / "Cargo.toml"]
    crate_names: List[str] = []

    crates_dir = root_dir / "crates"
    if crates_dir.is_dir():
        for crate_dir in sorted(p for p in crates_dir.iterdir() if p.is_dir()):
            cargo_toml = crate_dir / "Cargo.toml"
            if not cargo_toml.is_file():
                continue
            cargo_tomls.append(cargo_toml)

            crate_name = extract_crate_name(cargo_toml)
            if crate_name:
                crate_names.append(crate_name)

    return cargo_tomls, crate_names


def extract_crate_name(cargo_toml: Path) -> str | None:
    """Grab the crate name from the [package] section."""
    name_pattern = re.compile(r'^\s*name\s*=\s*"([^"]+)"\s*$')
    in_package = False
    for line in cargo_toml.read_text(encoding="utf-8").splitlines():
        stripped = line.strip()
        if stripped.startswith("[") and stripped.endswith("]"):
            in_package = stripped.strip("[]") == "package"
            continue
        if in_package:
            match = name_pattern.match(line)
            if match:
                return match.group(1)
    return None


def update_package_version(lines: List[str], new_version: str) -> None:
    """Rewrite the version in the [package] section."""
    in_package = False
    version_pattern = re.compile(r'^(\s*version\s*=\s*")([^"]*)(".*)$')

    for idx, line in enumerate(lines):
        stripped = line.strip()
        if stripped.startswith("[") and stripped.endswith("]"):
            in_package = stripped.strip("[]") == "package"
            continue
        if not in_package:
            continue

        match = version_pattern.match(line)
        if match:
            lines[idx] = f"{match.group(1)}{new_version}{match.group(3)}"
            break


def normalize_section_name(section: str) -> str:
    """Remove surrounding brackets and whitespace."""
    return section.strip().strip("[]").strip()


def update_table_dependency_versions(
    lines: List[str], crate_names: Iterable[str], new_version: str
) -> None:
    """
    Update `[dependencies.foo] version = "..."` style sections.
    """
    crate_name_set = set(crate_names)
    current_section: str | None = None

    version_pattern = re.compile(r'^(\s*version\s*=\s*")([^"]*)(".*)$')

    for idx, line in enumerate(lines):
        stripped = line.strip()
        if stripped.startswith("[") and stripped.endswith("]"):
            current_section = normalize_section_name(stripped)
            continue

        if not current_section:
            continue

        parts = [part.strip().strip("\"'") for part in current_section.split(".")]
        if not parts:
            continue
        maybe_crate = parts[-1]
        if maybe_crate not in crate_name_set:
            continue

        match = version_pattern.match(line)
        if match:
            lines[idx] = f"{match.group(1)}{new_version}{match.group(3)}"


def update_inline_dependency_versions(
    text: str, crate_names: Iterable[str], new_version: str
) -> str:
    """
    Rewrite inline dependency tables:
    `crate_name = { path = "...", version = "x.y.z", ... }`
    """
    for crate_name in crate_names:
        pattern = re.compile(
            rf'({re.escape(crate_name)}\s*=\s*{{[^}}]*?\bversion\s*=\s*")([^"]*)(")',
            re.MULTILINE,
        )
        text = pattern.sub(
            lambda match: f"{match.group(1)}{new_version}{match.group(3)}", text
        )
    return text


def update_file(cargo_toml: Path, crate_names: Iterable[str], new_version: str) -> None:
    print(f"Updating {cargo_toml}")

    original_text = cargo_toml.read_text(encoding="utf-8")
    has_trailing_newline = original_text.endswith("\n")
    lines = original_text.splitlines()

    update_package_version(lines, new_version)
    update_table_dependency_versions(lines, crate_names, new_version)

    updated_text = "\n".join(lines)
    if has_trailing_newline:
        updated_text += "\n"

    updated_text = update_inline_dependency_versions(
        updated_text, crate_names, new_version
    )

    if updated_text != original_text:
        cargo_toml.write_text(updated_text, encoding="utf-8")

    print(f"✓ Updated {cargo_toml}")


def main() -> None:
    if len(sys.argv) != 2:
        usage(Path(sys.argv[0]).name)

    new_version = sys.argv[1]
    root_dir = Path(__file__).resolve().parents[1]

    print(f"Updating all crates to version {new_version}")

    cargo_tomls, crate_names = find_cargo_tomls(root_dir)
    for cargo_toml in cargo_tomls:
        update_file(cargo_toml, crate_names, new_version)

    print(f"All crates updated to version {new_version}")


if __name__ == "__main__":
    main()
