from __future__ import annotations

import ast
import re
from pathlib import Path
from packaging.requirements import Requirement


SETUP_PATH = Path("setup.py")
FREEZE_PATH = Path("ocr-pip-freeze.txt")


# Packages you usually do NOT want in install_requires.
# These are dev/test/build tools or environment-specific packages.
EXCLUDE = {
    "pip",
    "setuptools",
    "wheel",
    "pytest",
    "black",
    "flake8",
    "mypy",
    "ipython",
    "jupyter",
    "notebook",
    "twine",
    "build",
}


# Optional: keep these unpinned even if pip freeze pins them.
# Good for libraries, because strict pins can break downstream users.
UNPIN = {
    "requests",
    "pdf2image",
    "PyPDF2",
    "easyocr",
    "pytesseract",
    "lxml",
    "spacy",
    "abstract_utilities",
    "abstract_hugpy",
}


def normalize_name(name: str) -> str:
    return name.lower().replace("_", "-")


def parse_requirement_name(req: str) -> str | None:
    req = req.strip()
    if not req or req.startswith("#"):
        return None

    # Skip editable installs, local paths, git URLs, etc.
    if req.startswith("-e ") or " @ file://" in req or req.startswith("git+"):
        return None

    try:
        return normalize_name(Requirement(req).name)
    except Exception:
        # Handles simple freeze lines like package==1.2.3
        match = re.match(r"^([A-Za-z0-9_.-]+)", req)
        return normalize_name(match.group(1)) if match else None


def load_freeze_requirements(path: Path) -> dict[str, str]:
    requirements = {}

    for line in path.read_text(encoding="utf-8").splitlines():
        req = line.strip()
        name = parse_requirement_name(req)

        if not name:
            continue

        if name in {normalize_name(x) for x in EXCLUDE}:
            continue

        if name in {normalize_name(x) for x in UNPIN}:
            req = Requirement(req).name

        requirements[name] = req

    return requirements


def extract_install_requires(setup_text: str) -> list[str]:
    tree = ast.parse(setup_text)

    for node in ast.walk(tree):
        if isinstance(node, ast.Call):
            func = node.func
            is_setup = (
                isinstance(func, ast.Attribute)
                and func.attr == "setup"
            ) or (
                isinstance(func, ast.Name)
                and func.id == "setup"
            )

            if not is_setup:
                continue

            for keyword in node.keywords:
                if keyword.arg == "install_requires":
                    return ast.literal_eval(keyword.value)

    raise RuntimeError("Could not find install_requires=[...] in setup.py")


def replace_install_requires(setup_text: str, requirements: list[str]) -> str:
    rendered = "install_requires=[\n"
    for req in requirements:
        rendered += f"        {req!r},\n"
    rendered += "        ]"

    pattern = re.compile(
        r"install_requires\s*=\s*\[(?:.|\n)*?\n\s*\]",
        re.MULTILINE,
    )

    new_text, count = pattern.subn(rendered, setup_text, count=1)

    if count != 1:
        raise RuntimeError("Could not replace install_requires block cleanly")

    return new_text


def main() -> None:
    setup_text = SETUP_PATH.read_text(encoding="utf-8")

    current_requirements = extract_install_requires(setup_text)
    freeze_requirements = load_freeze_requirements(FREEZE_PATH)

    merged = {}

    for req in current_requirements:
        name = parse_requirement_name(req)
        if name:
            merged[name] = req

    for name, req in freeze_requirements.items():
        merged.setdefault(name, req)

    final_requirements = sorted(
        merged.values(),
        key=lambda value: normalize_name(parse_requirement_name(value) or value),
    )

    updated = replace_install_requires(setup_text, final_requirements)
    SETUP_PATH.write_text(updated, encoding="utf-8")

    print(f"Merged {len(final_requirements)} requirements into {SETUP_PATH}")


if __name__ == "__main__":
    main()
