from __future__ import annotations

import json

import httpx
import pytest

from app.modules.data_process.generation import (
    ModelGenerationError,
    chat_completions_url,
    generate_model_records,
)


def test_chat_completions_url_accepts_host_base_and_complete_url() -> None:
    assert chat_completions_url("www.caoxiaozhu.com") == (
        "https://www.caoxiaozhu.com/v1/chat/completions"
    )
    assert chat_completions_url("https://model.example/v1") == (
        "https://model.example/v1/chat/completions"
    )
    complete = "https://model.example/openai/v1/chat/completions"
    assert chat_completions_url(complete) == complete


def test_generate_model_records_uses_prompt_auth_and_stable_split() -> None:
    requests: list[httpx.Request] = []
    progress_updates: list[tuple[int, int]] = []

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(request)
        payload = json.loads(request.content)
        assert payload["model"] == "qwen-plus"
        assert payload["response_format"] == {"type": "json_object"}
        assert "客户反馈页面加载慢" in payload["messages"][1]["content"]
        return httpx.Response(
            200,
            json={
                "choices": [
                    {
                        "message": {
                            "content": json.dumps(
                                {
                                    "items": [
                                        {
                                            "instruction": "请生成简洁客服回复",
                                            "input": "客户反馈页面加载慢",
                                            "output": "已收到反馈，我们正在排查。",
                                        }
                                    ]
                                },
                                ensure_ascii=False,
                            )
                        }
                    }
                ]
            },
        )

    client = httpx.Client(transport=httpx.MockTransport(handler))
    records = generate_model_records(
        [{"id": "preview-1", "edited_content": "客户反馈页面加载慢"}],
        model={
            "name": "Qwen",
            "online_model_name": "qwen-plus",
            "api_url": "model.example",
            "api_key": "test-secret",
        },
        config={
            "generation_prompt": "请处理：{{ content }}",
            "json_mode": True,
            "temperature": 0.2,
            "max_tokens": 512,
        },
        task_id="task-1",
        split={"train": 100, "validation": 0, "test": 0},
        qa_pairs_per_item=1,
        client=client,
        on_progress=lambda processed, total: progress_updates.append((processed, total)),
    )

    assert len(records) == 1
    assert records[0]["status"] == "valid"
    assert records[0]["split"] == "train"
    assert requests[0].headers["Authorization"] == "Bearer test-secret"
    assert progress_updates == [(1, 1)]


def test_generate_model_records_builds_reasoning_output_with_think_tags() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        payload = json.loads(request.content)
        system_prompt = payload["messages"][0]["content"]
        assert '"reasoning":"...","answer":"..."' in system_prompt
        assert "系统会在保存时统一组装" in system_prompt
        content = "<think>模型接口自己的分析</think>" + json.dumps(
            {
                "items": [
                    {
                        "instruction": "计算两项费用合计",
                        "input": "交通费 30 元，餐费 20 元",
                        "reasoning": "先识别两项费用，再计算 30 + 20。",
                        "answer": "合计 50 元。",
                    }
                ]
            },
            ensure_ascii=False,
        )
        return httpx.Response(
            200,
            json={"choices": [{"message": {"content": content}}]},
        )

    records = generate_model_records(
        [{"id": "preview-reasoning", "edited_content": "交通费 30 元，餐费 20 元"}],
        model={"name": "model", "api_url": "https://model.example/v1"},
        config={"output_type": "reasoning"},
        task_id="task-reasoning",
        split={"train": 100, "validation": 0, "test": 0},
        qa_pairs_per_item=1,
        client=httpx.Client(transport=httpx.MockTransport(handler)),
    )

    assert records[0]["status"] == "valid"
    assert records[0]["output"] == (
        "<think>\n先识别两项费用,再计算 30 + 20。\n</think>\n合计 50 元。"
    )


def test_generate_model_records_marks_reasoning_without_reasoning_field_invalid() -> None:
    response = {
        "choices": [
            {
                "message": {
                    "content": json.dumps(
                        {
                            "items": [
                                {
                                    "instruction": "问题",
                                    "answer": "只有最终答案",
                                }
                            ]
                        },
                        ensure_ascii=False,
                    )
                }
            }
        ]
    }
    client = httpx.Client(
        transport=httpx.MockTransport(lambda _: httpx.Response(200, json=response))
    )

    records = generate_model_records(
        [{"id": "preview-missing-reasoning", "edited_content": "来源正文"}],
        model={"name": "model", "api_url": "https://model.example/v1"},
        config={"output_type": "reasoning"},
        task_id="task-missing-reasoning",
        split={"train": 100, "validation": 0, "test": 0},
        qa_pairs_per_item=1,
        client=client,
    )

    assert records[0]["status"] == "invalid"
    assert records[0]["output"] == "只有最终答案"
    assert "reasoning" in records[0]["error"]


def test_standard_output_removes_model_think_block() -> None:
    content = json.dumps(
        {
            "items": [
                {
                    "instruction": "问题",
                    "output": "<think>不应保存的分析</think>最终答案",
                }
            ]
        },
        ensure_ascii=False,
    )
    client = httpx.Client(
        transport=httpx.MockTransport(
            lambda _: httpx.Response(
                200,
                json={"choices": [{"message": {"content": content}}]},
            )
        )
    )

    records = generate_model_records(
        [{"id": "preview-standard", "edited_content": "来源正文"}],
        model={"name": "model", "api_url": "https://model.example/v1"},
        config={"output_type": "standard"},
        task_id="task-standard",
        split={"train": 100, "validation": 0, "test": 0},
        qa_pairs_per_item=1,
        client=client,
    )

    assert records[0]["output"] == "最终答案"


def test_generate_model_records_keeps_partial_failure_for_manual_repair() -> None:
    client = httpx.Client(
        transport=httpx.MockTransport(
            lambda _: httpx.Response(200, json={"choices": [{"message": {"content": "not-json"}}]})
        )
    )
    records = generate_model_records(
        [{"id": "preview-1", "edited_content": "来源正文"}],
        model={"name": "model", "api_url": "https://model.example/v1"},
        config={"generation_retries": 1},
        task_id="task-1",
        split={"train": 80, "validation": 10, "test": 10},
        qa_pairs_per_item=1,
        client=client,
    )

    assert len(records) == 1
    assert records[0]["status"] == "invalid"
    assert records[0]["error"]


def test_generate_model_records_batches_fifty_results_with_unique_ids() -> None:
    requests: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(request)
        batch_start = (len(requests) - 1) * 10 + 1
        batch_end = batch_start + 9
        payload = json.loads(request.content)
        system_prompt = payload["messages"][0]["content"]
        assert "items 必须包含 10 条" in system_prompt
        assert f"第 {batch_start}-{batch_end} 条" in system_prompt
        return httpx.Response(
            200,
            json={
                "choices": [
                    {
                        "message": {
                            "content": json.dumps(
                                {
                                    "items": [
                                        {
                                            "instruction": "同一问题",
                                            "input": "来源正文",
                                            "output": "同一答案",
                                        }
                                        for _ in range(batch_start, batch_end + 1)
                                    ]
                                },
                                ensure_ascii=False,
                            )
                        }
                    }
                ]
            },
        )

    records = generate_model_records(
        [{"id": "preview-50", "edited_content": "来源正文"}],
        model={"name": "model", "api_url": "https://model.example/v1"},
        config={},
        task_id="task-50",
        split={"train": 100, "validation": 0, "test": 0},
        qa_pairs_per_item=50,
        client=httpx.Client(transport=httpx.MockTransport(handler)),
    )

    assert len(requests) == 5
    assert len(records) == 50
    assert len({record["id"] for record in records}) == 50
    assert {record["instruction"] for record in records} == {"同一问题"}
    assert all(record["status"] == "valid" for record in records)


def test_generate_model_records_preserves_successful_batches_when_one_fails() -> None:
    request_count = 0

    def handler(_: httpx.Request) -> httpx.Response:
        nonlocal request_count
        request_count += 1
        if request_count == 2:
            return httpx.Response(500)
        return httpx.Response(
            200,
            json={
                "choices": [
                    {
                        "message": {
                            "content": json.dumps(
                                {
                                    "items": [
                                        {
                                            "instruction": f"问题 {index}",
                                            "output": f"答案 {index}",
                                        }
                                        for index in range(1, 11)
                                    ]
                                },
                                ensure_ascii=False,
                            )
                        }
                    }
                ]
            },
        )

    records = generate_model_records(
        [{"id": "preview-partial", "edited_content": "来源正文"}],
        model={"name": "model", "api_url": "https://model.example/v1"},
        config={"generation_retries": 0},
        task_id="task-partial",
        split={"train": 100, "validation": 0, "test": 0},
        qa_pairs_per_item=20,
        client=httpx.Client(transport=httpx.MockTransport(handler)),
    )

    assert len(records) == 11
    assert sum(record["status"] == "valid" for record in records) == 10
    failed = next(record for record in records if record["status"] == "invalid")
    assert "第 11-20 条" in failed["instruction"]
    assert len({record["id"] for record in records}) == len(records)


def test_generate_model_records_retries_short_batch_then_marks_it_invalid() -> None:
    request_count = 0

    def handler(_: httpx.Request) -> httpx.Response:
        nonlocal request_count
        request_count += 1
        return httpx.Response(
            200,
            json={
                "choices": [
                    {
                        "message": {
                            "content": json.dumps(
                                {
                                    "items": [
                                        {
                                            "instruction": "只有一条",
                                            "output": "不足本批要求数量",
                                        }
                                    ]
                                },
                                ensure_ascii=False,
                            )
                        }
                    }
                ]
            },
        )

    records = generate_model_records(
        [{"id": "preview-short", "edited_content": "来源正文"}],
        model={"name": "model", "api_url": "https://model.example/v1"},
        config={"generation_retries": 1},
        task_id="task-short",
        split={"train": 100, "validation": 0, "test": 0},
        qa_pairs_per_item=10,
        client=httpx.Client(transport=httpx.MockTransport(handler)),
    )

    assert request_count == 2
    assert len(records) == 1
    assert records[0]["status"] == "invalid"
    assert "expected 10, got 1" in records[0]["error"]


@pytest.mark.parametrize("qa_pairs_per_item", [0, 51])
def test_generate_model_records_rejects_out_of_range_count(
    qa_pairs_per_item: int,
) -> None:
    with pytest.raises(ModelGenerationError, match=r"\[1, 50\]"):
        generate_model_records(
            [],
            model={"name": "model", "api_url": "https://model.example/v1"},
            config={},
            task_id="task-invalid",
            split={"train": 100, "validation": 0, "test": 0},
            qa_pairs_per_item=qa_pairs_per_item,
        )


def test_generate_model_records_rejects_unknown_output_type() -> None:
    with pytest.raises(ModelGenerationError, match="output_type"):
        generate_model_records(
            [],
            model={"name": "model", "api_url": "https://model.example/v1"},
            config={"output_type": "unknown"},
            task_id="task-invalid-output",
            split={"train": 100, "validation": 0, "test": 0},
            qa_pairs_per_item=1,
        )