"""separate built-ins, resources, and per-use tool bindings

Revision ID: 20260813_0021
Revises: 20260812_0020
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from copy import deepcopy
import json
import re

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql


revision: str = "20260813_0021"
down_revision: str | Sequence[str] | None = "20260812_0020"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


PLACEHOLDER = re.compile(r"{{\s*([A-Za-z][A-Za-z0-9_]*)\s*}}")
FULL_PLACEHOLDER = re.compile(r"^{{\s*([A-Za-z][A-Za-z0-9_]*)\s*}}$")


def _source_for(value) -> dict:
    raw = json.dumps(value, ensure_ascii=False)
    return (
        {"source": "template", "value": value}
        if PLACEHOLDER.search(raw)
        else {"source": "literal", "value": value}
    )


def _llm_rule(name: str, definition: Mapping | None) -> dict:
    value_type = str((definition or {}).get("type") or "string")
    return {
        "variableName": name,
        "expression": {
            "type": "llm",
            "parameterName": name,
            "valueType": value_type,
            "description": f"从当前对话中提取并确认 {name}",
        },
    }


def _builtin(
    kind: str,
    tool: Mapping,
    definitions: Mapping,
    allowed=None,
    *,
    camel_case: bool = False,
) -> dict:
    definition = tool.get("definition") or {}
    config = deepcopy(definition.get("config") or {})
    config.pop("kind", None)
    config["description"] = str(tool.get("description") or "")
    if kind == "update_state":
        names = list(allowed) if allowed is not None else list(definitions)
        config["updates"] = [
            _llm_rule(str(name), definitions.get(str(name)))
            for name in names
            if str(name) in definitions
        ]
    if camel_case:
        aliases = {
            "message_type": "messageType",
            "custom_message": "customMessage",
            "capture_reason": "captureReason",
            "allow_end_interruption": "allowEndInterruption",
        }
        config = {aliases.get(key, key): value for key, value in config.items()}
    return config


def _builtin_key(kind: str) -> str:
    return {
        "end_conversation": "endConversation",
        "update_state": "updateState",
        "skip_turn": "skipTurn",
        "request_human_handoff": "requestHumanHandoff",
    }.get(kind, kind)


def _binding(
    tool_id: str,
    tool: Mapping,
    arguments=None,
    assignments=None,
    inherited_assignments: Mapping | None = None,
) -> dict:
    config = ((tool.get("definition") or {}).get("config") or {})
    if tool.get("type") == "mcp":
        schema = config.get("input_schema") or {}
        names = list((schema.get("properties") or {}).keys())
    else:
        names = [
            str(parameter.get("name"))
            for parameter in config.get("parameters") or []
            if parameter.get("name")
        ]
    if arguments is None:
        argument_bindings = {name: {"source": "model"} for name in names}
    else:
        argument_bindings = {
            str(name): _source_for(value) for name, value in arguments.items()
        }
    return {
        "toolId": tool_id,
        "argumentBindings": argument_bindings,
        "templateBindings": {},
        "resultAssignments": dict(
            (inherited_assignments or {}) if assignments is None else assignments
        ),
    }


def _state_expression(value) -> dict:
    """Translate deterministic v3 values without silently changing semantics."""
    if isinstance(value, str):
        match = FULL_PLACEHOLDER.fullmatch(value)
        if match:
            return {"type": "variable", "variableName": match.group(1)}
        if PLACEHOLDER.search(value):
            raise RuntimeError(
                "Update State 节点中的混合字符串模板无法无损迁移；"
                "请在升级前改成完整变量占位符"
            )
    if not isinstance(value, (str, int, float, bool)):
        raise RuntimeError("Update State 节点只支持字符串、数字、布尔值或变量")
    return {"type": "constant", "value": value}


def _rewrite_templates(value, inputs: set[str]):
    if isinstance(value, str):
        def replace(match):
            name = match.group(1)
            if name.startswith(("system__", "secret__")):
                return match.group(0)
            if name.startswith("input__"):
                inputs.add(name.removeprefix("input__"))
                return match.group(0)
            inputs.add(name)
            return "{{input__" + name + "}}"
        return PLACEHOLDER.sub(replace, value)
    if isinstance(value, list):
        return [_rewrite_templates(item, inputs) for item in value]
    if isinstance(value, dict):
        return {key: _rewrite_templates(item, inputs) for key, item in value.items()}
    return value


def _add_template_bindings(binding: dict, inputs: set[str]) -> None:
    binding.setdefault("templateBindings", {}).update(
        {
            name: {"source": "variable", "variableName": name}
            for name in sorted(inputs)
        }
    )


def _freeze_inherited_agent_for_update_state(
    data: dict,
    settings: Mapping,
    definitions: Mapping,
    allowed: list,
) -> None:
    """Preserve the old per-Agent Update State write scope.

    Before v4 an Agent could inherit the global Update State resource while its
    writable variable names still lived on the node. A v4 built-in owns its
    complete configuration, so this one exceptional case must snapshot the
    current global stage configuration onto the Agent.
    """

    runtime_mode = str(settings.get("runtimeMode") or "pipeline")
    if runtime_mode != "realtime":
        for source, target in (
            ("defaultLlmResourceId", "llmResourceId"),
            ("defaultAsrResourceId", "asrResourceId"),
            ("defaultTtsResourceId", "ttsResourceId"),
        ):
            data[target] = settings.get(source)
        for key in ("enableInterrupt", "turnConfig"):
            if key in settings:
                data[key] = deepcopy(settings[key])
    for key in (
        "visionEnabled",
        "visionModelResourceId",
        "visionPromptMode",
        "visionPrompt",
        "knowledgeBaseId",
        "knowledgeMode",
        "knowledgeTopN",
        "knowledgeScoreThreshold",
    ):
        if key in settings:
            data[key] = deepcopy(settings[key])
    data["toolBindings"] = deepcopy(settings.get("toolBindings") or [])
    data["builtInTools"] = deepcopy(settings.get("builtInTools") or {})
    data["builtInTools"]["updateState"]["updates"] = [
        _llm_rule(str(name), definitions.get(str(name)))
        for name in allowed
        if str(name) in definitions
    ]
    global_prompt = str(settings.get("globalPrompt") or "").strip()
    node_prompt = str(data.get("prompt") or "").strip()
    if global_prompt:
        data["prompt"] = "\n\n".join(part for part in (global_prompt, node_prompt) if part)
    data["inheritGlobalConfig"] = False


def upgrade() -> None:
    connection = op.get_bind()
    op.add_column(
        "assistants",
        sa.Column(
            "built_in_tools",
            postgresql.JSONB(astext_type=sa.Text()),
            server_default=sa.text("'{}'::jsonb"),
            nullable=False,
        ),
    )
    op.add_column(
        "assistant_tool_bindings",
        sa.Column(
            "config",
            postgresql.JSONB(astext_type=sa.Text()),
            server_default=sa.text("'{}'::jsonb"),
            nullable=False,
        ),
    )

    tool_rows = connection.execute(
        sa.text("SELECT id, type, description, definition, mcp_server_id FROM tools")
    ).mappings().all()
    tools = {str(row["id"]): dict(row) for row in tool_rows}
    template_inputs: dict[str, set[str]] = {}
    resource_assignments: dict[str, dict] = {}
    server_inputs: dict[str, set[str]] = {}
    server_rows = connection.execute(
        sa.text("SELECT id, url, config FROM mcp_servers")
    ).mappings().all()
    for row in server_rows:
        inputs: set[str] = set()
        config = deepcopy(row["config"] or {})
        url = _rewrite_templates(str(row["url"] or ""), inputs)
        config["headers"] = _rewrite_templates(config.get("headers") or {}, inputs)
        server_id = str(row["id"])
        server_inputs[server_id] = inputs
        connection.execute(
            sa.text(
                "UPDATE mcp_servers SET url=:url, config=CAST(:config AS jsonb) "
                "WHERE id=:id"
            ),
            {"id": server_id, "url": url, "config": json.dumps(config)},
        )
    for tool_id, tool in tools.items():
        if tool.get("type") == "system":
            continue
        definition = deepcopy(tool.get("definition") or {})
        config = definition.get("config") or {}
        definition["config"] = config
        resource_assignments[tool_id] = dict(
            config.get("dynamic_variable_assignments") or {}
        )
        config.pop("dynamic_variable_assignments", None)
        inputs: set[str] = set()
        if tool.get("type") == "http":
            for key in ("url", "headers", "body"):
                config[key] = _rewrite_templates(config.get(key), inputs)
            config["template_inputs"] = [
                {"name": name, "type": "string", "location": "body", "required": True}
                for name in sorted(inputs)
            ]
        elif tool.get("type") == "mcp":
            inputs.update(server_inputs.get(str(tool.get("mcp_server_id") or ""), set()))
            config["template_inputs"] = [
                {
                    "name": name,
                    "type": "string",
                    "location": "body",
                    "required": True,
                }
                for name in sorted(inputs)
            ]
        definition["schema_version"] = 2
        template_inputs[tool_id] = inputs
        connection.execute(
            sa.text("UPDATE tools SET definition=CAST(:definition AS jsonb) WHERE id=:id"),
            {"id": tool_id, "definition": json.dumps(definition)},
        )
        tool["definition"] = definition

    assistant_rows = connection.execute(
        sa.text(
            "SELECT id, type, runtime_mode, graph, startup, "
            "dynamic_variable_definitions FROM assistants"
        )
    ).mappings().all()
    for assistant in assistant_rows:
        assistant_id = str(assistant["id"])
        definitions = assistant["dynamic_variable_definitions"] or {}
        binding_rows = connection.execute(
            sa.text(
                "SELECT tool_id FROM assistant_tool_bindings "
                "WHERE assistant_id=:assistant_id ORDER BY created_at, tool_id"
            ),
            {"assistant_id": assistant_id},
        ).mappings().all()
        selected_ids = [str(row["tool_id"]) for row in binding_rows]
        system_by_kind: dict[str, Mapping] = {}
        for tool_id in selected_ids:
            tool = tools.get(tool_id) or {}
            kind = ((tool.get("definition") or {}).get("config") or {}).get("kind")
            if tool.get("type") == "system" and kind:
                if kind in system_by_kind:
                    raise RuntimeError(
                        f"Assistant {assistant_id} 同时绑定多个 {kind} System Tool"
                    )
                system_by_kind[str(kind)] = tool

        builtins = (
            {
                kind: _builtin(kind, tool, definitions)
                for kind, tool in system_by_kind.items()
            }
            if (
                assistant["type"] == "prompt"
                and assistant["runtime_mode"] != "realtime"
            )
            else {}
        )
        for tool_id in selected_ids:
            tool = tools.get(tool_id) or {}
            if tool.get("type") == "system":
                continue
            binding = _binding(
                tool_id,
                tool,
                inherited_assignments=resource_assignments.get(tool_id),
            )
            _add_template_bindings(binding, template_inputs.get(tool_id, set()))
            connection.execute(
                sa.text(
                    "UPDATE assistant_tool_bindings SET config=CAST(:config AS jsonb) "
                    "WHERE assistant_id=:assistant_id AND tool_id=:tool_id"
                ),
                {
                    "assistant_id": assistant_id,
                    "tool_id": tool_id,
                    "config": json.dumps(
                        {
                            "argument_bindings": binding["argumentBindings"],
                            "template_bindings": binding["templateBindings"],
                            "result_assignments": binding["resultAssignments"],
                        }
                    ),
                },
            )

        startup = deepcopy(assistant["startup"] or {})
        for action in startup.get("actions") or []:
            tool_id = str(action.pop("toolId", action.pop("tool_id", "")) or "")
            arguments = action.pop("arguments", {}) or {}
            tool = tools.get(tool_id) or {}
            binding = _binding(
                tool_id,
                tool,
                arguments=arguments,
                inherited_assignments=resource_assignments.get(tool_id),
            )
            _add_template_bindings(binding, template_inputs.get(tool_id, set()))
            action["toolBinding"] = binding

        graph = deepcopy(assistant["graph"] or {})
        if assistant["type"] == "workflow":
            graph["specVersion"] = 4
            settings = graph.setdefault("settings", {})
            global_ids = list(settings.pop("toolIds", []) or [])
            settings["toolBindings"] = []
            settings["builtInTools"] = {}
            for tool_id in global_ids:
                tool_id = str(tool_id)
                tool = tools.get(tool_id) or {}
                if tool.get("type") == "system":
                    kind = ((tool.get("definition") or {}).get("config") or {}).get("kind")
                    if kind:
                        settings["builtInTools"][_builtin_key(str(kind))] = _builtin(
                            str(kind), tool, definitions, camel_case=True
                        )
                    continue
                binding = _binding(
                    tool_id,
                    tool,
                    inherited_assignments=resource_assignments.get(tool_id),
                )
                _add_template_bindings(binding, template_inputs.get(tool_id, set()))
                settings["toolBindings"].append(binding)
            for node in graph.get("nodes") or []:
                data = node.setdefault("data", {})
                if node.get("type") == "agent":
                    inherits_global = bool(data.get("inheritGlobalConfig", True))
                    ids = list(data.pop("toolIds", []) or [])
                    allowed = list(data.pop("stateVariableNames", []) or [])
                    data["toolBindings"] = []
                    data["builtInTools"] = {}
                    if inherits_global:
                        if "updateState" in settings["builtInTools"]:
                            _freeze_inherited_agent_for_update_state(
                                data, settings, definitions, allowed
                            )
                    else:
                        for tool_id in ids:
                            tool_id = str(tool_id)
                            tool = tools.get(tool_id) or {}
                            if tool.get("type") == "system":
                                kind = ((tool.get("definition") or {}).get("config") or {}).get("kind")
                                if kind:
                                    data["builtInTools"][_builtin_key(str(kind))] = _builtin(
                                        str(kind),
                                        tool,
                                        definitions,
                                        allowed if kind == "update_state" else None,
                                        camel_case=True,
                                    )
                                continue
                            binding = _binding(
                                tool_id,
                                tool,
                                inherited_assignments=resource_assignments.get(tool_id),
                            )
                            _add_template_bindings(
                                binding, template_inputs.get(tool_id, set())
                            )
                            data["toolBindings"].append(binding)
                elif node.get("type") == "action":
                    tool_id = str(data.pop("toolId", "") or "")
                    arguments = data.pop("arguments", {}) or {}
                    mode = str(data.pop("resultAssignmentMode", "none") or "none")
                    overrides = data.pop("resultAssignments", {}) or {}
                    assignments = (
                        None if mode == "inherit" else overrides if mode == "override" else {}
                    )
                    binding = _binding(
                        tool_id,
                        tools.get(tool_id) or {},
                        arguments,
                        assignments,
                        resource_assignments.get(tool_id),
                    )
                    _add_template_bindings(binding, template_inputs.get(tool_id, set()))
                    data["toolBinding"] = binding
                elif node.get("type") == "update_state":
                    assignments = data.pop("assignments", {}) or {}
                    data["updates"] = [
                        {
                            "variableName": str(name),
                            "expression": _state_expression(value),
                        }
                        for name, value in assignments.items()
                    ]

        connection.execute(
            sa.text(
                "UPDATE assistants SET built_in_tools=CAST(:builtins AS jsonb), "
                "startup=CAST(:startup AS json), graph=CAST(:graph AS json) WHERE id=:id"
            ),
            {
                "id": assistant_id,
                "builtins": json.dumps(builtins),
                "startup": json.dumps(startup),
                "graph": json.dumps(graph),
            },
        )

        if assistant["type"] == "workflow":
            connection.execute(
                sa.text(
                    "UPDATE assistant_tool_bindings SET config='{}'::jsonb "
                    "WHERE assistant_id=:assistant_id"
                ),
                {"assistant_id": assistant_id},
            )

    connection.execute(
        sa.text(
            "DELETE FROM assistant_tool_bindings "
            "WHERE tool_id IN (SELECT id FROM tools WHERE type='system')"
        )
    )
    connection.execute(
        sa.text("DELETE FROM tools WHERE type='system'")
    )
    op.create_check_constraint(
        "ck_tools_external_type",
        "tools",
        "type IN ('http', 'mcp', 'client')",
    )


def downgrade() -> None:
    raise RuntimeError(
        "Tool Resource v4 是一次性边界迁移；请从部署前数据库备份回滚。"
    )