from __future__ import annotations

import argparse
import inspect
from pathlib import Path
from textwrap import dedent
from typing import Any

import orjson

from tracecat.registry.repository import Repository

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

EXAMPLES: dict[str, dict[str, str]] = {
    "ai_action_extract_json": {
        "title": "Extract structured fields from an alert",
        "language": "yaml",
        "code": dedent(
            """
            - ref: extract_alert
              action: ai.action
              args:
                model_name: gpt-4.1-mini
                model_provider: openai
                instructions: |
                  Extract the alert into the requested schema.
                  Use null for fields that are missing.
                user_prompt: |
                  Parse this alert and normalize the key fields:

                  ${{ TRIGGER.raw_alert }}
                output_type:
                  type: object
                  properties:
                    severity:
                      type: string
                    user:
                      type:
                        - string
                        - "null"
                    src_ip:
                      type:
                        - string
                        - "null"
                    needs_escalation:
                      type: boolean
                  required:
                    - severity
                    - needs_escalation
            """
        ).strip(),
    },
    "ai_action_classify_email": {
        "title": "Classify an email triage decision",
        "language": "yaml",
        "code": dedent(
            """
            - ref: classify_email
              action: ai.action
              args:
                model_name: claude-3-5-sonnet-latest
                model_provider: anthropic
                instructions: |
                  Return one of: phishing, benign, spam, unknown.
                user_prompt: |
                  Subject: ${{ TRIGGER.subject }}
                  Sender: ${{ TRIGGER.sender }}
                  Body:
                  ${{ TRIGGER.body }}
                output_type:
                  type: string
                  enum:
                    - phishing
                    - benign
                    - spam
                    - unknown
            """
        ).strip(),
    },
    "ai_agent_investigate_alert": {
        "title": "Investigate an alert with tools",
        "language": "yaml",
        "code": dedent(
            """
            - ref: investigate_alert
              action: ai.agent
              args:
                model_name: gpt-4.1
                model_provider: openai
                instructions: |
                  You are a security triage agent.
                  Review the alert, gather missing context, and write a concise analyst summary.
                user_prompt: |
                  Investigate this alert and explain whether it needs escalation:

                  ${{ TRIGGER.alert }}
                actions:
                  - tools.slack.lookup_users_by_email
                  - tools.github.search_code
                  - core.cases.create_comment
                max_tool_calls: 6
                output_type:
                  type: object
                  properties:
                    verdict:
                      type: string
                    summary:
                      type: string
                  required:
                    - verdict
                    - summary
            """
        ).strip(),
    },
    "ai_agent_with_approvals": {
        "title": "Require approval for sensitive tools",
        "language": "yaml",
        "code": dedent(
            """
            - ref: draft_case_update
              action: ai.agent
              args:
                model_name: claude-3-5-sonnet-latest
                model_provider: anthropic
                instructions: |
                  Investigate the incident, then propose the next action.
                  Create or update records only if the required approval is granted.
                user_prompt: |
                  Review the latest evidence for case ${{ TRIGGER.case_id }} and decide what to do next.
                actions:
                  - core.cases.get
                  - core.cases.create_comment
                  - core.cases.update
                tool_approvals:
                  core.cases.update: true
                max_tool_calls: 4
            """
        ).strip(),
    },
    "preset_agent_run": {
        "title": "Run a saved agent preset",
        "language": "yaml",
        "code": dedent(
            """
            - ref: run_security_analyst
              action: ai.preset_agent
              args:
                preset: security-analyst
                preset_version: 3
                user_prompt: |
                  Review the incident below and decide whether to escalate it:

                  ${{ TRIGGER.incident }}
                max_tool_calls: 8
            """
        ).strip(),
    },
}

PAGES: list[dict[str, Any]] = [
    {
        "slug": "ai-action",
        "title": "AI action",
        "description": "Use `ai.action` to run a single LLM call in a workflow.",
        "body": dedent(
            """
            Use `ai.action` when you only need a prompt and an output. It runs one model call and does not let the model call tools.

            ## Capabilities

            - Pass task instructions in `instructions` and task data in `user_prompt`.
            - Use `output_type` when you need a string, list, or JSON-shaped response for downstream actions.
            - Tune provider-specific behavior with `model_settings`.

            ## Structured outputs

            Use `output_type` when a downstream action needs a predictable shape.

            - Use a scalar such as `str` or `int` for simple classification or routing.
            - Use a JSON schema object when you need named fields.
            - Tracecat parses valid JSON before storing it in the action result.

            ## Reference
            """
        ).strip(),
        "actions": [
            {
                "id": "ai.action",
                "examples": [
                    "ai_action_extract_json",
                    "ai_action_classify_email",
                ],
            }
        ],
    },
    {
        "slug": "ai-agent",
        "title": "AI agent",
        "description": "Use `ai.agent` and `ai.preset_agent` for tool-calling agent runs in workflows.",
        "body": dedent(
            """
            Use `ai.agent` when the model needs tool calls. You give the agent a prompt, instructions, and an allowlist of actions it can call during the run.

            ## Capabilities

            - `ai.agent`: Prompt plus tool calls. Use the `actions` list to control which Tracecat actions the agent can use.
            - <Badge icon="lock" color="blue" size="sm" shape="pill">EE</Badge> `ai.preset_agent`: Prompt plus a saved agent configuration. Use this when you want reusable instructions, tools, and MCP integrations across workflows.
            - `tool_approvals`: Require approval before selected tools run. This is an enterprise feature.
            - `max_tool_calls` and `max_requests`: Bound how much work the agent can do in a single run.

            ## Structured outputs

            `ai.agent` also supports `output_type`.

            - Use it when the agent should return a final object or typed value after tool use.
            - Keep the schema focused on the final answer, not the intermediate tool steps.

            ## MCP

            <Badge icon="lock" color="blue" size="lg" shape="pill">Enterprise Edition</Badge>

            `ai.agent` does not take MCP servers directly in the workflow action. If you need MCP, save that configuration in an agent preset and run it with `ai.preset_agent`.
            `ai.preset_agent` supports both remote and stdio MCP servers.

            Internet access is controlled by the root preset for the shared sandbox process. Subagent presets can define their own tools and MCP integrations, but their internet setting does not grant network access unless the root preset also enables it.

            See [MCP integrations](/automations/integrations/mcp-integrations) to learn more.

            ## Reference
            """
        ).strip(),
        "actions": [
            {
                "id": "ai.agent",
                "examples": [
                    "ai_agent_investigate_alert",
                    "ai_agent_with_approvals",
                ],
            },
            {
                "id": "ai.preset_agent",
                "examples": ["preset_agent_run"],
            },
        ],
    },
]


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, definitions: dict[str, Any] | None = 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, definitions)
            if part not in parts:
                parts.append(part)
        return " | ".join(parts)

    if ref := schema.get("$ref"):
        name = str(ref).rsplit("/", maxsplit=1)[-1]
        if definitions and isinstance(definitions.get(name), dict):
            return _format_type(definitions[name], definitions)
        return name

    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'), definitions)}]"
    if schema_type == "object":
        additional = schema.get("additionalProperties")
        if isinstance(additional, dict):
            return f"map[string, {_format_type(additional, definitions)}]"
        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], definitions: dict[str, Any] | None = None
) -> 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, definitions)
        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,
    definitions: dict[str, Any] | None = None,
) -> list[str]:
    field_type = _escape_text(_format_type(schema, definitions))
    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, definitions))
    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.", ""]

    definitions = expects_schema.get("$defs")
    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,
                definitions=definitions if isinstance(definitions, dict) else None,
            )
        )
    return lines


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


def _render_examples(example_ids: list[str]) -> list[str]:
    lines: list[str] = []
    for example_id in example_ids:
        example = EXAMPLES.get(example_id)
        if not isinstance(example, dict):
            raise ValueError(f"Unknown example {example_id!r}.")
        language = str(example.get("language", "yaml"))
        title = example.get("title")
        code = str(example.get("code", "")).rstrip()
        if not code:
            raise ValueError(f"Example {example_id!r} is missing code.")
        if title:
            lines.extend([f"**{_escape_text(str(title))}**", ""])
        lines.extend([f"```{language}", code, "```", ""])
    return lines


def _render_action_section(action_entry: dict[str, Any], action: Any) -> list[str]:
    action_id = action_entry["id"]
    interface = action.get_interface()
    expects_schema = interface["expects"]

    lines = [f"### `{action_id}`", ""]

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

    lines.extend([_escape_text(_get_description(action)), ""])

    lines.extend(["#### Inputs", ""])
    lines.extend(_render_input_fields(expects_schema))
    lines.extend(["#### Examples", ""])
    lines.extend(_render_examples(action_entry["examples"]))
    return lines


def _render_page(page: dict[str, Any], actions_by_id: dict[str, Any]) -> str:
    lines = [
        "---",
        f'title: "{page["title"]}"',
        f'description: "{page["description"]}"',
        "---",
        "",
        AUTO_GENERATED_COMMENT,
        "",
        str(page["body"]).strip(),
        "",
    ]

    for action_entry in page["actions"]:
        action_id = action_entry["id"]
        lines.extend(_render_action_section(action_entry, actions_by_id[action_id]))

    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 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=False)

    actions_by_id: dict[str, Any] = {}
    for page in PAGES:
        for action_entry in page["actions"]:
            action_id = action_entry["id"]
            if not isinstance(action_id, str):
                raise ValueError(f"Action entry {action_entry!r} is missing an id.")
            actions_by_id[action_id] = repo.get(action_id)
            example_ids = action_entry.get("examples")
            if not isinstance(example_ids, list) or not example_ids:
                raise ValueError(
                    f"Action {action_id!r} must define at least one example."
                )

    changed_paths: list[str] = []
    for page in PAGES:
        output_path = DOCS_DIR / f"{page['slug']}.mdx"
        content = _render_page(page=page, actions_by_id=actions_by_id)
        if _write_if_changed(output_path, content, check=args.check):
            changed_paths.append(str(output_path.relative_to(ROOT_DIR)))

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


if __name__ == "__main__":
    main()
