"""原生 function-calling 阶段 1:provider chat_with_tools + 工具 JSON schema。

替代 react-over-text 的"发文本求 JSON、正则抠"——从根上消灭 0 执行/格式飘类 bug。
见 docs/NATIVE_FUNCTION_CALLING_DESIGN.md。阶段 1 只测基础层,不动 react loop(零回归)。
"""
from __future__ import annotations

import json
import unittest
from unittest.mock import patch, MagicMock

from lambdagent.providers.openai_compat_provider import (
    OpenAICompatProvider, _parse_tool_calls,
)
from lambdagent.providers.base import ProviderConfig
from lambdagent.fromconfig.compiler import _tools_json_schema


class TestToolCallParsing(unittest.TestCase):
    def test_dict_args(self):
        out = _parse_tool_calls([{"id": "c1", "function": {
            "name": "ReadFile", "arguments": {"file_path": "/a.txt"}}}])
        self.assertEqual(out, [{"id": "c1", "name": "ReadFile",
                                "arguments": {"file_path": "/a.txt"}}])

    def test_string_json_args(self):
        out = _parse_tool_calls([{"id": "c2", "function": {
            "name": "Bash", "arguments": '{"command": "ls -la"}'}}])
        self.assertEqual(out[0]["arguments"], {"command": "ls -la"})

    def test_bad_or_empty_args(self):
        self.assertEqual(_parse_tool_calls([{"function": {"name": "X", "arguments": "not json"}}])[0]["arguments"], {})
        self.assertEqual(_parse_tool_calls([{"function": {"name": "X", "arguments": ""}}])[0]["arguments"], {})
        self.assertEqual(_parse_tool_calls(None), [])
        self.assertEqual(_parse_tool_calls([]), [])


class TestToolsJsonSchema(unittest.TestCase):
    def test_schema_for_file_tools(self):
        cfg = {"mcp": {"localTools": ["ReadFile", "WriteFile", "terminate"]}}
        specs = _tools_json_schema(cfg)
        by_name = {s["function"]["name"]: s["function"] for s in specs}
        self.assertIn("ReadFile", by_name)
        self.assertIn("WriteFile", by_name)
        self.assertIn("terminate", by_name)
        # ReadFile: file_path required, offset/limit optional
        rf = by_name["ReadFile"]
        self.assertEqual(rf["parameters"]["type"], "object")
        self.assertIn("file_path", rf["parameters"]["properties"])
        self.assertIn("file_path", rf["parameters"]["required"])
        self.assertNotIn("offset", rf["parameters"]["required"])
        # 类型推断:offset → integer
        self.assertEqual(rf["parameters"]["properties"]["offset"]["type"], "integer")
        # WriteFile: file_path + content 都 required
        wf = by_name["WriteFile"]
        self.assertIn("content", wf["parameters"]["required"])

    def test_empty_when_no_tools(self):
        self.assertEqual(_tools_json_schema({}), [])


class TestChatWithTools(unittest.TestCase):
    def _provider(self):
        p = OpenAICompatProvider.__new__(OpenAICompatProvider)
        p.config = ProviderConfig(model="qwen-plus", temperature=0.0, max_tokens=4096, timeout=60)
        p._provider_name = "dashscope"
        p.base_url = "https://x/v1"
        p.api_key = "k"
        p._usage_input = 0
        p._usage_output = 0
        return p

    def _mock_resp(self, payload):
        cm = MagicMock()
        cm.__enter__.return_value.read.return_value = json.dumps(payload).encode()
        return cm

    def test_returns_structured_tool_calls(self):
        p = self._provider()
        payload = {"choices": [{"finish_reason": "tool_calls", "message": {
            "content": None,
            "tool_calls": [{"id": "c1", "function": {
                "name": "WriteFile", "arguments": '{"file_path":"/out.md","content":"x"}'}}]}}],
            "usage": {"prompt_tokens": 100, "completion_tokens": 20}}
        with patch("urllib.request.urlopen", return_value=self._mock_resp(payload)):
            r = p.chat_with_tools([{"role": "user", "content": "write it"}],
                                  tools=[{"type": "function", "function": {"name": "WriteFile"}}])
        self.assertEqual(r["tool_calls"][0]["name"], "WriteFile")
        self.assertEqual(r["tool_calls"][0]["arguments"], {"file_path": "/out.md", "content": "x"})
        self.assertIsNone(r["content"])
        self.assertEqual(r["usage"]["input_tokens"], 100)
        self.assertEqual(p._usage_input, 100)  # 累计进 get_usage

    def test_content_only_no_tool_calls(self):
        p = self._provider()
        payload = {"choices": [{"finish_reason": "stop", "message": {"content": "答案"}}],
                   "usage": {"prompt_tokens": 5, "completion_tokens": 3}}
        with patch("urllib.request.urlopen", return_value=self._mock_resp(payload)):
            r = p.chat_with_tools([{"role": "user", "content": "hi"}], tools=None)
        self.assertEqual(r["content"], "答案")
        self.assertEqual(r["tool_calls"], [])

    def test_supports_fc(self):
        self.assertTrue(self._provider().supports_function_calling())


class _FakeFCProvider:
    """mock chat_with_tools,按序返回 canned 响应。"""
    def __init__(self, responses):
        self.responses = responses; self.i = 0
    def chat_with_tools(self, messages, tools_schema=None, **kw):
        r = self.responses[min(self.i, len(self.responses) - 1)]; self.i += 1; return r
    def supports_function_calling(self): return True
    def get_usage(self): return {"input_tokens": 0, "output_tokens": 0}


class TestFCLoop(unittest.TestCase):
    def test_executes_tool_then_terminates(self):
        from lambdagent.fromconfig.compiler import _run_fc_loop
        calls = []
        tools = {"WriteFile": lambda args: calls.append(args) or "wrote ok"}
        prov = _FakeFCProvider([
            {"content": None, "tool_calls": [{"id": "c1", "name": "WriteFile",
              "arguments": {"file_path": "/a.md", "content": "x"}}]},
            {"content": None, "tool_calls": [{"id": "c2", "name": "terminate",
              "arguments": {"summary": "完成"}}]},
        ])
        out = _run_fc_loop(prov, "sys", "do it", tools, [], max_steps=5, tool_timeout=10)
        self.assertEqual(calls, [{"file_path": "/a.md", "content": "x"}])  # 工具真执行
        self.assertEqual(out, "完成")

    def test_content_only_is_final(self):
        from lambdagent.fromconfig.compiler import _run_fc_loop
        prov = _FakeFCProvider([{"content": "答案", "tool_calls": []}])
        self.assertEqual(_run_fc_loop(prov, "s", "q", {}, [], max_steps=5, tool_timeout=10), "答案")

    def test_unknown_tool_then_terminate(self):
        from lambdagent.fromconfig.compiler import _run_fc_loop
        prov = _FakeFCProvider([
            {"content": None, "tool_calls": [{"id": "c1", "name": "Nope", "arguments": {}}]},
            {"content": None, "tool_calls": [{"id": "c2", "name": "terminate", "arguments": {"summary": "end"}}]},
        ])
        # 未知工具 → error obs 回灌,循环继续到 terminate(不崩)
        self.assertEqual(_run_fc_loop(prov, "s", "q", {}, [], max_steps=5, tool_timeout=10), "end")


class TestFCRouting(unittest.TestCase):
    def _cfg(self, fc):
        return {"name": "t", "type": "react", "systemPrompt": "s",
                "model": {"provider": "dashscope", "name": "qwen-plus"},
                "react": {"maxSteps": 5, "nativeToolCalls": fc},
                "mcp": {"localTools": ["ReadFile", "terminate"]}}

    def test_fc_term_when_flag_and_capable(self):
        from lambdagent.fromconfig.compiler import _compile_react
        term = _compile_react(self._cfg(True), {})
        self.assertTrue(term._name.endswith("fc_react"))

    def test_text_term_when_flag_off(self):
        from lambdagent.fromconfig.compiler import _compile_react
        term = _compile_react(self._cfg(False), {})
        self.assertFalse(term._name.endswith("fc_react"))


class TestFCEnforceLoop(unittest.TestCase):
    def test_terminate_blocked_until_min_count(self):
        from lambdagent.fromconfig.compiler import _run_fc_loop
        calls = []
        tools = {"WriteFile": lambda a: calls.append(a) or "ok"}
        prov = _FakeFCProvider([
            {"content": None, "tool_calls": [{"id": "t1", "name": "terminate", "arguments": {"summary": "早退"}}]},
            {"content": None, "tool_calls": [{"id": "w1", "name": "WriteFile", "arguments": {"file_path": "/a", "content": "x"}}]},
            {"content": None, "tool_calls": [{"id": "t2", "name": "terminate", "arguments": {"summary": "完成"}}]},
        ])
        out = _run_fc_loop(prov, "s", "q", tools, [], max_steps=6, tool_timeout=10,
                           enforce_tool="WriteFile", enforce_min=1)
        self.assertEqual(len(calls), 1)   # 第一次 terminate 被挡 → 先调了 WriteFile
        self.assertEqual(out, "完成")      # 满足后第二次 terminate 才放行


class TestFCTrace(unittest.TestCase):
    def test_logs_per_tool_to_ctx_trace(self):
        from lambdagent.fromconfig.compiler import _run_fc_loop
        from lambdagent.core import Context
        ctx = Context()
        prov = _FakeFCProvider([
            {"content": None, "tool_calls": [{"id": "w1", "name": "WriteFile", "arguments": {"file_path": "/a", "content": "x"}}]},
            {"content": None, "tool_calls": [{"id": "t1", "name": "terminate", "arguments": {"summary": "d"}}]},
        ])
        _run_fc_loop(prov, "s", "q", {"WriteFile": lambda a: "ok"}, [],
                     max_steps=4, tool_timeout=10, ctx=ctx)
        assert any(e.term_name == "fc:WriteFile" for e in ctx.trace)  # per-tool trace 写了
