from __future__ import annotations

import ast
from collections import defaultdict
from pathlib import Path

NAMESPACE_ORDER = (
    "workspaces",
    "workflows",
    "cases",
    "tables",
    "variables",
    "secrets",
    "integrations",
    "agents",
)


def _extract_namespace_map(module: ast.Module) -> dict[str, str]:
    for node in module.body:
        target_name: str | None = None
        value: ast.AST | None = None
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name):
                    target_name = target.id
                    value = node.value
                    break
        elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
            target_name = node.target.id
            value = node.value

        if target_name != "_TOOL_NAMESPACE_BY_NAME":
            continue
        if not isinstance(value, ast.Dict):
            raise ValueError("_TOOL_NAMESPACE_BY_NAME must be a dict literal")
        mapping: dict[str, str] = {}
        for key_node, value_node in zip(value.keys, value.values, strict=True):
            if not isinstance(key_node, ast.Constant) or not isinstance(
                key_node.value, str
            ):
                raise ValueError("Tool namespace key must be a string literal")
            if not isinstance(value_node, ast.Constant) or not isinstance(
                value_node.value, str
            ):
                raise ValueError("Tool namespace value must be a string literal")
            mapping[key_node.value] = value_node.value
        return mapping
    raise ValueError("Could not find _TOOL_NAMESPACE_BY_NAME in tracecat/mcp/server.py")


def _clean_docstring(text: str | None) -> str:
    if not text:
        return "No description provided."
    lines: list[str] = []
    for raw in text.splitlines():
        line = raw.rstrip()
        stripped = line.strip()
        if stripped in {"Args:", "Returns:", "Raises:"}:
            break
        lines.append(line)

    while lines and not lines[0].strip():
        lines.pop(0)
    while lines and not lines[-1].strip():
        lines.pop()
    if not lines:
        return "No description provided."

    normalized: list[str] = []
    for line in lines:
        stripped = line.strip()
        if not stripped:
            if normalized and normalized[-1]:
                normalized.append("")
            continue

        if stripped.startswith("- "):
            normalized.append(stripped)
            continue

        if (
            normalized
            and normalized[-1]
            and (line.startswith((" ", "\t")) or not normalized[-1].startswith("- "))
        ):
            normalized[-1] = f"{normalized[-1]} {stripped}"
        else:
            normalized.append(stripped)

    expanded: list[str] = []
    for line in normalized:
        if not line:
            expanded.append("")
            continue
        if " - " not in line:
            expanded.append(line)
            continue
        lead, *tail = line.split(" - ")
        if not tail:
            expanded.append(line)
            continue
        if lead.strip():
            expanded.append(lead.strip())
        expanded.extend(f"- {chunk.strip()}" for chunk in tail if chunk.strip())

    while expanded and not expanded[-1]:
        expanded.pop()
    return "\n".join(expanded) if expanded else "No description provided."


def _format_namespace_title(namespace: str) -> str:
    return namespace.replace("_", " ").title()


def main() -> None:
    root_dir = Path(__file__).resolve().parent.parent
    server_file = root_dir / "tracecat" / "mcp" / "server.py"
    output_file = root_dir / "docs" / "snippets" / "mcp-tools.mdx"

    source = server_file.read_text(encoding="utf-8")
    module = ast.parse(source)
    namespace_by_tool = _extract_namespace_map(module)

    tools_by_namespace: dict[str, list[tuple[str, str]]] = defaultdict(list)
    for node in module.body:
        if not isinstance(node, ast.AsyncFunctionDef):
            continue
        tool_namespace = namespace_by_tool.get(node.name)
        if not tool_namespace:
            continue
        tool_name = f"{tool_namespace}_{node.name}"
        tools_by_namespace[tool_namespace].append(
            (tool_name, _clean_docstring(ast.get_docstring(node)))
        )

    lines = [
        "{/* Auto-generated by scripts/generate_mcp_docs.py; do not edit by hand. */}",
        "",
    ]
    for namespace in NAMESPACE_ORDER:
        tools = tools_by_namespace.get(namespace)
        if not tools:
            continue
        title = _format_namespace_title(namespace)
        lines.append(f'<ResponseField name="{title}" type="Namespace">')
        lines.append('  <Expandable title="Tools">')
        for tool_name, docstring in tools:
            lines.append(f'    <ResponseField name="{tool_name}" type="tool">')
            for doc_line in docstring.splitlines():
                if doc_line:
                    rendered_line = (
                        f"-> {doc_line[2:]}" if doc_line.startswith("- ") else doc_line
                    )
                    lines.append(f"  {rendered_line}")
                else:
                    lines.append("")
            lines.append("    </ResponseField>")
        lines.append("  </Expandable>")
        lines.append("</ResponseField>")
        lines.append("")

    output_file.write_text("\n".join(lines), encoding="utf-8")


if __name__ == "__main__":
    main()
