"""Build source-only Python reference pages with Griffe, without importing FastVideo."""
from __future__ import annotations

import html
import ast
import json
import re
import runpy
import shutil
import shlex
from pathlib import Path
from urllib.parse import urljoin, urlsplit, urlunsplit

import griffe
import markdown
import yaml

ROOT = Path(__file__).resolve().parents[2]
SITE = ROOT / "docs-site"
GENERATED = SITE / ".generated"


def source_controls(source: str) -> list[dict]:
    """Expose literal settings in maintained examples without importing them."""
    labels = {
        "prompt": "Prompt", "num_gpus": "GPU count", "height": "Height", "width": "Width",
        "num_frames": "Frames", "seed": "Seed", "num_inference_steps": "Inference steps",
        "guidance_scale": "Guidance scale", "fps": "FPS", "audio_end_in_s": "Audio duration",
        "text_encoder_cpu_offload": "Offload text encoder", "vae_cpu_offload": "Offload VAE",
        "dit_cpu_offload": "Offload transformer", "inference_torch_compile": "Compile transformer",
        "video_decode_backend": "Video decoder",
    }
    lines = source.splitlines(keepends=True)
    offsets = [0]
    for line in lines:
        offsets.append(offsets[-1] + len(line))

    def position(line, column):
        prefix = source[:offsets[line - 1]] + lines[line - 1].encode()[:column].decode()
        return len(prefix.encode("utf-16-le")) // 2

    settings = {}

    def add(name, value):
        if name not in labels or not isinstance(value, ast.Constant) or type(value.value) not in (str, int, float, bool):
            return
        key = json.dumps([name, value.value])
        entry = settings.setdefault(key, {"name": name, "label": labels[name], "value": value.value,
                                          "integer": type(value.value) is int, "spans": []})
        entry["spans"].append([position(value.lineno, value.col_offset), position(value.end_lineno, value.end_col_offset)])

    for node in ast.walk(ast.parse(source)):
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
            if node.func.attr in {"from_pretrained", "generate_video"}:
                for keyword in node.keywords:
                    add(keyword.arg, keyword.value)
            elif node.func.attr == "add_argument":
                flags = [arg.value for arg in node.args if isinstance(arg, ast.Constant) and isinstance(arg.value, str)]
                flag = next((flag for flag in flags if flag.startswith("--")), "")
                for keyword in node.keywords:
                    if keyword.arg == "default":
                        add(flag.removeprefix("--").replace("-", "_"), keyword.value)
                        key = json.dumps([flag.removeprefix("--").replace("-", "_"), getattr(keyword.value, "value", None)])
                        if key in settings:
                            settings[key]["flag"] = flag
                            if any(item.arg == "type" and isinstance(item.value, ast.Name) and item.value.id == "float"
                                   for item in node.keywords):
                                settings[key]["integer"] = False
                            settings[key]["booleanOptional"] = any(
                                item.arg == "action" and isinstance(item.value, ast.Attribute) and item.value.attr == "BooleanOptionalAction"
                                for item in node.keywords)
        elif isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id.isupper():
                    add(target.id.lower(), node.value)
    result = list(settings.values())
    seen = {}
    for field in result:
        seen[field["label"]] = seen.get(field["label"], 0) + 1
        if seen[field["label"]] > 1:
            field["label"] += f' ({seen[field["label"]]})'
    return result


def signature(obj) -> str:
    """Preserve parameter kinds, annotations, defaults, and return annotations."""
    if obj.is_attribute:
        annotation = f": {obj.annotation}" if obj.annotation else ""
        value = f" = {obj.value}" if obj.value is not None else ""
        return f"{obj.name}{annotation}{value}"
    parameters = list(getattr(obj, "parameters", []))
    parts = []
    keyword_separator = False
    for i, param in enumerate(parameters):
        if param.name in {"self", "cls"}:
            continue
        kind = param.kind.value
        if kind == "keyword-only" and not keyword_separator:
            parts.append("*")
            keyword_separator = True
        prefix = "**" if kind == "variadic keyword" else "*" if kind == "variadic positional" else ""
        if prefix == "*":
            keyword_separator = True
        text = prefix + param.name
        if param.annotation is not None:
            text += f": {param.annotation}"
        if param.default is not None and not prefix:
            text += f" = {param.default}"
        parts.append(text)
        if kind == "positional-only" and (i + 1 == len(parameters) or parameters[i + 1].kind.value != kind):
            parts.append("/")
    result = f"{obj.name}({', '.join(parts)})"
    if obj.is_function and obj.returns is not None:
        result += f" -> {obj.returns}"
    return result


def code_html(text: str) -> str:
    return html.escape(text).replace("[", "[").replace("]", "]")


def prose(text: str) -> str:
    # Docstrings are trusted repository source, rendered as Markdown rather than executed.
    return markdown.markdown(text, extensions=["fenced_code", "tables", "sane_lists"])


def api_overview(text: str, package) -> str:
    def reference(match):
        label, target = match.groups()
        path = (target or label).strip("`")
        obj = package if path == package.path else package[path.removeprefix(package.path + ".")]
        module = obj
        while not module.is_module:
            module = module.parent
        route = "/api/" + module.path.replace(".", "/")
        return f"[{label}]({route}#{path})"

    text = re.sub(r"\[([^\]\n]+)\]\[([^\]\n]*)\]", reference, text)

    def guide(match):
        label, target = match.groups()
        url = urlsplit(target)
        if url.scheme or url.netloc or not re.search(r"\.mdx?$", url.path):
            return match[0]
        url = urlsplit(urljoin("/api/index.md", target))
        route = re.sub(r"\.mdx?$", "", url.path)
        route = re.sub(r"/index$", "", route) or "/"
        return f"[{label}]({urlunsplit(('', '', route, url.query, url.fragment))})"

    return prose(re.sub(r"\[([^\]\n]+)\]\(([^\s)]+)\)", guide, text))


def docstring(obj, level: int) -> str:
    if not obj.docstring:
        return ""
    output = []
    for section in obj.docstring.parsed:
        kind = section.kind.value
        if kind == "text":
            output.append(prose(section.value))
            continue
        title = section.title or kind.replace("_", " ").capitalize()
        output.append(f"<h{level}>{html.escape(title)}</h{level}>")
        if kind == "examples":
            for example_kind, value in section.value:
                output.append(f"<pre><code>{html.escape(value)}</code></pre>" if example_kind.value == "examples" else prose(value))
        elif isinstance(section.value, list):
            output.append("<dl>")
            for item in section.value:
                name = getattr(item, "name", "") or ""
                annotation = getattr(item, "annotation", "") or ""
                label = f"{name}: {annotation}" if name and annotation else str(name or annotation)
                output.append(f"<dt><code>{html.escape(label)}</code></dt><dd>{prose(getattr(item, 'description', ''))}</dd>")
            output.append("</dl>")
        else:
            output.append(prose(getattr(section.value, "description", str(section.value))))
    return "\n".join(output)


def symbol(obj, path: str, level: int = 2, members: list[str] | None = None, source: bool = True) -> str:
    level = min(level, 6)
    output = [f'<h{level} id="{html.escape(path)}"><code>{html.escape(path)}</code></h{level}>']
    output.append(f'<pre class="doc-signature"><code>{code_html(signature(obj))}</code></pre>')
    output.append(docstring(obj, min(level + 1, 6)))
    if obj.is_class:
        constructor = obj.members.get("__init__")
        if constructor is not None:
            output.append(docstring(constructor, min(level + 1, 6)))
    if source and not obj.is_alias:
        try:
            code = obj.source
        except (ValueError, KeyError):
            code = ""
        if code:
            output.append(f'<details><summary>Source</summary><pre><code>{code_html(code)}</code></pre></details>')
    if obj.is_class:
        for name, member in obj.members.items():
            if members is not None and name not in members:
                continue
            if name.startswith("_") or member.is_module:
                continue
            if member.is_attribute and not member.docstring:
                continue
            try:
                output.append(symbol(member, f"{path}.{name}", min(level + 1, 6), source=source))
            except griffe.AliasResolutionError:
                continue
    return "\n".join(output)


def write_page(route: str, title: str, body: str, anchors: dict) -> None:
    # Preserve incoming fragments, including historical category labels whose layout changed.
    present = set(re.findall(r'id="([^"]+)"', body))
    aliases = "\n".join(f'<a id="{html.escape(h["id"])}"></a>' for h in anchors.get(route, []) if h["id"] not in present)
    target = GENERATED / "python-reference" / route / "index.md"
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text("---\n" + yaml.safe_dump({"title": title}, sort_keys=False) +
                      f'---\n\n<div class="python-reference">\n{aliases}\n{body}\n</div>\n')


def prepare() -> None:
    GENERATED.mkdir(exist_ok=True)
    examples = runpy.run_path(str(ROOT / "docs/generate_examples.py"))
    examples["validate_cookbook"]()
    examples["generate_examples"](generate_main_index=True)
    source_dir = GENERATED / "source-examples"
    shutil.rmtree(source_dir, ignore_errors=True)
    source_dir.mkdir()
    for recipe in json.loads((ROOT / "docs/assets/cookbook-recipes.json").read_text())["recipes"]:
        path = ROOT / recipe["source"]
        if path.suffix == ".py":
            source = path.read_text()
            controls = source_controls(source)
            argv = shlex.split(recipe["command"])
            for field in controls:
                flag = field.get("flag")
                positions = [i for i, argument in enumerate(argv[:-1]) if argument == flag]
                if positions and type(field["value"]) is not bool:
                    converter = float if isinstance(field["value"], (int, float)) and not field["integer"] else type(field["value"])
                    field["value"] = converter(argv[positions[-1] + 1])
                elif field.get("booleanOptional"):
                    if flag in argv:
                        field["value"] = True
                    elif flag.replace("--", "--no-", 1) in argv:
                        field["value"] = False
            (source_dir / f'{recipe["id"]}.json').write_text(json.dumps({"source": source, "controls": controls}))
    anchors = json.loads((SITE / "migration-anchors.json").read_text())
    shutil.rmtree(GENERATED / "python-reference", ignore_errors=True)
    package = griffe.load("fastvideo", search_paths=[ROOT], allow_inspection=False, resolve_aliases=True,
                          resolve_external=False, docstring_parser="google", docstring_options={"warnings": False})
    modules = []

    def visit(module):
        if module.is_alias or any(part.startswith("_") or part in {"tests", "third_party"} for part in module.path.split(".")):
            return
        modules.append(module)
        for child in module.modules.values():
            visit(child)

    visit(package)
    for module in sorted(modules, key=lambda item: item.path):
        body = [f'<a id="{module.path}"></a>', docstring(module, 2)]
        for name, member in module.members.items():
            if name.startswith("_"):
                continue
            try:
                # External imports belong to their upstream reference, not FastVideo's API.
                if member.is_alias and not member.target_path.startswith("fastvideo."):
                    continue
                if member.is_module or member.is_attribute and not member.docstring:
                    continue
                body.append(symbol(member, f"{module.path}.{name}"))
            except griffe.AliasResolutionError:
                continue
        write_page("api/" + module.path.replace(".", "/"), module.name, "\n".join(body), anchors)
    write_page("api", "Python API reference", api_overview((ROOT / "docs/api/index.md").read_text(), package), anchors)
    v1 = (ROOT / "docs/getting_started/v1_api.md").read_text()

    def reference(match):
        path, level, members = match.groups()
        return "\n" + symbol(package[path.removeprefix("fastvideo.")], path,
                               int(level), members.split(","), source=False) + "\n"

    v1 = re.sub(r'<api-reference path="([^"]+)" heading="(\d+)" members="([^"]+)" />', reference, v1)
    write_page("getting_started/v1_api", "V1 API", prose(v1.removeprefix("# V1 API\n")), anchors)
    (GENERATED / "old-headings.json").write_text(json.dumps(anchors))
    print(f"Rendered {len(modules) + 1} Python reference pages directly from source, without MkDocs or runtime imports.")


if __name__ == "__main__":
    prepare()