from __future__ import annotations

from genesis.inference.engine import InferenceEngine
from genesis.inference.exceptions import LLMError
from genesis.inference.prompt_registry import PromptRegistry
from genesis.inference.token import approximate_token_count
from genesis.inference.types import Prompt
from tests.inference_helpers import FakeLLMClient


class Models:
    """模拟 config.InferenceModels（pydantic 结构，测试中直接构建）"""
    def __init__(self):
        import types as t
        self.primary = t.SimpleNamespace(name="deepseek-chat", provider="deepseek")
        self.fallback = t.SimpleNamespace(name="qwen-max", provider="qwen")


class ModelsNoFallback:
    """无 fallback 的模型配置（覆盖 getattr(fallback) 假值分支）"""
    def __init__(self):
        import types as t
        self.primary = t.SimpleNamespace(name="deepseek-chat", provider="deepseek")
        self.fallback = None


class ModelsEmpty:
    """主/备用模型均为空（覆盖 names 为空 → 返回内置默认）"""
    def __init__(self):
        self.primary = None
        self.fallback = None


def make_engine(client=None, *, constants=None):
    eng = InferenceEngine(
        client=client or FakeLLMClient(),
        models=Models(),
        registry=PromptRegistry(),
        estimator=approximate_token_count,
    )
    if constants:
        eng._max_context_tokens = constants  # 内部测试钩子
    return eng


# ---------- chat 主路径 ----------

def test_chat_ok():
    client = FakeLLMClient([("ok", "正文")])
    eng = make_engine(client)
    r = eng.chat(
        session_id="s1",
        prompt=Prompt(name="writer", version="v1", template="章节：{{ chapter }}"),
        variables={"chapter": "DB設計"},
    )
    assert r.status == "ok" and r.text == "正文"
    assert client.calls[0]["model"] == "deepseek-chat"


def test_chat_fallback_after_primary_failure():
    client = FakeLLMClient([("raise_timeout", ""), ("ok", "备用输出")])
    eng = make_engine(client)
    r = eng.chat(
        session_id="s1",
        prompt=Prompt(name="p", version="v1", template="t:{{ x }}"),
        variables={"x": "1"},
    )
    assert r.status == "fallback"
    # 先主模型，后备模型（修正：断言首次调用为主模型）
    assert client.calls[0]["model"] == "deepseek-chat"


def test_chat_all_failed_returns_failed():
    client = FakeLLMClient([("raise_network", ""), ("raise_network", "")])
    eng = make_engine(client)
    r = eng.chat(
        session_id="s1",
        prompt=Prompt(name="p", version="v1", template="t"),
        variables={},
    )
    assert r.status == "failed" and r.error


# ---------- error_code 透传（与 error 同源取最后一次异常） ----------

def test_chat_failed_error_code_timeout():
    # 主模型超时 + 备用也失败：error_code 与 error 同源（取最后一次异常）
    client = FakeLLMClient([("raise_timeout", ""), ("raise_timeout", "")])
    eng = InferenceEngine(client=client, models=Models())
    r = eng.chat(session_id="s1", prompt=Prompt(name="p", version="v1", template="t"), variables={})
    assert r.status == "failed"
    assert r.error_code == "LLM_TIMEOUT"


def test_chat_failed_error_code_network_last():
    # 主模型超时（第一次）、备用网络失败（最后一次）→ error_code 取最后一次 = LLM_NETWORK_ERROR
    client = FakeLLMClient([("raise_timeout", ""), ("raise_network", "")])
    eng = InferenceEngine(client=client, models=Models())
    r = eng.chat(session_id="s1", prompt=Prompt(name="p", version="v1", template="t"), variables={})
    assert r.status == "failed"
    assert r.error_code == "LLM_NETWORK_ERROR"


def test_chat_failed_error_code_not_configured():
    # 备用模型 Key 缺失（最后一次）→ LLM_NOT_CONFIGURED
    # 用自定义异常客户端模拟 NotConfigured
    class NotConfiguredClient:
        def __init__(self):
            self.calls = []
        def chat(self, *, model, messages, temperature, max_tokens):
            self.calls.append(model)
            from genesis.inference.exceptions import LLMNotConfiguredError
            raise LLMNotConfiguredError("no key")
    eng = InferenceEngine(client=NotConfiguredClient(), models=Models())
    r = eng.chat(session_id="s1", prompt=Prompt(name="p", version="v1", template="t"), variables={})
    assert r.status == "failed"
    assert r.error_code == "LLM_NOT_CONFIGURED"


def test_chat_structured_parse_error_code():
    client = FakeLLMClient([("parse_fail", ""), ("parse_fail", "")])
    eng = InferenceEngine(client=client, models=Models())
    r = eng.chat_structured(
        session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
        variables={}, schema={}, retry_count=1,
    )
    assert r.status == "parse_error"
    assert r.error_code == "LLM_PARSE_ERROR"


def test_chat_structured_failed_error_code_network():
    client = FakeLLMClient([("raise_network", "")])
    eng = InferenceEngine(client=client, models=Models())
    r = eng.chat_structured(
        session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
        variables={}, schema={},
    )
    assert r.status == "failed"
    assert r.error_code == "LLM_NETWORK_ERROR"


def test_chat_plain_string_prompt():
    eng = make_engine(FakeLLMClient([("ok", "hi")]))
    r = eng.chat(session_id="s1", prompt="直接文本", variables={})
    assert r.text == "hi" and r.status == "ok"


def test_chat_truncation_callback_triggered():
    seen = {}

    def truncate_cb(prompt_text, variables):
        seen["called"] = True
        seen["len"] = len(prompt_text)
        return {**variables, "chapter": "裁剪版"}

    eng = InferenceEngine(
        client=FakeLLMClient([("ok", "x")]),
        models=Models(),
        registry=PromptRegistry(),
        estimator=approximate_token_count,
        truncate_cb=truncate_cb,
    )
    eng._max_context_tokens = 2  # 强制超限（复习：'abcd很长很长的标题' 约 3 token > 2）
    r = eng.chat(
        session_id="s1",
        prompt=Prompt(name="p", version="v1", template="abcd{{ chapter }}"),
        variables={"chapter": "很长很长的标题"},
    )
    assert seen["called"] is True
    assert r.status == "ok"


# ---------- chat_structured 主路径 ----------

def test_chat_structured_ok():
    client = FakeLLMClient([("ok", '{"a": 1}')])
    eng = make_engine(client)
    r = eng.chat_structured(
        session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
        variables={"text": "内容"},
        schema={"type": "object", "properties": {"a": {"type": "number"}}},
    )
    assert r.status == "ok" and r.data == {"a": 1}


def test_chat_structured_retry_parse():
    client = FakeLLMClient([("parse_fail", ""), ("ok", '{"a": 2}')])
    eng = make_engine(client)
    r = eng.chat_structured(
        session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
        variables={}, schema={},
    )
    assert r.status == "ok" and r.data == {"a": 2} and r.parse_attempts == 2


def test_chat_structured_parse_error_returns_raw():
    client = FakeLLMClient([("parse_fail", ""), ("parse_fail", "")])
    eng = make_engine(client)
    r = eng.chat_structured(
        session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
        variables={}, schema={}, retry_count=1,
    )
    assert r.status == "parse_error"
    assert r.raw_text == "NOT JSON"
    assert r.parse_attempts == 2


def test_chat_structured_failed_on_network():
    client = FakeLLMClient([("raise_network", "")])
    eng = make_engine(client)
    r = eng.chat_structured(
        session_id="s1", prompt=Prompt(name="p", version="v1", template="提取"),
        variables={}, schema={},
    )
    assert r.status == "failed"


def test_chat_structured_truncation_callback_triggered():
    # chat_structured 超限时同样触发 truncate_cb（与 chat 流程一致）
    seen = {}

    def truncate_cb(prompt_text, variables):
        seen["called"] = True
        return {**variables, "chapter": "裁剪版"}

    eng = InferenceEngine(
        client=FakeLLMClient([("ok", '{"a": 1}')]),
        models=Models(),
        registry=PromptRegistry(),
        estimator=approximate_token_count,
        truncate_cb=truncate_cb,
    )
    eng._max_context_tokens = 2  # 强制超限（渲染约 3 token > 2）
    r = eng.chat_structured(
        session_id="s1",
        prompt=Prompt(name="p", version="v1", template="abcd{{ chapter }}"),
        variables={"chapter": "很长很长的标题"},
        schema={"type": "object", "properties": {"a": {"type": "number"}}},
    )
    assert seen["called"] is True
    assert r.status == "ok"


# ---------- 补充分支覆盖（defensive / 缺失配置） ----------

def test_chat_fallback_all_failed_when_only_primary():
    """models 无 fallback：仅尝试主模型，失败后直接 failed。"""
    client = FakeLLMClient([("raise_network", "")])
    eng = InferenceEngine(
        client=client,
        models=ModelsNoFallback(),
        registry=PromptRegistry(),
        estimator=approximate_token_count,
    )
    r = eng.chat(session_id="s1", prompt="t", variables={})
    assert r.status == "failed"
    assert len(client.calls) == 1
    assert client.calls[0]["model"] == "deepseek-chat"


def test_chat_empty_models_uses_default():
    """主/备用均为空配置时回退内置默认模型 deepseek-chat。"""
    eng = InferenceEngine(
        client=FakeLLMClient([("ok", "x")]),
        models=ModelsEmpty(),
        registry=PromptRegistry(),
        estimator=approximate_token_count,
    )
    r = eng.chat(session_id="s1", prompt="t", variables={})
    assert r.status == "ok"
    assert eng._model_names(None) == ["deepseek-chat"]


def test_chat_defaults_without_registry_estimator():
    """registry/estimator/models 均未配置时使用内置默认（PromptRegistry + approximate + deepseek-chat）。"""
    client = FakeLLMClient([("ok", "默认")])
    eng = InferenceEngine(client=client, models=None)
    r = eng.chat(
        session_id="s1",
        prompt=Prompt(name="p", version="v1", template="tt"),
        variables={},
    )
    assert r.status == "ok" and r.text == "默认"
    assert eng._model_names(None) == ["deepseek-chat"]


def test_chat_explicit_model_param():
    """显式指定 model，跳过主/备选择，仅调用该模型。"""
    client = FakeLLMClient([("ok", "c")])
    eng = make_engine(client)
    r = eng.chat(session_id="s1", prompt="t", variables={}, model="qwen-custom")
    assert r.status == "ok"
    assert client.calls[0]["model"] == "qwen-custom"
    assert len(client.calls) == 1


def test_chat_truncation_without_callback_keeps_variables():
    """超限但无 truncate_cb，variables 原样保留（裁剪分支不触发）。"""
    client = FakeLLMClient([("ok", "ok")])
    eng = make_engine(client)
    eng._max_context_tokens = 2  # 强制超限（渲染后约 11 token）
    r = eng.chat(
        session_id="s1",
        prompt=Prompt(name="p", version="v1", template="abcd{{ chapter }}"),
        variables={"chapter": "很长很长的标题"},
    )
    assert r.status == "ok"
    assert client.calls[0]["model"] == "deepseek-chat"


def test_chat_structured_empty_schema_no_hint():
    """schema 为空时不追加 schema 提示，仍可正常解析。"""
    client = FakeLLMClient([("ok", '{"v": true}')])
    eng = make_engine(client)
    r = eng.chat_structured(
        session_id="s1", prompt="直接文本", variables={}, schema={},
    )
    assert r.status == "ok" and r.data == {"v": True}


def test_chat_truncation_callback_returns_none_keeps_variables():
    """truncate_cb 返回 None 时回退原 variables（覆盖 new_vars is None 分支）。"""
    def truncate_cb(prompt_text, variables):
        return None

    client = FakeLLMClient([("ok", "x")])
    eng = InferenceEngine(
        client=client,
        models=Models(),
        registry=PromptRegistry(),
        estimator=approximate_token_count,
        truncate_cb=truncate_cb,
    )
    eng._max_context_tokens = 2
    r = eng.chat(
        session_id="s1",
        prompt=Prompt(name="p", version="v1", template="abcd{{ chapter }}"),
        variables={"chapter": "很长很长的标题"},
    )
    assert r.status == "ok"
    # 未替换变量：发送内容仍为原样渲染
    assert "很长很长的标题" in client.calls[0]["messages"][0]