#!/usr/bin/env python3
"""Generate llms.txt and llms-full.txt from the SDK documentation.

Reads all English Markdown docs in the order defined by mkdocs.yml nav,
cleans MkDocs-specific markup, and produces two files:

- llms.txt: lightweight index with links and one-line descriptions
- llms-full.txt: full documentation concatenated in plain Markdown

Follows the llms.txt spec from https://llmstxt.org/
"""

from __future__ import annotations

import re
from pathlib import Path

import yaml

BASE_URL = "https://pe-menezes.github.io/arandu"
SDK_ROOT = Path(__file__).resolve().parent.parent
DOCS_EN = SDK_ROOT / "docs" / "en"
MKDOCS_YML = SDK_ROOT / "mkdocs.yml"

# Page descriptions for llms.txt index (one-line per page)
PAGE_DESCRIPTIONS: dict[str, str] = {
    "index.md": "Project overview, features, installation, and quickstart.",
    "getting-started.md": "Step-by-step guide: install, setup PostgreSQL+pgvector, first write and retrieve.",
    "concepts/write-pipeline.md": "How the write pipeline works: extraction, entity resolution, reconciliation, upsert.",
    "concepts/read-pipeline.md": "How the read pipeline works: planning, retrieval, reranking, formatting.",
    "concepts/background-jobs.md": "Background jobs: clustering, consolidation, memify, sleep-time compute.",
    "concepts/design-philosophy.md": "Design decisions and neuroscience parallels behind the architecture.",
    "advanced/write-api.md": "Advanced write pipeline API: extraction strategy, canonicalization, entity helpers, corrections.",
    "advanced/read-api.md": "Advanced read pipeline API: query expansion, graph traversal, spreading activation, directives.",
    "advanced/database.md": "Database utilities: create_engine, create_session_factory, init_db.",
    "advanced/data-types.md": "All enums, dataclasses, and result types reference.",
    "configuration.md": "All MemoryConfig parameters grouped by pipeline stage with defaults.",
    "custom-providers.md": "How to implement custom LLMProvider and EmbeddingProvider.",
    "cookbook.md": "Practical examples: basic usage, custom providers, config tuning, background jobs, multi-user.",
    "reference/index.md": "Auto-generated API reference for all public classes and functions.",
}


class _IgnoreUnknownTagsLoader(yaml.SafeLoader):
    """YAML loader that ignores unknown tags (e.g. !!python/name:)."""


_IgnoreUnknownTagsLoader.add_multi_constructor(
    "tag:yaml.org,2002:python/",
    lambda loader, suffix, node: None,
)


def parse_nav_order() -> list[str]:
    """Extract page file paths from mkdocs.yml nav in order."""
    with open(MKDOCS_YML) as f:
        config = yaml.load(f, Loader=_IgnoreUnknownTagsLoader)  # noqa: S506

    pages: list[str] = []

    def _walk(node: object) -> None:
        if isinstance(node, str):
            pages.append(node)
        elif isinstance(node, dict):
            for v in node.values():
                _walk(v)
        elif isinstance(node, list):
            for item in node:
                _walk(item)

    _walk(config.get("nav", []))
    return pages


def extract_title(content: str) -> str:
    """Extract the first H1 title from Markdown content."""
    match = re.search(r"^#\s+(.+)$", content, re.MULTILINE)
    return match.group(1).strip() if match else "Untitled"


def clean_mkdocs_markup(content: str) -> str:
    """Remove MkDocs-specific markup, producing plain Markdown."""
    # Remove mkdocstrings directives (::: arandu.X)
    content = re.sub(r"^:::.*$", "", content, flags=re.MULTILINE)

    # Convert admonitions to blockquotes
    # !!! note "Title"     or    !!! note
    #     Content                     Content
    def _convert_admonition(m: re.Match[str]) -> str:
        admon_type = m.group(1).strip().capitalize()
        title = m.group(2).strip() if m.group(2) else admon_type
        body = m.group(3)
        # Dedent body (remove 4-space indent)
        lines = [line[4:] if line.startswith("    ") else line for line in body.splitlines()]
        body_text = "\n".join(lines).strip()
        return f"> **{title}:** {body_text}\n"

    content = re.sub(
        r'^(?:!!!|\?\?\?)\s+(\w+)(?:\s+"([^"]*)")?\n((?:    .+\n?)*)',
        _convert_admonition,
        content,
        flags=re.MULTILINE,
    )

    # Remove tab syntax (=== "Tab Name") — keep content
    content = re.sub(r'^===\s+"[^"]*"\s*$', "", content, flags=re.MULTILINE)

    # Remove {.class} and {: .class} attribute syntax
    content = re.sub(r"\{[:.][^}]*\}", "", content)

    # Clean up excessive blank lines (3+ → 2)
    content = re.sub(r"\n{3,}", "\n\n", content)

    return content.strip()


def page_url(page_path: str) -> str:
    """Convert a doc file path to its full URL."""
    # reference/index.md → reference/
    # getting-started.md → getting-started/
    # index.md → (root)
    url_path = page_path.replace(".md", "").replace("index", "").strip("/")
    if url_path:
        return f"{BASE_URL}/{url_path}/"
    return f"{BASE_URL}/"


def generate_llms_txt(pages: list[str]) -> str:
    """Generate llms.txt index following the llmstxt.org spec."""
    sections: dict[str, list[str]] = {
        "Guides": [],
        "Advanced": [],
        "Reference": [],
    }

    for page_path in pages:
        desc = PAGE_DESCRIPTIONS.get(page_path, "")
        url = page_url(page_path)
        title = page_path.replace(".md", "").replace("/", " — ").replace("-", " ").title()

        # Read actual title from file
        doc_file = DOCS_EN / page_path
        if doc_file.exists():
            title = extract_title(doc_file.read_text(encoding="utf-8"))

        entry = f"- [{title}]({url}): {desc}"

        if page_path.startswith("advanced/"):
            sections["Advanced"].append(entry)
        elif page_path.startswith("reference/") or page_path in ("configuration.md", "custom-providers.md"):
            sections["Reference"].append(entry)
        else:
            sections["Guides"].append(entry)

    lines = [
        "# arandu",
        "",
        "> Long-term memory for AI agents. Extract, store, and retrieve structured facts from conversations "
        "using PostgreSQL + pgvector. Pluggable LLM and embedding providers. "
        "Inspired by how the human brain encodes, consolidates, and retrieves memories.",
        "",
    ]

    for section_name, entries in sections.items():
        if entries:
            lines.append(f"## {section_name}")
            lines.append("")
            lines.extend(entries)
            lines.append("")

    return "\n".join(lines).strip() + "\n"


def generate_llms_full_txt(pages: list[str]) -> str:
    """Generate llms-full.txt with all docs concatenated."""
    parts = [
        "# arandu — Complete Documentation",
        "",
        "> This file contains the complete documentation for the arandu SDK, "
        "a long-term memory system for AI agents. Generated automatically from the source docs.",
        "",
        "---",
        "",
    ]

    for page_path in pages:
        doc_file = DOCS_EN / page_path
        if not doc_file.exists():
            continue

        content = doc_file.read_text(encoding="utf-8")
        cleaned = clean_mkdocs_markup(content)

        if cleaned:
            parts.append(cleaned)
            parts.append("")
            parts.append("---")
            parts.append("")

    return "\n".join(parts).strip() + "\n"


def main() -> None:
    """Generate llms.txt and llms-full.txt."""
    pages = parse_nav_order()
    print(f"Found {len(pages)} pages in nav")

    # Generate llms.txt
    llms_txt = generate_llms_txt(pages)
    output_llms = DOCS_EN / "llms.txt"
    output_llms.write_text(llms_txt, encoding="utf-8")
    print(f"Generated {output_llms} ({len(llms_txt)} bytes)")

    # Generate llms-full.txt
    llms_full = generate_llms_full_txt(pages)
    output_full = DOCS_EN / "llms-full.txt"
    output_full.write_text(llms_full, encoding="utf-8")
    print(f"Generated {output_full} ({len(llms_full)} bytes)")


if __name__ == "__main__":
    main()
