import configparser
import importlib.metadata
import json
import os
import re
import subprocess
import sys
from pathlib import Path
import tomllib

from packaging.markers import default_environment
from packaging.requirements import Requirement


def normalise(name: str) -> str:
    """PEP 503 normalisation: collapse [-_.] runs to _, lowercase."""
    return re.sub(r"[-_.]+", "_", name).lower()


def find_active_python() -> Path:
    """
    Return (python_path, site_packages_dirs) for the environment to measure.

    Priority:
      1. $VIRTUAL_ENV environment variable (set by `activate`, uv, poetry shell, etc.)
      2. sys.executable (the interpreter running this script, e.g. uv tool install)
    """
    venv = os.environ.get("VIRTUAL_ENV")
    if venv:
        venv_path = Path(venv)
        for rel in ["bin/python", "bin/python3", "Scripts/python.exe"]:
            p = venv_path / rel
            if p.exists():
                return p

    return Path(sys.executable)


def _site_packages(python: Path | None) -> list[Path]:
    if python is None:
        python = find_active_python()

    return _site_packages_for(python)


def _site_packages_for(python: Path) -> list[Path]:
    """Ask an interpreter for its site-packages directories."""
    try:
        out = subprocess.check_output(
            [
                str(python),
                "-c",
                "import site, json; print(json.dumps(site.getsitepackages()))",
            ],
            stderr=subprocess.DEVNULL,
            text=True,
        )
        return [Path(p) for p in json.loads(out.strip())]
    except Exception:
        return []


def build_dist_map(
    site_packages: list[Path],
) -> dict[str, importlib.metadata.Distribution]:
    """Return normalised-name → Distribution for every package in site_packages."""
    mapping: dict[str, importlib.metadata.Distribution] = {}
    search_path = [str(p) for p in site_packages if p.exists()]
    dists = (
        importlib.metadata.distributions(path=search_path)
        if search_path
        else importlib.metadata.distributions()
    )
    for dist in dists:
        try:
            name = dist.metadata["Name"]
            if name:
                mapping[normalise(name)] = dist
        except Exception:
            pass
    return mapping


# ── Disk-size calculation ─────────────────────────────────────────────────────


def _walk_paths(paths: list[Path]) -> int:
    """Sum sizes of all regular files under the given paths, dedup by inode."""
    total = 0
    seen: set[int] = set()
    for root in paths:
        items = root.rglob("*") if root.is_dir() else [root]
        for p in items:
            try:
                if p.is_file():
                    st = p.stat()
                    if st.st_ino not in seen:
                        seen.add(st.st_ino)
                        total += st.st_size
            except OSError:
                pass
    return total


def _dist_info_path(dist: importlib.metadata.Distribution) -> Path | None:
    """Return the .dist-info / .egg-info directory for a distribution."""
    # _path is a private but stable attribute on PathDistribution (CPython impl)
    raw = getattr(dist, "_path", None)
    if raw is not None:
        p = Path(raw)
        if p.exists():
            return p
    return None


def _package_dirs(dist_info: Path, pkg_name: str) -> list[Path]:
    """
    Given a .dist-info dir, return the list of actual package dirs/files
    that belong to this distribution (via top_level.txt or name guessing).
    """
    site = dist_info.parent
    candidates: list[Path] = [dist_info]

    tl = dist_info / "top_level.txt"
    if tl.exists():
        tops = tl.read_text().split()
    else:
        normed = normalise(pkg_name)
        tops = list(dict.fromkeys([pkg_name, pkg_name.lower(), normed]))

    for top in tops:
        for p in [site / top, site / f"{top}.py"]:
            if p.exists():
                candidates.append(p)

    return candidates


def get_package_size(dist: importlib.metadata.Distribution) -> int:
    """
    Return on-disk size (bytes) for an installed distribution.

    Strategy 1 – RECORD file (standard pip installs).
      Most packages ship a RECORD that lists every installed file.

    Strategy 2 – walk dist-info + top_level dirs (no RECORD).
      Packages like PyTorch omit the RECORD file for size reasons.
      We find the .dist-info dir via the private _path attribute, read
      top_level.txt to get the real package directory names, then walk them.

    Strategy 3 – `pip show --files` subprocess fallback.
      When _path is unavailable or wrong (rare edge cases).
    """
    name = dist.metadata.get("Name", "?")

    # ── Strategy 1: RECORD ───────────────────────────────────────────────────
    files = dist.files
    if files:
        paths: list[Path] = []
        for f in files:
            try:
                p = f.locate()
                if p.exists():
                    paths.append(p)
            except (OSError, AttributeError):
                pass
        if paths:
            size = _walk_paths(paths)
            if size > 0:
                return size

    # ── Strategy 2: dist-info + top_level walk ───────────────────────────────
    dist_info = _dist_info_path(dist)
    if dist_info:
        dirs = _package_dirs(dist_info, name)
        size = _walk_paths(dirs)
        if size > 0:
            return size

    # ── Strategy 3: pip show --files ─────────────────────────────────────────
    try:
        out = subprocess.check_output(
            [sys.executable, "-m", "pip", "show", "--files", name],
            stderr=subprocess.DEVNULL,
            text=True,
            timeout=15,
        )
        location: str | None = None
        file_lines: list[str] = []
        in_files = False
        for line in out.splitlines():
            if line.startswith("Location:"):
                location = line.split(":", 1)[1].strip()
                in_files = False
            elif line.startswith("Files:"):
                in_files = True
            elif in_files:
                stripped = line.strip()
                if stripped:
                    file_lines.append(stripped)
                else:
                    in_files = False
        if location and file_lines:
            loc = Path(location)
            paths = [loc / f for f in file_lines if (loc / f).exists()]
            size = _walk_paths(paths)
            return size
    except Exception:
        pass

    return 0


def _req_name(req_str: str) -> str | None:
    """Extract the package name from a requirement string."""
    return Requirement(req_str).name


def parse_project_deps(project_dir: Path) -> list[str]:
    """
    Return top-level dependency names declared by the project at project_dir.

    Sources tried in order:
      1. pyproject.toml  – PEP 621 [project.dependencies]
      2. pyproject.toml  – Poetry  [tool.poetry.dependencies]
      3. setup.cfg       – [options] install_requires
      4. requirements.txt
    """

    toml_path = project_dir / "pyproject.toml"
    if toml_path.exists():
        try:
            with open(toml_path, "rb") as f:
                data = tomllib.load(f)

            pep621: list[str] = data.get("project", {}).get("dependencies", [])
            if pep621:
                names = [_req_name(d) for d in pep621]
                result = [n for n in names if n]
                if result:
                    return result

            poetry: dict = (
                data.get("tool", {}).get("poetry", {}).get("dependencies", {})
            )
            if poetry:
                return [k for k in poetry if normalise(k) != "python"]

        except Exception as e:
            print(f"  Warning: could not parse pyproject.toml: {e}", file=sys.stderr)

    cfg_path = project_dir / "setup.cfg"
    if cfg_path.exists():
        try:
            cfg = configparser.ConfigParser()
            cfg.read(cfg_path)
            raw = cfg.get("options", "install_requires", fallback="")
            lines = [line.strip() for line in raw.splitlines() if line.strip()]
            names = [_req_name(line) for line in lines]
            result = [n for n in names if n]
            if result:
                return result
        except Exception as e:
            print(f"  Warning: could not parse setup.cfg: {e}", file=sys.stderr)

    req_path = project_dir / "requirements.txt"
    if req_path.exists():
        with open(req_path) as f:
            lines = [
                line.strip()
                for line in f
                if line.strip() and not line.startswith(("#", "-"))
            ]
        names = [_req_name(line) for line in lines]
        result = [n for n in names if n]
        if result:
            return result

    return []


def get_requires(dist: importlib.metadata.Distribution) -> list[str]:
    """Return direct dependency names for an installed distribution."""
    requires = dist.requires
    if not requires:
        return []
    names: list[str] = []
    for req_str in requires:
        try:
            req = Requirement(req_str)
            # Skip deps whose markers don't match this environment
            if req.marker and not req.marker.evaluate(default_environment()):
                continue
            names.append(req.name)
        except Exception:
            name = _req_name(req_str)
            if name:
                names.append(name)
    return names


def parse_project_info(project_dir: Path) -> tuple[str, str]:
    """
    Return (name, version) for the project itself.
    Name and version are read from pyproject.toml / setup.cfg, falling
    back to the directory name / "dev" if not found.
    """
    name = project_dir.name
    version = "dev"

    # Try to read name + version from pyproject.toml
    toml_path = project_dir / "pyproject.toml"
    if toml_path.exists():
        try:
            with open(toml_path, "rb") as f:
                data = tomllib.load(f)
            proj = data.get("project") or data.get("tool", {}).get("poetry", {})
            name = proj.get("name", name)
            version = proj.get("version", version)
        except Exception:
            pass

    # Try setup.cfg as fallback
    if version == "dev":
        cfg_path = project_dir / "setup.cfg"
        if cfg_path.exists():
            try:
                cfg = configparser.ConfigParser()
                cfg.read(cfg_path)
                name = cfg.get("metadata", "name", fallback=name)
                version = cfg.get("metadata", "version", fallback=version)
            except Exception:
                pass

    return name, version


def resolve_tree(
    roots: list[str],
    dist_map: dict[str, importlib.metadata.Distribution],
) -> list[dict]:
    visited: set[str] = set()

    def resolve(name: str) -> dict:
        key = normalise(name)
        dist = dist_map.get(key)
        node: dict = {
            "name": name,
            "version": dist.metadata["Version"] if dist else "?",
            "sizeBytes": get_package_size(dist) if dist else 0,
            "dependencies": [],
        }
        if key in visited:
            return node
        visited.add(key)
        if dist:
            for child in get_requires(dist):
                resolved = resolve(child)
                node["dependencies"].append(resolved)
                node["sizeBytes"] += resolved["sizeBytes"]
        return node

    return [resolve(r) for r in roots]


class NoDependenciesException(Exception):
    pass


def calculate_tree(project_dir: Path, python: Path | None) -> dict:
    site_pkgs = _site_packages(python)

    dist_map = build_dist_map(site_pkgs)

    proj_name, proj_version = parse_project_info(project_dir)
    deps = parse_project_deps(project_dir)
    if not deps:
        raise NoDependenciesException

    tree = resolve_tree(deps, dist_map)

    proj_size = sum((dep["sizeBytes"] for dep in tree))

    return {
        "name": proj_name,
        "version": proj_version,
        "sizeBytes": proj_size,
        "dependencies": tree,
    }
