"""顧客対応エージェント。LLM（OpenAI 互換 API。既定はローカルの Ollama）と、2つの MCP サーバーを使う。

- 参照系ツール: http://localhost:8080/mcp/read   （banking-reader）
- 更新系ツール: http://localhost:8080/mcp/write  （banking-writer。権限がなければ接続自体が 403）
"""

import json
import os
import re
import uuid
from dataclasses import dataclass, field

from openai import OpenAI
from opentelemetry import trace
from opentelemetry.trace import SpanKind

from . import auth
from .mcp_client import McpClient, McpHttpError, ToolResult
from .telemetry import trace_id_hex

TOOLS_URL = os.getenv("AGENT_TOOLS_URL", "http://localhost:8080")
LLM_BASE_URL = os.getenv("LLM_BASE_URL", "http://localhost:11434/v1")
LLM_MODEL = os.getenv("LLM_MODEL", "qwen3:8b")
LLM_API_KEY = os.getenv("LLM_API_KEY", "ollama")

SYSTEM_PROMPT = """あなたは銀行の顧客対応エージェントです。サインイン中の顧客本人のために、ツールだけを使って口座の照会と振込の依頼を行います。
守ること:
- 残高・明細・振込の状態は必ずツールで確認する。推測で答えない。
- 振込は request_transfer を1回だけ呼ぶ。金額は日本円の整数。
- 承認待ち（PENDING_APPROVAL）になったら、人間の承認が必要だと伝える。承認を回避しようとしない。
- ツールがエラーを返したら、その理由を顧客に分かる言葉で伝える。別の方法で同じことを試みない。
- 回答は日本語で簡潔に。"""

tracer = trace.get_tracer("asc_agent")


@dataclass
class ToolCall:
    name: str
    arguments: dict
    is_error: bool
    result: str


@dataclass
class AgentRun:
    answer: str
    tool_calls: list[ToolCall] = field(default_factory=list)
    trace_id: str = ""


class ConciergeAgent:
    def __init__(self, user: str, password: str | None = None, model: str = LLM_MODEL):
        self.user = user
        self.model = model
        self.session_id = uuid.uuid4()
        token = auth.token(user, password)
        self.clients: dict[str, McpClient] = {}
        self.routes: dict[str, McpClient] = {}
        self.unavailable: dict[str, str] = {}
        for name in ("read", "write"):
            client = McpClient(f"{TOOLS_URL}/mcp/{name}", token, name)
            try:
                client.initialize()
            except McpHttpError as e:
                # 権限のないツール群は、そもそもエージェントに見せない（最小権限）
                self.unavailable[name] = f"HTTP {e.status}"
                continue
            self.clients[name] = client
            for t in client.list_tools():
                self.routes[t["name"]] = client
        self._tools = self._openai_tools()
        self.llm = OpenAI(base_url=LLM_BASE_URL, api_key=LLM_API_KEY)

    def _openai_tools(self) -> list[dict]:
        out = []
        for client in self.clients.values():
            for t in client.list_tools():
                schema = dict(t.get("inputSchema") or {"type": "object", "properties": {}})
                if t["name"] == "request_transfer":
                    # 冪等キーはエージェント側で決める（LLM に任せない）
                    schema = json.loads(json.dumps(schema))
                    schema.get("properties", {}).pop("idempotencyKey", None)
                    if "required" in schema:
                        schema["required"] = [r for r in schema["required"] if r != "idempotencyKey"]
                out.append({"type": "function", "function": {
                    "name": t["name"], "description": t.get("description", ""), "parameters": schema}})
        return out

    def idempotency_key(self, args: dict) -> str:
        """同じ会話で同じ振込内容なら同じキー。LLM が同じ依頼を二度出しても二重送金にならない。"""
        canonical = json.dumps({k: v for k, v in sorted(args.items()) if k != "idempotencyKey"}, ensure_ascii=False)
        return str(uuid.uuid5(self.session_id, canonical))

    def call(self, name: str, args: dict) -> tuple[dict, ToolResult]:
        """ツールを呼ぶ。実際に送った引数（冪等キーを差し替えた後）と結果を返す。"""
        client = self.routes.get(name)
        if client is None:
            return args, ToolResult(True, f"tool '{name}' is not available to this agent")
        if name == "request_transfer":
            args = {**args, "idempotencyKey": self.idempotency_key(args)}
        return args, client.call_tool(name, args)

    def run(self, prompt: str, max_steps: int = 6) -> AgentRun:
        with tracer.start_as_current_span("invoke_agent concierge", kind=SpanKind.INTERNAL) as span:
            span.set_attribute("gen_ai.operation.name", "invoke_agent")
            span.set_attribute("gen_ai.agent.name", "concierge")
            span.set_attribute("enduser.id", self.user)
            run = AgentRun(answer="", trace_id=trace_id_hex())
            messages = [{"role": "system", "content": SYSTEM_PROMPT},
                        {"role": "user", "content": prompt + " /no_think"}]
            for _ in range(max_steps):
                with tracer.start_as_current_span(f"chat {self.model}", kind=SpanKind.CLIENT) as chat:
                    chat.set_attribute("gen_ai.operation.name", "chat")
                    chat.set_attribute("gen_ai.request.model", self.model)
                    resp = self.llm.chat.completions.create(model=self.model, messages=messages,
                                                            tools=self._tools, temperature=0)
                    if resp.usage:
                        chat.set_attribute("gen_ai.usage.input_tokens", resp.usage.prompt_tokens)
                        chat.set_attribute("gen_ai.usage.output_tokens", resp.usage.completion_tokens)
                msg = resp.choices[0].message
                if not msg.tool_calls:
                    run.answer = re.sub(r"<think>.*?</think>", "", msg.content or "", flags=re.S).strip()
                    return run
                messages.append({"role": "assistant", "content": msg.content or "",
                                 "tool_calls": [tc.model_dump() for tc in msg.tool_calls]})
                for tc in msg.tool_calls:
                    try:
                        args = json.loads(tc.function.arguments or "{}")
                    except ValueError:
                        args = {}
                    sent, result = self.call(tc.function.name, args)
                    run.tool_calls.append(ToolCall(tc.function.name, sent, result.is_error, result.text))
                    messages.append({"role": "tool", "tool_call_id": tc.id,
                                     "content": ("ERROR: " if result.is_error else "") + result.text})
            run.answer = "（ステップ数の上限に達したため、処理を打ち切りました）"
            return run

    def close(self) -> None:
        for c in self.clients.values():
            c.close()