"""
Nexa Language — Documentation Generator (v23)
`nexa doc [directory] [--output docs/]`

Extracts docstrings, function signatures, class hierarchies,
and interface definitions from Nexa source files.
Generates:
  - docs/index.md     — project overview
  - docs/<file>.md    — per-file reference
  - docs/api.md       — full API reference

Docstring format:
  def fetch(url) {
      \"\"\"
      Fetch data from a URL.
      @param url  string  URL to fetch
      @returns    dict    Response data
      @example    let r = fetch("https://example.com")
      \"\"\"
  }
"""
from __future__ import annotations
import os, re
from typing import Any, List, Optional
from dataclasses import dataclass, field


@dataclass
class DocEntry:
    kind: str          # "function" | "class" | "interface" | "const" | "enum"
    name: str
    signature: str
    docstring: str
    line: int = 0
    members: List["DocEntry"] = field(default_factory=list)
    source_file: str = ""


class DocGenerator:
    """Parse Nexa AST and generate Markdown documentation."""

    def __init__(self, output_dir: str = "docs"):
        self.output_dir = output_dir
        self._all_entries: List[DocEntry] = []

    def generate(self, target: str = ".") -> str:
        """Generate docs for all .nexa files. Returns output dir path."""
        os.makedirs(self.output_dir, exist_ok=True)

        files = self._discover(target)
        if not files:
            print(f"  ⚠️  No .nexa files found in: {target}")
            return self.output_dir

        print(f"\n  📚 Nexa Doc Generator v23")
        print(f"  {'─' * 50}")

        for fp, src in files:
            entries = self._parse_file(fp, src)
            self._all_entries.extend(entries)
            self._write_file_doc(fp, entries, src)

        self._write_index(files)
        self._write_api_ref()

        total = len(self._all_entries)
        print(f"  ✅ Generated docs for {len(files)} file(s), {total} symbol(s) → {self.output_dir}/")
        return self.output_dir

    # ── File discovery ─────────────────────────────────────────────────────────
    def _discover(self, path: str):
        if os.path.isfile(path) and path.endswith(".nexa"):
            return [(path, open(path, encoding="utf-8").read())]
        result = []
        for root, _, names in os.walk(path):
            for name in sorted(names):
                if name.endswith(".nexa") and not name.startswith("test_"):
                    fp = os.path.join(root, name)
                    result.append((fp, open(fp, encoding="utf-8").read()))
        return result

    # ── AST → DocEntry ─────────────────────────────────────────────────────────
    def _parse_file(self, path: str, src: str) -> List[DocEntry]:
        """Parse a .nexa source file and extract doc entries."""
        from nexa_lang.lexer import Lexer
        from nexa_lang.parser import Parser

        entries = []
        try:
            tokens = Lexer(src).tokenize()
            ast = Parser(tokens).parse()
        except Exception as e:
            return [DocEntry("error", path, "", f"Parse error: {e}")]

        for stmt in ast.statements:
            entry = self._entry_from_stmt(stmt, path)
            if entry: entries.append(entry)
        return entries

    def _entry_from_stmt(self, stmt, source_file: str) -> Optional[DocEntry]:
        t = type(stmt).__name__
        line = getattr(stmt, 'line', 0)

        if t == "FunctionDef":
            tps = f"<{', '.join(getattr(stmt, 'type_params', []) or [])}>" if getattr(stmt, 'type_params', None) else ""
            params = ", ".join(stmt.params)
            ret = f" -> {stmt.return_type}" if getattr(stmt, 'return_type', None) else ""
            sig = f"def {stmt.name}{tps}({params}){ret}"
            doc = self._extract_docstring(stmt.body)
            return DocEntry("function", stmt.name, sig, doc, line, source_file=source_file)

        elif t == "ClassDef":
            base = f"({stmt.base})" if stmt.base else ""
            impls = f" implements {', '.join(stmt.implements)}" if getattr(stmt, 'implements', None) else ""
            abstract = "abstract " if getattr(stmt, 'is_abstract', False) else ""
            sig = f"{abstract}class {stmt.name}{base}{impls}"
            doc = self._extract_docstring(stmt.body)
            members = []
            for m in stmt.body:
                me = self._entry_from_stmt(m, source_file)
                if me: members.append(me)
            return DocEntry("class", stmt.name, sig, doc, line, members, source_file)

        elif t == "InterfaceDef":
            methods = ", ".join(stmt.method_names)
            sig = f"interface {stmt.name} {{ {methods} }}"
            return DocEntry("interface", stmt.name, sig, "", line, source_file=source_file)

        elif t == "ConstStatement":
            sig = f"const {stmt.name} = {_safe_repr(stmt.value)}"
            return DocEntry("const", stmt.name, sig, "", line, source_file=source_file)

        elif t == "EnumDef":
            variants = " | ".join(
                v.name + (f"({', '.join(v.fields)})" if v.fields else "")
                for v in stmt.variants
            )
            sig = f"enum {stmt.name} {{ {variants} }}"
            return DocEntry("enum", stmt.name, sig, "", line, source_file=source_file)

        return None

    def _extract_docstring(self, body: list) -> str:
        """Extract the first string literal from a function/class body as docstring."""
        if not body: return ""
        first = body[0]
        t = type(first).__name__
        if t == "ExprStatement":
            inner = type(first.expr).__name__
            if inner in ("StringLiteral", "DocString"):
                return getattr(first.expr, 'value', '')
        elif t in ("StringLiteral", "DocString"):
            return getattr(first, 'value', '')
        return ""

    # ── Markdown Writers ───────────────────────────────────────────────────────
    def _write_file_doc(self, fp: str, entries: List[DocEntry], src: str):
        """Write per-file .md documentation."""
        base = os.path.splitext(os.path.basename(fp))[0]
        out_path = os.path.join(self.output_dir, base + ".md")

        lines = [f"# `{base}.nexa`\n", f"> Source: `{fp}`\n", "---\n"]

        for entry in entries:
            lines.append(self._render_entry(entry, level=2))

        with open(out_path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))
        print(f"  📄 {out_path}  ({len(entries)} symbol(s))")

    def _render_entry(self, entry: DocEntry, level: int = 2) -> str:
        h = "#" * level
        kind_icon = {"function": "🔧", "class": "📦", "interface": "🔷",
                     "const": "🔒", "enum": "🎯", "error": "❌"}.get(entry.kind, "▪")
        out = [f"\n{h} {kind_icon} `{entry.name}`\n"]
        out.append(f"```nexa\n{entry.signature}\n```\n")
        if entry.docstring:
            out.append(self._format_docstring(entry.docstring))
        for m in entry.members:
            out.append(self._render_entry(m, level + 1))
        return "\n".join(out)

    def _format_docstring(self, doc: str) -> str:
        """Parse @param/@returns/@example tags from docstring."""
        lines = doc.strip().splitlines()
        desc_lines = []
        tag_lines = []
        for line in lines:
            stripped = line.strip()
            if stripped.startswith("@"):
                tag_lines.append(stripped)
            else:
                desc_lines.append(stripped)
        out = []
        if desc_lines:
            out.append(" ".join(l for l in desc_lines if l) + "\n")
        if tag_lines:
            out.append("| Tag | Name | Type | Description |")
            out.append("|-----|------|------|-------------|")
            for tag in tag_lines:
                parts = tag.split(None, 3)
                tag_name = parts[0].lstrip("@") if parts else ""
                name = parts[1] if len(parts) > 1 else ""
                typ  = parts[2] if len(parts) > 2 else ""
                desc = parts[3] if len(parts) > 3 else ""
                out.append(f"| `@{tag_name}` | `{name}` | `{typ}` | {desc} |")
            out.append("")
        return "\n".join(out)

    def _write_index(self, files: list):
        """Write docs/index.md with overview table."""
        out_path = os.path.join(self.output_dir, "index.md")
        lines = [
            "# Nexa Project Docs\n",
            "> Auto-generated by `nexa doc` — Nexa v23\n",
            "---\n",
            "## Modules\n",
            "| File | Symbols |",
            "|------|---------|",
        ]
        for fp, _ in files:
            base = os.path.splitext(os.path.basename(fp))[0]
            count = sum(1 for e in self._all_entries if e.source_file == fp)
            lines.append(f"| [`{base}.nexa`]({base}.md) | {count} |")
        lines.append("\n---\n")
        lines.append("## All Symbols\n")
        for entry in sorted(self._all_entries, key=lambda e: e.name):
            base = os.path.splitext(os.path.basename(entry.source_file))[0]
            lines.append(f"- [`{entry.name}`]({base}.md) — *{entry.kind}*")
        with open(out_path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines) + "\n")
        print(f"  📋 {out_path}")

    def _write_api_ref(self):
        """Write docs/api.md — full consolidated API reference."""
        out_path = os.path.join(self.output_dir, "api.md")
        lines = ["# Nexa API Reference\n", "> All public symbols\n", "---\n"]

        for kind, icon in [("function", "🔧"), ("class", "📦"),
                             ("interface", "🔷"), ("enum", "🎯"), ("const", "🔒")]:
            group = [e for e in self._all_entries if e.kind == kind]
            if not group: continue
            lines.append(f"\n## {icon} {kind.capitalize()}s\n")
            for entry in sorted(group, key=lambda e: e.name):
                lines.append(f"### `{entry.name}`")
                lines.append(f"```nexa\n{entry.signature}\n```")
                if entry.docstring:
                    lines.append(entry.docstring.strip())
                lines.append("")

        with open(out_path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines) + "\n")
        print(f"  📖 {out_path}")


def _safe_repr(node) -> str:
    t = type(node).__name__
    if t == "NumberLiteral": return str(node.value)
    if t == "StringLiteral": return f'"{node.value}"'
    if t == "BoolLiteral":   return "true" if node.value else "false"
    if t == "NullLiteral":   return "null"
    return "<expr>"


def generate_docs(target: str = ".", output_dir: str = "docs") -> str:
    return DocGenerator(output_dir).generate(target)
