"""Vorim SDK — OpenAI Agents SDK Python integration.

Wraps OpenAI function calling with Vorim permission checks, audit trails,
and agent identity. Works with the OpenAI Python SDK (chat completions
with tool use).

Requires: ``pip install vorim[openai]``

Usage::

    from vorim import Vorim
    from vorim.integrations.openai_agents import VorimToolRegistry

    vorim = Vorim(api_key="agid_sk_live_...")

    registry = VorimToolRegistry(vorim=vorim, agent_id="agid_acme_...")
    registry.add(
        name="search_docs",
        description="Search internal documents",
        parameters={"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]},
        execute=lambda args: search_docs(args["query"]),
        permission="agent:read",
    )
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from typing import Any, Callable

from ..client import Vorim
from ..types import AuditEventInput, PermissionScope


# ── Tool Definition ──────────────────────────────────────────────────────


@dataclass
class VorimToolDefinition:
    """Definition of a tool managed by the Vorim registry."""

    name: str
    description: str
    parameters: dict[str, Any]
    execute: Callable[[Any], Any]
    permission: PermissionScope = "agent:execute"
    strict: bool = False


# ── Tool Registry ────────────────────────────────────────────────────────


class VorimToolRegistry:
    """Manages tools with Vorim permission checks and audit logging.

    Converts tools to OpenAI's ``ChatCompletionTool`` format and handles
    execution of tool calls from the model's response.

    Usage::

        from openai import OpenAI
        from vorim import Vorim
        from vorim.integrations.openai_agents import VorimToolRegistry

        vorim = Vorim(api_key="agid_sk_live_...")
        client = OpenAI()

        registry = VorimToolRegistry(vorim=vorim, agent_id="agid_acme_a1b2c3d4")

        registry.add(
            name="search_docs",
            description="Search internal documents",
            parameters={
                "type": "object",
                "properties": {"query": {"type": "string"}},
                "required": ["query"],
            },
            execute=lambda args: search_docs(args["query"]),
            permission="agent:read",
        )

        response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            tools=registry.to_openai_tools(),
        )

        tool_messages = registry.execute_tool_calls(
            response.choices[0].message.tool_calls or []
        )
    """

    def __init__(
        self,
        vorim: Vorim,
        agent_id: str,
        default_permission: PermissionScope = "agent:execute",
    ) -> None:
        self._vorim = vorim
        self._agent_id = agent_id
        self._default_permission = default_permission
        self._tools: dict[str, VorimToolDefinition] = {}

    def add(
        self,
        name: str,
        description: str,
        parameters: dict[str, Any],
        execute: Callable[[Any], Any],
        permission: PermissionScope | None = None,
        strict: bool = False,
    ) -> "VorimToolRegistry":
        """Register a tool."""
        self._tools[name] = VorimToolDefinition(
            name=name,
            description=description,
            parameters=parameters,
            execute=execute,
            permission=permission or self._default_permission,
            strict=strict,
        )
        return self

    def add_all(self, definitions: list[VorimToolDefinition]) -> "VorimToolRegistry":
        """Register multiple tools."""
        for d in definitions:
            self._tools[d.name] = d
        return self

    def to_openai_tools(self) -> list[dict[str, Any]]:
        """Convert registered tools to OpenAI's ChatCompletionTool format."""
        tools = []
        for t in self._tools.values():
            tool: dict[str, Any] = {
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.parameters,
                },
            }
            if t.strict:
                tool["function"]["strict"] = True
            tools.append(tool)
        return tools

    def execute_tool_calls(self, tool_calls: list[Any]) -> list[dict[str, Any]]:
        """Execute tool calls from an OpenAI chat completion response.

        Each call is checked against Vorim permissions and audited.
        Returns a list of tool messages ready to append to the conversation.
        """
        return [self._execute_single(tc) for tc in tool_calls]

    def _execute_single(self, tool_call: Any) -> dict[str, Any]:
        call_id = tool_call.id
        fn_name = tool_call.function.name
        fn_args_raw = tool_call.function.arguments

        definition = self._tools.get(fn_name)
        if not definition:
            return {
                "role": "tool",
                "tool_call_id": call_id,
                "content": json.dumps({"error": f"Unknown tool: {fn_name}"}),
            }

        scope = definition.permission

        try:
            args = json.loads(fn_args_raw)
        except (json.JSONDecodeError, TypeError):
            return {
                "role": "tool",
                "tool_call_id": call_id,
                "content": json.dumps({"error": f"Invalid JSON arguments for {fn_name}"}),
            }

        # 1. Permission check
        check = self._vorim.check(self._agent_id, scope)
        if not check.allowed:
            self._vorim.emit(
                agent_id=self._agent_id,
                event_type="tool_call",
                action=fn_name,
                resource=_truncate(fn_args_raw, 500),
                result="denied",
                permission=scope,
                metadata={"reason": check.reason, "framework": "openai"},
            )
            reason_suffix = f" — {check.reason}" if check.reason else ""
            return {
                "role": "tool",
                "tool_call_id": call_id,
                "content": json.dumps({"error": f"Permission denied: {scope}{reason_suffix}"}),
            }

        # 2. Execute
        start = time.monotonic()
        try:
            result = definition.execute(args)
            content = result if isinstance(result, str) else json.dumps(result)

            self._vorim.emit(
                agent_id=self._agent_id,
                event_type="tool_call",
                action=fn_name,
                resource=_truncate(fn_args_raw, 500),
                result="success",
                permission=scope,
                latency_ms=int((time.monotonic() - start) * 1000),
                metadata={"framework": "openai"},
            )

            return {"role": "tool", "tool_call_id": call_id, "content": content}
        except Exception as exc:
            err_msg = str(exc)
            self._vorim.emit(
                agent_id=self._agent_id,
                event_type="tool_call",
                action=fn_name,
                resource=_truncate(fn_args_raw, 500),
                result="error",
                permission=scope,
                latency_ms=int((time.monotonic() - start) * 1000),
                error_code=type(exc).__name__,
                metadata={"error": err_msg, "framework": "openai"},
            )
            return {
                "role": "tool",
                "tool_call_id": call_id,
                "content": json.dumps({"error": err_msg}),
            }


# ── Agent Loop ───────────────────────────────────────────────────────────


def run_agent_loop(
    *,
    vorim: Vorim,
    agent_id: str,
    openai_client: Any,
    registry: VorimToolRegistry,
    user_message: str,
    model: str = "gpt-4o",
    system_prompt: str | None = None,
    max_iterations: int = 10,
) -> str:
    """Run a complete agent loop with OpenAI function calling, Vorim
    permission enforcement, and audit logging.

    Usage::

        from openai import OpenAI

        response = run_agent_loop(
            vorim=vorim,
            agent_id="agid_acme_...",
            openai_client=OpenAI(),
            registry=registry,
            user_message="Find docs about onboarding",
            system_prompt="You are a helpful assistant.",
        )
    """
    messages: list[dict[str, Any]] = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": user_message})

    tools = registry.to_openai_tools()

    for _ in range(max_iterations):
        create_kwargs: dict[str, Any] = {"model": model, "messages": messages}
        if tools:
            create_kwargs["tools"] = tools

        response = openai_client.chat.completions.create(**create_kwargs)
        choice = response.choices[0]
        messages.append(choice.message)

        if choice.finish_reason == "stop" or not getattr(choice.message, "tool_calls", None):
            return choice.message.content or ""

        tool_messages = registry.execute_tool_calls(choice.message.tool_calls)
        messages.extend(tool_messages)

    last = messages[-1]
    return last.get("content", "") if isinstance(last, dict) else getattr(last, "content", "")


# ── Agent Registration Helper ────────────────────────────────────────────


def create_vorim_openai_agent(
    vorim: Vorim,
    *,
    name: str,
    capabilities: list[str],
    scopes: list[PermissionScope],
    tools: list[VorimToolDefinition],
    description: str | None = None,
) -> dict[str, Any]:
    """Register a Vorim agent and return a ready-to-use tool registry.

    Returns a dict with ``agent_id``, ``registration``, ``registry``,
    and ``private_key``.

    Usage::

        result = create_vorim_openai_agent(
            vorim,
            name="support-agent",
            capabilities=["search", "email"],
            scopes=["agent:read", "agent:execute"],
            tools=[search_tool_def, email_tool_def],
        )
        registry = result["registry"]
    """
    reg = vorim.register(
        name=name,
        description=description,
        capabilities=capabilities,
        scopes=scopes,
    )
    agent_id = reg.agent.agent_id

    registry = VorimToolRegistry(vorim=vorim, agent_id=agent_id)
    registry.add_all(tools)

    return {
        "agent_id": agent_id,
        "registration": reg,
        "registry": registry,
        "private_key": reg.private_key,
    }


# ── Helpers ──────────────────────────────────────────────────────────────


def _truncate(s: str, max_len: int) -> str:
    return s[:max_len] + "…" if len(s) > max_len else s
