"""真实 fallback 换模型重试单测。

只用注入的假 runner（假 CLI），不跑真实模型回合、不联网、
不登录任何服务；env 映射也是注入的普通 dict，里面的值只是
占位用的假名，绝不是真实密钥。

假 runner 的约定：host.py 风格的 JSON 成功时 returncode=0，
失败且可 fallback 时同样 returncode=0 但 JSON 的 status 字段
是 "404"/"500"/"timeout"（_result_from_runner 据此把结果标
failed 并把 host_status 置成该字段，这正是 worker 判定
fallback 的依据）；returncode 非 0 时 host_status 为空。
"""

import json
import os
import sqlite3
import sys
import tempfile
import time
import unittest
from pathlib import Path
from types import SimpleNamespace

try:
    from cloud.cost import CostStore
    from cloud.models import ModelRouter
    from cloud.server import init_db
    from cloud.worker import Worker
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.cost import CostStore
    from cloud.models import ModelRouter
    from cloud.server import init_db
    from cloud.worker import Worker


def _insert_task(db_path, task_id, prompt="do work", status="queued",
                 project_id="p1"):
    now = time.time()
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            "INSERT INTO tasks(id, project_id, prompt, status,"
            " idempotency_key, depends_on, priority, created_by,"
            " created_at, updated_at, last_heartbeat)"
            " VALUES(?, ?, ?, ?, NULL, '[]', 0, 'tester', ?, ?, NULL)",
            (task_id, project_id, prompt, status, now, now),
        )
        conn.commit()
    finally:
        conn.close()


def _status(db_path, task_id):
    conn = sqlite3.connect(db_path)
    try:
        row = conn.execute(
            "SELECT status FROM tasks WHERE id = ?",
            (task_id,)).fetchone()
    finally:
        conn.close()
    return row[0] if row else None


def _execution(db_path, task_id):
    conn = sqlite3.connect(db_path)
    try:
        row = conn.execute(
            "SELECT error_text, actual_model FROM tasks WHERE id = ?",
            (task_id,)).fetchone()
        cost = conn.execute(
            "SELECT model, fallback_count FROM cost_records"
            " WHERE task_id = ?",
            (task_id,)).fetchone()
        # 测试库里只有一个项目/任务，审计按项目全取：
        # worker_fallback 的 detail 是“模型 -> 模型”，不含任务 id，
        # 按 detail LIKE 任务 id 查会漏掉它。
        audits = [r[0] for r in conn.execute(
            "SELECT action FROM audit_log WHERE project_id = ?",
            ("p1",)).fetchall()]
    finally:
        conn.close()
    return row, cost, audits


def _host_json(name, status, text):
    return SimpleNamespace(
        returncode=0,
        stdout=json.dumps({"task": name, "status": status, "text": text}),
        stderr="")


def _model_of(cmd):
    """假 CLI 命令里的 --model 值（没有时返回 None，即配置模型）。"""
    if "--model" in cmd:
        return cmd[cmd.index("--model") + 1]
    return None


class FallbackRetryTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        init_db(self.db_path)
        self.env = {"XCM_API_KEY": "test-placeholder"}
        self.calls = []

    def tearDown(self):
        self.tmp.cleanup()

    def _worker(self, runner=None, **kwargs):
        kwargs.setdefault("env", self.env)
        kwargs.setdefault("retry_backoff_seconds", 0)
        return Worker(self.db_path, "/fake/host.py",
                      runner=runner, **kwargs)

    def _recording(self, inner):
        def fake(cmd):
            self.calls.append(list(cmd))
            return inner(cmd)
        return fake

    def _models_seen(self):
        """每次调用里实际传给 host.py 的模型（None 表示配置模型）。"""
        seen = []
        for cmd in self.calls:
            if "dispatch" in cmd or "run" in cmd:
                seen.append(_model_of(cmd))
        return seen

    def test_primary_404_fallback_succeeds(self):
        # 主模型 404，配了 fallback 列表：真的换模型跑第二次且成功。
        def fake(cmd):
            if _model_of(cmd) == "gpt-6-sol":
                return _host_json("A", "completed", "done by fallback")
            return _host_json("A", "404", "primary 404 boom")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        self.assertEqual(result["actual_model"], "gpt-6-sol")
        self.assertEqual(len(self.calls), 2)
        self.assertEqual(self._models_seen(),
                         ["gpt-6.1-sol", "gpt-6-sol"])
        self.assertEqual(len(result["events"]), 1)
        event = result["events"][0]
        self.assertEqual(event["type"], "model_fallback")
        self.assertEqual(event["from_model"], "gpt-6.1-sol")
        self.assertEqual(event["to_model"], "gpt-6-sol")
        self.assertEqual(_status(self.db_path, "A"), "completed")
        row, cost, audits = _execution(self.db_path, "A")
        self.assertEqual(row[1], "gpt-6-sol")
        self.assertEqual(cost[0], "gpt-6-sol")
        self.assertEqual(cost[1], 1)
        self.assertIn("worker_fallback", audits)

    def test_primary_401_never_fallbacks(self):
        # 401 绝不 fallback：只调一次，直接 failed，无 fallback 事件。
        def fake(cmd):
            return _host_json("A", "401", "bad credentials")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        self.assertEqual(len(self.calls), 1)
        self.assertEqual(result["events"], [])
        row, cost, audits = _execution(self.db_path, "A")
        self.assertEqual(cost[1], 0)
        self.assertNotIn("worker_fallback", audits)

    def test_no_fallback_list_keeps_old_event_only_behavior(self):
        # 未配 fallback 列表：老行为——只记事件、直接 failed、调一次。
        def fake(cmd):
            return _host_json("A", "404", "primary 404 boom")

        _insert_task(self.db_path, "A")
        router = ModelRouter(["gpt-6.1-sol", "gpt-6-sol"])
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            router=router,
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        self.assertEqual(len(self.calls), 1)
        self.assertEqual(len(result["events"]), 1)
        self.assertEqual(result["events"][0]["type"], "model_fallback")
        self.assertEqual(result["actual_model"], "gpt-6.1-sol")
        row, cost, audits = _execution(self.db_path, "A")
        self.assertEqual(cost[1], 1)
        self.assertIn("worker_fallback", audits)

    def test_both_attempts_fail_error_contains_both_reasons(self):
        # 第二次也失败：error_text 同时含两次原因，fallback_count=1。
        def fake(cmd):
            if _model_of(cmd) == "gpt-6-sol":
                return _host_json("A", "500", "fallback 500 boom")
            return _host_json("A", "404", "primary 404 boom")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        self.assertIn("primary 404 boom", result["error"])
        self.assertIn("fallback 500 boom", result["error"])
        self.assertIn("gpt-6.1-sol", result["error"])
        self.assertIn("gpt-6-sol", result["error"])
        self.assertEqual(len(self.calls), 2)
        row, cost, audits = _execution(self.db_path, "A")
        self.assertEqual(cost[1], 1)
        self.assertEqual(row[1], "gpt-6-sol")

    def test_fallback_at_most_once_even_with_two_models(self):
        # 配两个 fallback 模型也只试第一个：总共调两次，绝不试第三个。
        def fake(cmd):
            return _host_json("A", "429", "always rate limited")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol", "gpt-5.6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        self.assertEqual(len(self.calls), 2)
        self.assertEqual(self._models_seen(),
                         ["gpt-6.1-sol", "gpt-6-sol"])
        self.assertNotIn("gpt-5.6-sol", self._models_seen())

    def test_timeout_first_attempt_fallback_succeeds(self):
        # 超时也算可 fallback：第一次 timeout，换模型后成功。
        def fake(cmd):
            if _model_of(cmd) == "gpt-6-sol":
                return _host_json("A", "completed", "done after timeout")
            return _host_json("A", "timeout", "")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        self.assertEqual(result["actual_model"], "gpt-6-sol")
        self.assertEqual(len(self.calls), 2)
        self.assertEqual(result["events"][0]["reason"], "timeout")

    def test_thread_start_crash_with_404_in_stderr_fallback_succeeds(self):
        # 坏模型最常见的真实失败：thread/start 抛错、host.py 非 0
        # 退出、不打印 JSON，host_status 为 None；stderr 里有显式
        # 404 时必须 fallback，第二次换模型成功（主机实测回归）。
        def fake(cmd):
            if _model_of(cmd) == "gpt-6-sol":
                return _host_json("A", "completed", "done by fallback")
            return SimpleNamespace(
                returncode=1, stdout="",
                stderr="RuntimeError: thread/start 失败: {'error': "
                       "{'code': 404, 'message': 'model not found'}}")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        self.assertEqual(result["actual_model"], "gpt-6-sol")
        self.assertEqual(len(self.calls), 2)
        self.assertEqual(result["events"][0]["reason"], "404")
        row, cost, audits = _execution(self.db_path, "A")
        self.assertEqual(cost[1], 1)
        self.assertIn("worker_fallback", audits)

    def test_thread_start_crash_with_401_in_stderr_never_fallbacks(self):
        # stderr 里是 401 时绝不 fallback，只调一次。
        def fake(cmd):
            return SimpleNamespace(
                returncode=1, stdout="",
                stderr="RuntimeError: thread/start 失败: {'error': "
                       "{'code': 401, 'message': 'unauthorized'}}")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        self.assertEqual(len(self.calls), 1)
        self.assertEqual(result["events"], [])

    def test_thread_start_crash_without_code_never_fallbacks(self):
        # 没有显式状态码的通用失败不猜测：不 fallback，只调一次。
        def fake(cmd):
            return SimpleNamespace(
                returncode=1, stdout="", stderr="host execution failed")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        self.assertEqual(len(self.calls), 1)
        self.assertEqual(result["events"], [])

    def test_crash_without_code_error_text_keeps_stderr(self):
        # stdout 为空、stderr 有真实报错时，最终 error_text 必须
        # 留住 stderr，不能被通用的 host execution failed 盖掉
        # （主机实测：缺 repo 的 1 秒失败曾因此完全无法排障）。
        def fake(cmd):
            return SimpleNamespace(
                returncode=1, stdout="",
                stderr="Traceback (most recent call last):\n"
                       "RuntimeError: 缺少 --repo，也没有在 host-config.json 里配 repo")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        row, cost, audits = _execution(self.db_path, "A")
        self.assertIn("缺少 --repo", row[0])
        self.assertEqual(result["events"], [])

    def test_failed_turn_json_with_404_in_text_fallback_succeeds(self):
        # 真实回合失败形态：host.py JSON 的 status 只是 "failed"，
        # 真正的 404 藏在 text 里；必须从 text 里提取出来 fallback。
        def fake(cmd):
            if _model_of(cmd) == "gpt-6-sol":
                return _host_json("A", "completed", "done by fallback")
            return _host_json(
                "A", "failed",
                "turn failed: 404 model not found: gpt-9.9-nonexistent")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        self.assertEqual(result["actual_model"], "gpt-6-sol")
        self.assertEqual(len(self.calls), 2)
        self.assertEqual(result["events"][0]["reason"], "404")

    def test_failed_turn_json_with_401_in_text_never_fallbacks(self):
        # 回合 JSON 里是 401 时绝不 fallback，只调一次。
        def fake(cmd):
            return _host_json("A", "failed", "turn failed: 401 unauthorized")

        _insert_task(self.db_path, "A")
        worker = self._worker(
            runner=self._recording(fake),
            model="gpt-6.1-sol",
            fallback_models=["gpt-6-sol"],
            max_retries=0)
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        self.assertEqual(len(self.calls), 1)
        self.assertEqual(result["events"], [])

    def test_load_tasks_only_picks_relevant_statuses(self):
        # load_tasks 只取 queued/completed/merged：running/failed
        # 的任务对 pick_next 既不是就绪也不是已完成，排除它们
        # 不影响结果，但省掉每轮 poll 全表 + 重字段的开销。
        for tid, status in [("Q", "queued"), ("C", "completed"),
                            ("R", "running"), ("F", "failed")]:
            _insert_task(self.db_path, tid)
            import sqlite3
            conn = sqlite3.connect(self.db_path)
            conn.execute(
                "UPDATE tasks SET status = ? WHERE id = ?", (status, tid))
            conn.commit()
            conn.close()
        worker = self._worker(runner=lambda cmd: None)
        got = {t["id"] for t in worker.load_tasks()}
        self.assertEqual(got, {"Q", "C"})


if __name__ == "__main__":
    unittest.main()