from __future__ import annotations

import argparse
import inspect
from collections import defaultdict
from pathlib import Path
from typing import Any

import orjson

from tracecat.registry.repository import Repository

ROOT_DIR = Path(__file__).resolve().parent.parent
DOCS_DIR = ROOT_DIR / "docs"
TOOLS_DOCS_DIR = DOCS_DIR / "tools"
DOCS_JSON_PATH = DOCS_DIR / "docs.json"
AUTO_GENERATED_COMMENT = (
    "{/* Auto-generated by scripts/generate_tool_docs.py; do not edit by hand. */}"
)
ENTERPRISE_BADGE = (
    '<Badge icon="lock" color="blue" size="lg" shape="pill">Enterprise Edition</Badge>'
)


def _escape_text(text: str) -> str:
    return (
        text.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace("{", "&#123;")
        .replace("}", "&#125;")
    )


def _clean_docstring(text: str | None) -> str | None:
    if not text:
        return None

    lines: list[str] = []
    seen_content = False
    for raw in text.splitlines():
        line = raw.rstrip()
        stripped = line.strip()
        if stripped in {"Args:", "Returns:", "Raises:"}:
            break
        if seen_content and not stripped:
            break
        if stripped:
            seen_content = True
        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 None

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

    return "\n".join(normalized) if normalized else None


def _format_value(value: Any) -> str:
    return orjson.dumps(value, option=orjson.OPT_INDENT_2).decode("utf-8")


def _format_type(schema: dict[str, Any] | None) -> str:
    if not schema:
        return "any"

    if enum_values := schema.get("enum"):
        base_type = schema.get("type")
        if base_type:
            return str(base_type)
        if enum_values:
            return "enum"

    variants = schema.get("anyOf") or schema.get("oneOf")
    if isinstance(variants, list) and variants:
        parts: list[str] = []
        for variant in variants:
            part = _format_type(variant)
            if part not in parts:
                parts.append(part)
        return " | ".join(parts)

    if ref := schema.get("$ref"):
        return str(ref).rsplit("/", maxsplit=1)[-1]

    schema_type = schema.get("type")
    if isinstance(schema_type, list):
        return " | ".join(str(item) for item in schema_type)
    if schema_type == "array":
        return f"array[{_format_type(schema.get('items'))}]"
    if schema_type == "object":
        additional = schema.get("additionalProperties")
        if isinstance(additional, dict):
            return f"map[string, {_format_type(additional)}]"
        return "object"
    if isinstance(schema_type, str):
        return schema_type

    return "any"


def _find_object_schema(schema: dict[str, Any] | None) -> dict[str, Any] | None:
    if not schema:
        return None
    if schema.get("type") == "object" or schema.get("properties"):
        return schema
    for key in ("anyOf", "oneOf"):
        variants = schema.get(key)
        if not isinstance(variants, list):
            continue
        for variant in variants:
            if obj_schema := _find_object_schema(variant):
                return obj_schema
    return None


def _render_nested_fields(schema: dict[str, Any]) -> list[str]:
    obj_schema = _find_object_schema(schema)
    if not obj_schema:
        return []

    properties = obj_schema.get("properties", {})
    if not isinstance(properties, dict) or not properties:
        return []

    required = set(obj_schema.get("required", []))
    lines = ["", "Nested fields:"]
    for name in sorted(properties):
        child = properties[name]
        if not isinstance(child, dict):
            continue
        child_type = _format_type(child)
        label = f"- `{name}` (`{child_type}`"
        if name in required:
            label += ", required"
        label += ")"
        if description := child.get("description"):
            label += f": {_escape_text(str(description).strip())}"
        lines.append(label)
    return lines


def _order_field_names(
    properties: dict[str, Any], required_fields: set[str]
) -> list[str]:
    required = sorted(name for name in properties if name in required_fields)
    optional = sorted(name for name in properties if name not in required_fields)
    return [*required, *optional]


def _render_field(name: str, schema: dict[str, Any], required: bool) -> list[str]:
    field_type = _escape_text(_format_type(schema))
    required_attr = " required" if required else ""
    lines = [f'<ParamField path="{name}" type="{field_type}"{required_attr}>', ""]

    if description := schema.get("description"):
        lines.append(_escape_text(str(description).strip()))

    if "default" in schema:
        lines.extend(["", f"Default: `{_format_value(schema['default'])}`."])

    if enum_values := schema.get("enum"):
        allowed_values = ", ".join(f"`{value}`" for value in enum_values)
        lines.extend(["", f"Allowed values: {allowed_values}."])

    lines.extend(_render_nested_fields(schema))
    lines.extend(["", "</ParamField>", ""])
    return lines


def _render_input_fields(expects_schema: dict[str, Any]) -> list[str]:
    properties = expects_schema.get("properties", {})
    if not isinstance(properties, dict) or not properties:
        return ["This action does not take input fields.", ""]

    required_fields = set(expects_schema.get("required", []))
    lines: list[str] = []
    for name in _order_field_names(properties, required_fields):
        schema = properties[name]
        if not isinstance(schema, dict):
            continue
        lines.extend(
            _render_field(name=name, schema=schema, required=name in required_fields)
        )
    return lines


def _render_secret(secret: Any) -> str:
    secret_type = getattr(secret, "type", None)
    name = f"`{_escape_text(str(secret.name))}`"

    if secret_type == "custom":
        required_keys = getattr(secret, "keys", None) or []
        optional_keys = getattr(secret, "optional_keys", None) or []
        details: list[str] = []
        if required_keys:
            details.append(
                "required values "
                + ", ".join(f"`{_escape_text(str(key))}`" for key in required_keys)
            )
        if optional_keys:
            details.append(
                "optional values "
                + ", ".join(f"`{_escape_text(str(key))}`" for key in optional_keys)
            )
        return f"- {name}: {'; '.join(details)}."

    if secret_type == "oauth":
        return f"- {name}: OAuth token `{_escape_text(str(secret.token_name))}`."

    raise ValueError(f"Unsupported secret type: {secret_type!r}")


def _render_secrets(action: Any) -> list[str]:
    secrets = getattr(action, "secrets", None) or []
    if not secrets:
        return []

    lines = ["### Secrets", ""]
    required_secrets = [
        secret for secret in secrets if not getattr(secret, "optional", False)
    ]
    optional_secrets = [
        secret for secret in secrets if getattr(secret, "optional", False)
    ]

    if required_secrets:
        lines.extend(["Required secrets:", ""])
        lines.extend(_render_secret(secret) for secret in required_secrets)
        lines.append("")

    if optional_secrets:
        lines.extend(["Optional secrets:", ""])
        lines.extend(_render_secret(secret) for secret in optional_secrets)
        lines.append("")

    return lines


def _get_description(action: Any) -> str:
    docstring = _clean_docstring(inspect.getdoc(action.fn))
    return docstring or action.description or "No description provided."


def _namespace_slug(namespace: str) -> str:
    return namespace.removeprefix("tools.")


def _namespace_title(namespace: str, actions: list[Any]) -> str:
    for action in actions:
        if action.display_group:
            return action.display_group
    return _namespace_slug(namespace).replace("_", " ").title()


def _sort_actions(actions: list[Any]) -> list[Any]:
    return sorted(
        actions,
        key=lambda action: (
            (action.default_title or "").casefold(),
            action.name.casefold(),
            action.action.casefold(),
        ),
    )


def _render_action(action: Any) -> list[str]:
    lines = [
        f"## {_escape_text(action.default_title or action.name.replace('_', ' ').title())}",
        "",
        f"Action ID: `{action.action}`",
        "",
        _escape_text(_get_description(action)),
        "",
    ]

    if action.required_entitlements:
        lines.extend([ENTERPRISE_BADGE, ""])

    if action.doc_url:
        lines.extend(
            [
                f"Reference: [{_escape_text(action.doc_url)}]({_escape_text(action.doc_url)})",
                "",
            ]
        )

    if action.deprecated:
        lines.extend([f"**Deprecated:** {_escape_text(action.deprecated)}", ""])

    lines.extend(_render_secrets(action))
    lines.extend(["### Input fields", ""])
    lines.extend(_render_input_fields(action.get_interface()["expects"]))
    return lines


def _render_page(namespace: str, actions: list[Any]) -> str:
    title = _namespace_title(namespace, actions)
    description = (
        f"Reference for the Tracecat {title} integration: registered actions, "
        f"required secrets, expected inputs, and example workflow usage."
    )
    lines = [
        "---",
        f'title: "{title}"',
        f'description: "{description}"',
        "---",
        "",
        AUTO_GENERATED_COMMENT,
    ]

    for action in _sort_actions(actions):
        lines.extend(_render_action(action))

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


def _write_if_changed(path: Path, content: str, check: bool) -> bool:
    existing = path.read_text(encoding="utf-8") if path.exists() else None
    changed = existing != content
    if changed and not check:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(content, encoding="utf-8")
    return changed


def _delete_if_exists(path: Path, check: bool) -> bool:
    if not path.exists():
        return False
    if not check:
        path.unlink()
    return True


def _load_docs_json() -> dict[str, Any]:
    return orjson.loads(DOCS_JSON_PATH.read_bytes())


def _update_tools_nav(doc: dict[str, Any], tool_pages: list[str]) -> None:
    navigation = doc.get("navigation")
    if not isinstance(navigation, dict):
        raise ValueError("docs.json is missing `navigation`.")

    tabs = navigation.get("tabs")
    if not isinstance(tabs, list):
        raise ValueError("docs.json is missing `navigation.tabs`.")

    for tab in tabs:
        if not isinstance(tab, dict) or tab.get("tab") != "Tools":
            continue
        groups = tab.get("groups")
        if not isinstance(groups, list):
            raise ValueError("The Tools tab is missing `groups`.")
        for group in groups:
            if not isinstance(group, dict) or group.get("group") != "Tools":
                continue
            group["pages"] = tool_pages
            return

    raise ValueError("Could not find navigation entry for the Tools tab.")


def _dump_docs_json(doc: dict[str, Any]) -> str:
    return orjson.dumps(doc, option=orjson.OPT_INDENT_2).decode("utf-8") + "\n"


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--check",
        action="store_true",
        help="Fail if generated output differs from the checked-in files.",
    )
    args = parser.parse_args()

    repo = Repository()
    repo.init(include_base=True, include_templates=True)

    actions_by_namespace: dict[str, list[Any]] = defaultdict(list)
    for action in repo.store.values():
        if action.namespace.startswith("tools.") and action.include_in_schema:
            actions_by_namespace[action.namespace].append(action)

    if not actions_by_namespace:
        raise ValueError("No tool namespaces were found in the registry.")

    changed_paths: list[str] = []
    expected_paths: set[Path] = set()
    tool_pages: list[str] = []

    for namespace in sorted(actions_by_namespace):
        output_path = TOOLS_DOCS_DIR / f"{_namespace_slug(namespace)}.mdx"
        expected_paths.add(output_path)
        tool_pages.append(f"tools/{_namespace_slug(namespace)}")
        content = _render_page(namespace, actions_by_namespace[namespace])
        if _write_if_changed(output_path, content, check=args.check):
            changed_paths.append(str(output_path.relative_to(ROOT_DIR)))

    existing_tool_pages = {
        path for path in TOOLS_DOCS_DIR.glob("*.mdx") if path.name not in {"_meta.mdx"}
    }
    stale_paths = sorted(existing_tool_pages - expected_paths)
    for stale_path in stale_paths:
        if _delete_if_exists(stale_path, check=args.check):
            changed_paths.append(str(stale_path.relative_to(ROOT_DIR)))

    docs_json = _load_docs_json()
    _update_tools_nav(docs_json, tool_pages)
    rendered_docs_json = _dump_docs_json(docs_json)
    if _write_if_changed(DOCS_JSON_PATH, rendered_docs_json, check=args.check):
        changed_paths.append(str(DOCS_JSON_PATH.relative_to(ROOT_DIR)))

    if args.check and changed_paths:
        print("Generated tool docs are out of date:")
        for path in changed_paths:
            print(f"- {path}")
        raise SystemExit(1)


if __name__ == "__main__":
    main()
