"""Turn the three input shapes (package directory, unified diff, snippet) into analysis units.

A unit is one file's text plus an optional focus set of 1-based line numbers. In diff mode the text is the new side
of the file as far as the diff shows it (context and added lines at their real line numbers, unknown regions blank)
and the focus is the added lines; behaviours are reported only where they touch the focus.
"""
from __future__ import annotations

import os
import re
from dataclasses import dataclass, field
from pathlib import Path

MAX_PARSE_BYTES = 1_000_000

EXT_LANG = {
    ".py": "python", ".pyw": "python", ".pth": "pth",
    ".js": "javascript", ".mjs": "javascript", ".cjs": "javascript", ".jsx": "javascript",
    ".ts": "typescript", ".mts": "typescript", ".cts": "typescript", ".tsx": "tsx",
    ".vue": "sfc", ".svelte": "sfc",
    ".sh": "shell", ".bash": "shell", ".zsh": "shell", ".ksh": "shell", ".command": "shell",
    ".yml": "yaml", ".yaml": "yaml", ".json": "json", ".toml": "toml", ".cfg": "ini", ".ini": "ini",
}
DOC_EXT = {".md", ".rst", ".txt", ".adoc", ".markdown", ".svg", ".html", ".htm", ".css", ".scss", ".lock", ".csv",
           ".ipynb", ".sql", ".xml", ".po", ".pot"}
SKIP_DIRS = {".git", "node_modules", "__pycache__", ".venv", "venv", ".tox", ".mypy_cache", ".pytest_cache"}
BINARY_EXT = {".node", ".so", ".dll", ".exe", ".dylib", ".pyd", ".pyc", ".pyo", ".wasm", ".bin", ".o", ".a",
              ".class", ".jar", ".dat"}
ARCHIVE_EXT = (".zip", ".tar", ".tgz", ".tar.gz", ".gz", ".bz2", ".xz", ".7z", ".rar", ".whl", ".egg", ".jar",
               ".nupkg", ".crx", ".vsix", ".tar.bz2", ".tar.xz", ".zst")


@dataclass
class Unit:
    path: str                       # relative, forward slashes
    lang: str                       # python|javascript|typescript|tsx|shell|yaml|json|... |other|doc|binary
    text: str = ""
    focus: set | None = None        # None = every line is in scope
    status: str = "file"            # file | added | modified | deleted | snippet
    size: int = 0
    raw_head: bytes = b""           # first bytes for binary/archive typing (package mode)
    binary_note: str = ""           # diff mode: "binary patch" etc.

    def in_focus(self, a: int, b: int | None = None) -> bool:
        if self.focus is None:
            return True
        b = a if b is None else b
        return any(x in self.focus for x in range(a, b + 1)) if b - a < 400 else any(a <= x <= b for x in self.focus)

    def first_focus(self) -> int:
        return min(self.focus) if self.focus else 1


def lang_of(path: str, text: str = "") -> str:
    p = path.lower()
    name = p.rsplit("/", 1)[-1]
    if name in ("dockerfile", "containerfile") or name.endswith(".dockerfile") or name.startswith("dockerfile."):
        return "dockerfile"
    if name in ("makefile", "gnumakefile") or name.endswith(".mk"):
        return "makefile"
    if name in (".npmrc", ".yarnrc", "pip.conf", "pip.ini", ".pypirc", ".netrc", ".env") or name.startswith(".env."):
        return "config"
    if name.startswith(".") and name.endswith("ignore") or name in ("license", "licence", "copying", "notice", "authors",
                                                                    "changelog", "changes", "codeowners", ".gitattributes",
                                                                    ".editorconfig", ".mailmap"):
        return "doc"
    if re.fullmatch(r"(requirements|constraints)[\w.-]*\.(txt|in)", name):
        return "requirements"
    for ext in sorted(EXT_LANG, key=len, reverse=True):
        if p.endswith(ext):
            return EXT_LANG[ext]
    if any(p.endswith(e) for e in BINARY_EXT):
        return "binary"
    if any(p.endswith(e) for e in DOC_EXT):
        return "doc"
    first = text.split("\n", 1)[0] if text else ""
    if first.startswith("#!"):
        if "python" in first:
            return "python"
        if "node" in first or "deno" in first or "bun" in first:
            return "javascript"
        if re.search(r"\b(ba|z|k|da)?sh\b", first):
            return "shell"
    if "/.husky/" in "/" + p or "/.git/hooks/" in "/" + p:
        return "shell"
    return "other"


def guess_snippet_lang(text: str) -> str:
    first = text.lstrip().split("\n", 1)[0]
    if first.startswith("#!"):
        return lang_of("snippet", first)
    py = len(re.findall(r"^\s*(def |import \w|from [\w.]+ import|print\(|with .+:\s*$|if __name__)", text, re.M))
    js = len(re.findall(r"\b(const|let|var|function|=>|require\(|console\.|await |import .+ from )", text))
    ts = len(re.findall(r":\s*(string|number|boolean|any|Promise<)|\binterface\s+\w+|\btype\s+\w+\s*=", text))
    sh = len(re.findall(r"^\s*(\$ )?(curl|wget|cd|ls|echo|export|sudo|apt|npm|pip|git|chmod|rm|mkdir|cat|grep)\b",
                        text, re.M)) + text.count("&&") + text.count(" | ")
    best = max((py, "python"), (js + ts, "typescript" if ts else "javascript"), (sh, "shell"))
    return best[1] if best[0] else "shell"


def read_package(root: Path, include_node_modules: bool = False):
    units = []
    root = root.resolve()
    for dirpath, dirnames, filenames in os.walk(root, followlinks=False):
        dirnames[:] = sorted(d for d in dirnames if include_node_modules or d not in SKIP_DIRS)
        for fn in sorted(filenames):
            full = Path(dirpath) / fn
            rel = full.relative_to(root).as_posix()
            if full.is_symlink():
                units.append(Unit(rel, "other", text="", size=0, binary_note="symlink -> " + os.readlink(full)))
                continue
            try:
                size = full.stat().st_size
                with open(full, "rb") as fh:
                    head = fh.read(MAX_PARSE_BYTES + 1)
            except OSError:
                continue
            u = Unit(rel, "other", size=size, raw_head=head[:4096])
            if b"\x00" in head[:8192]:
                u.lang = "binary"
            else:
                u.text = head[:MAX_PARSE_BYTES].decode("utf-8", errors="replace")
                u.lang = lang_of(rel, u.text)
                if u.lang == "binary":
                    u.lang = "other"
            units.append(u)
    return units


_HUNK = re.compile(r"^@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@")


def _strip_prefix(p: str) -> str:
    p = p.strip().split("\t", 1)[0]
    if p.startswith('"') and p.endswith('"'):
        p = p[1:-1]
    if p in ("/dev/null",):
        return p
    return p[2:] if p[:2] in ("a/", "b/") else p


def parse_diff(diff: str):
    """Return units for every file in a unified (git) diff. Only the new side is kept."""
    files = []
    cur = None

    def start(path):
        nonlocal cur
        cur = {"path": path, "lines": {}, "added": set(), "status": "modified", "binary": "", "old": path}
        files.append(cur)

    new_ln = 0
    old_left = new_left = 0          # remaining hunk body lines, so "--- x" inside a hunk is a removed line
    git_header = False
    for raw in diff.splitlines():
        if cur is not None and (old_left > 0 or new_left > 0):
            tag, body = raw[:1], raw[1:]
            if tag == "\\":
                continue
            if tag == "+":
                cur["lines"][new_ln] = body
                cur["added"].add(new_ln)
                new_ln += 1
                new_left -= 1
                continue
            if tag == "-":
                old_left -= 1
                continue
            if tag == " " or raw == "":
                cur["lines"][new_ln] = body
                new_ln += 1
                new_left -= 1
                old_left -= 1
                continue
            old_left = new_left = 0      # malformed hunk: fall through to header handling
        if raw.startswith("diff --git "):
            m = re.match(r'diff --git "?a/(.+?)"? "?b/(.+?)"?$', raw)
            start(m.group(2) if m else raw.split()[-1].removeprefix("b/"))
            git_header = True
            continue
        if raw.startswith("--- "):
            old = _strip_prefix(raw[4:])
            if not git_header:
                start(old)
            git_header = False
            cur["old"] = old
            if old == "/dev/null":
                cur["status"] = "added"
            continue
        if raw.startswith("+++ "):
            new = _strip_prefix(raw[4:])
            if cur is None:
                start(new)
            if new == "/dev/null":
                cur["status"] = "deleted"
            else:
                cur["path"] = new
            continue
        m = _HUNK.match(raw)
        if m and cur is not None:
            new_ln = int(m.group(3))
            old_left = int(m.group(2)) if m.group(2) is not None else 1
            new_left = int(m.group(4)) if m.group(4) is not None else 1
            git_header = False
            continue
        if cur is None:
            continue
        if raw.startswith("new file mode"):
            cur["status"] = "added"
        elif raw.startswith("deleted file mode"):
            cur["status"] = "deleted"
        elif raw.startswith("rename to "):
            cur["path"] = raw[len("rename to "):]
        elif raw.startswith("Binary files ") or raw.startswith("GIT binary patch"):
            cur["binary"] = raw.strip()
        elif raw.startswith("literal ") and cur.get("binary") and "bytes" not in cur["binary"]:
            cur["binary"] += f" ({raw.split()[1]} bytes)"
    units = []
    for f in files:
        path = f["path"]
        if f["binary"]:
            if f["status"] != "deleted":
                units.append(Unit(path, "binary", status=f["status"], binary_note=f["binary"]))
            continue
        if f["status"] == "deleted":
            continue
        n = max(f["lines"], default=0)
        text = "\n".join(f["lines"].get(i, "") for i in range(1, n + 1))
        lang = lang_of(path, f["lines"].get(1, ""))
        units.append(Unit(path, lang, text=text, focus=set(f["added"]), status=f["status"],
                          size=sum(len(f["lines"][i]) + 1 for i in f["added"])))
    return units


def snippet_unit(text: str, lang: str | None = None, name: str = "snippet") -> Unit:
    lang = {"js": "javascript", "ts": "typescript", "py": "python", "sh": "shell", "bash": "shell"}.get(lang or "", lang)
    return Unit(name, lang or guess_snippet_lang(text), text=text, status="snippet", size=len(text))