"""数据处理任务的大模型生成适配器。"""

from __future__ import annotations

import hashlib
import json
import re
from collections.abc import Callable, Iterable, Mapping
from typing import Any
from urllib.parse import urlsplit, urlunsplit

import httpx

from app.modules.data_process.algorithms import normalize_text, stable_split_assignments
from app.modules.data_process.constants import (
    MAX_QA_PAIRS_PER_ITEM,
    MODEL_GENERATION_BATCH_SIZE,
)


class ModelGenerationError(ValueError):
    """模型配置、响应或调用失败。"""


OUTPUT_TYPE_STANDARD = "standard"
OUTPUT_TYPE_REASONING = "reasoning"
SUPPORTED_OUTPUT_TYPES = {OUTPUT_TYPE_STANDARD, OUTPUT_TYPE_REASONING}


def chat_completions_url(value: str) -> str:
    """把域名、基础 URL 或完整地址统一为 chat completions 地址。"""

    raw = (value or "").strip()
    if not raw:
        raise ModelGenerationError("generation model api_url is required")
    if "://" not in raw:
        raw = f"https://{raw}"
    parsed = urlsplit(raw)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        raise ModelGenerationError("generation model api_url must be an HTTP(S) host or URL")
    if parsed.username or parsed.password:
        raise ModelGenerationError("generation model api_url must not contain credentials")

    path = parsed.path.rstrip("/")
    if path.endswith("/chat/completions"):
        target_path = path
    elif path.endswith("/v1"):
        target_path = f"{path}/chat/completions"
    elif not path:
        target_path = "/v1/chat/completions"
    else:
        target_path = f"{path}/v1/chat/completions"
    return urlunsplit((parsed.scheme, parsed.netloc, target_path, "", ""))


def _message_content(payload: Mapping[str, Any]) -> str:
    try:
        content = payload["choices"][0]["message"]["content"]
    except (KeyError, IndexError, TypeError) as exc:
        raise ModelGenerationError(
            "model response does not contain choices[0].message.content"
        ) from exc
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = [
            str(item.get("text") or "")
            for item in content
            if isinstance(item, Mapping) and item.get("type") in {None, "text", "output_text"}
        ]
        if parts:
            return "".join(parts)
    raise ModelGenerationError("model response content must be text")


def _json_payload(content: str) -> Any:
    # 只移除模型在 JSON 之前自行输出的思考过程，不能破坏 JSON 字段中的训练内容。
    cleaned = re.sub(
        r"^\s*<think>[\s\S]*?</think>\s*",
        "",
        content,
        count=1,
        flags=re.IGNORECASE,
    ).strip()
    fenced = re.fullmatch(r"```(?:json)?\s*([\s\S]*?)\s*```", cleaned, flags=re.IGNORECASE)
    if fenced:
        cleaned = fenced.group(1).strip()
    try:
        return json.loads(cleaned)
    except json.JSONDecodeError as exc:
        raise ModelGenerationError(
            f"model response is not valid JSON at line {exc.lineno}, column {exc.colno}"
        ) from exc


def _result_items(payload: Any) -> list[Mapping[str, Any]]:
    if isinstance(payload, list):
        values = payload
    elif isinstance(payload, Mapping):
        nested = next(
            (
                payload[key]
                for key in ("items", "results", "data", "records")
                if isinstance(payload.get(key), list)
            ),
            None,
        )
        values = nested if isinstance(nested, list) else [payload]
    else:
        raise ModelGenerationError("model JSON must be an object or array")
    items = [item for item in values if isinstance(item, Mapping)]
    if not items:
        raise ModelGenerationError("model JSON does not contain result objects")
    return items


def _prompt_messages(
    prompt: str,
    content: str,
    count: int,
    *,
    start_index: int,
    total_count: int,
    output_type: str,
) -> list[dict[str, str]]:
    end_index = start_index + count - 1
    if output_type == OUTPUT_TYPE_REASONING:
        schema = '{"items":[{"instruction":"...","input":"...","reasoning":"...","answer":"..."}]}'
        output_rule = (
            "instruction、reasoning 和 answer 均不得为空；reasoning 必须是基于来源内容、"
            "可核对且简洁的推理步骤，answer 只写最终答案。"
            "这是思维链输出模式，即使其他提示语要求省略分析，也不得省略 reasoning。"
            "不要自行添加 <think> 标签，系统会在保存时统一组装。"
        )
    else:
        schema = '{"items":[{"instruction":"...","input":"...","output":"..."}]}'
        output_rule = "instruction 和 output 不得为空，不要输出分析过程。"
    schema_instruction = (
        f"必须只返回 JSON 对象，格式为 {schema}；items 必须包含 {count} 条。"
        f"这是总计 {total_count} 条中的第 {start_index}-{end_index} 条，"
        "各条必须使用不同的提问角度和表述，避免重复。"
        f"{output_rule}不要输出 Markdown 代码围栏或 JSON 之外的说明。"
    )
    base_prompt = normalize_text(prompt) or "请根据来源内容生成可用于监督微调的问答数据。"
    if "{{ content }}" in base_prompt:
        user_prompt = base_prompt.replace("{{ content }}", content)
        return [
            {"role": "system", "content": schema_instruction},
            {"role": "user", "content": user_prompt},
        ]
    return [
        {"role": "system", "content": f"{base_prompt}\n{schema_instruction}"},
        {"role": "user", "content": f"来源内容：\n{content}"},
    ]


def generate_model_records(
    preview_items: Iterable[Mapping[str, Any]],
    *,
    model: Mapping[str, Any],
    config: Mapping[str, Any],
    task_id: str,
    split: Mapping[str, int],
    qa_pairs_per_item: int,
    client: httpx.Client | None = None,
    on_progress: Callable[[int, int], None] | None = None,
) -> list[dict[str, Any]]:
    """调用 OpenAI 兼容接口，将预览切片生成标准训练记录。

    每个切片按安全批次调用模型；失败批次会产生一条可人工修复的
    invalid 结果，已经成功的批次不会丢失。
    """

    if not 1 <= qa_pairs_per_item <= MAX_QA_PAIRS_PER_ITEM:
        raise ModelGenerationError(f"qa_pairs_per_item must be in [1, {MAX_QA_PAIRS_PER_ITEM}]")
    output_type = str(config.get("output_type") or OUTPUT_TYPE_STANDARD).strip().lower()
    if output_type not in SUPPORTED_OUTPUT_TYPES:
        raise ModelGenerationError(f"output_type must be one of {sorted(SUPPORTED_OUTPUT_TYPES)}")
    endpoint = chat_completions_url(str(model.get("api_url") or ""))
    model_name = str(model.get("online_model_name") or model.get("name") or "").strip()
    if not model_name:
        raise ModelGenerationError("generation model name is required")

    temperature = float(config.get("temperature", 0.7))
    max_tokens = int(config.get("max_tokens", 1024))
    timeout = max(1.0, min(120.0, float(config.get("request_timeout_seconds", 60))))
    retries = max(0, min(5, int(config.get("generation_retries", 2))))
    headers = {"Content-Type": "application/json"}
    api_key = str(model.get("api_key") or "").strip()
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    owns_client = client is None
    http_client = client or httpx.Client(timeout=timeout)
    results: list[dict[str, Any]] = []
    try:
        preview_list = list(preview_items)
        total_items = len(preview_list)
        for item_index, item in enumerate(preview_list):
            preview_id = str(item.get("id") or f"preview-{item_index + 1}")
            content = normalize_text(
                str(item.get("edited_content") or item.get("original_content") or "")
            )
            for batch_offset in range(0, qa_pairs_per_item, MODEL_GENERATION_BATCH_SIZE):
                batch_count = min(
                    MODEL_GENERATION_BATCH_SIZE,
                    qa_pairs_per_item - batch_offset,
                )
                batch_start = batch_offset + 1
                batch_end = batch_offset + batch_count
                request_payload: dict[str, Any] = {
                    "model": model_name,
                    "messages": _prompt_messages(
                        str(config.get("generation_prompt") or ""),
                        content,
                        batch_count,
                        start_index=batch_start,
                        total_count=qa_pairs_per_item,
                        output_type=output_type,
                    ),
                    "temperature": temperature,
                    "max_tokens": max_tokens,
                }
                if bool(config.get("json_mode", False)):
                    request_payload["response_format"] = {"type": "json_object"}

                last_error: Exception | None = None
                generated_items: list[Mapping[str, Any]] | None = None
                for _ in range(retries + 1):
                    try:
                        response = http_client.post(
                            endpoint,
                            headers=headers,
                            json=request_payload,
                        )
                        response.raise_for_status()
                        body = response.json()
                        if not isinstance(body, Mapping):
                            raise ModelGenerationError("model response body must be a JSON object")
                        candidate_items = _result_items(_json_payload(_message_content(body)))
                        if len(candidate_items) < batch_count:
                            raise ModelGenerationError(
                                "model response contains fewer result objects than requested: "
                                f"expected {batch_count}, got {len(candidate_items)}"
                            )
                        generated_items = candidate_items
                        break
                    except (
                        httpx.HTTPError,
                        json.JSONDecodeError,
                        ModelGenerationError,
                    ) as exc:
                        last_error = exc

                if generated_items is None:
                    error_message = str(last_error or "model generation failed")[:2000]
                    failure_instruction = (
                        f"模型生成失败，请人工补充（第 {batch_start}-{batch_end} 条）"
                    )
                    result_id = (
                        "result_"
                        f"{hashlib.sha256(f'{preview_id}:error:{batch_start}'.encode()).hexdigest()[:16]}"
                    )
                    results.append(
                        {
                            "id": result_id,
                            "preview_item_id": preview_id,
                            "instruction": failure_instruction,
                            "input": content,
                            "output": "",
                            "original_instruction": failure_instruction,
                            "original_input": content,
                            "original_output": "",
                            "status": "invalid",
                            "error": error_message,
                            "split": "train",
                        }
                    )
                    continue

                for batch_index, value in enumerate(generated_items[:batch_count]):
                    variant_index = batch_offset + batch_index
                    instruction = normalize_text(
                        str(value.get("instruction") or value.get("question") or "")
                    )
                    input_text = normalize_text(
                        str(value.get("input") or value.get("context") or "")
                    )
                    if output_type == OUTPUT_TYPE_REASONING:
                        reasoning = normalize_text(
                            re.sub(
                                r"</?think>",
                                "",
                                str(value.get("reasoning") or value.get("analysis") or ""),
                                flags=re.IGNORECASE,
                            )
                        )
                        answer = normalize_text(
                            re.sub(
                                r"</?think>",
                                "",
                                str(
                                    value.get("answer")
                                    or value.get("final_answer")
                                    or value.get("output")
                                    or ""
                                ),
                                flags=re.IGNORECASE,
                            )
                        )
                        output = (
                            f"<think>\n{reasoning}\n</think>\n{answer}"
                            if reasoning and answer
                            else answer or (f"<think>\n{reasoning}\n</think>" if reasoning else "")
                        )
                        valid = bool(instruction and reasoning and answer)
                        missing_error = "model result is missing instruction, reasoning or answer"
                    else:
                        output = normalize_text(
                            str(
                                value.get("output")
                                or value.get("answer")
                                or value.get("response")
                                or ""
                            )
                        )
                        output = normalize_text(
                            re.sub(
                                r"<think>[\s\S]*?(?:</think>|$)",
                                "",
                                output,
                                flags=re.IGNORECASE,
                            )
                        )
                        valid = bool(instruction and output)
                        missing_error = "model result is missing instruction or output"
                    raw_id = f"{preview_id}:{variant_index + 1}:{instruction}:{output}"
                    result_id = f"result_{hashlib.sha256(raw_id.encode()).hexdigest()[:16]}"
                    results.append(
                        {
                            "id": result_id,
                            "preview_item_id": preview_id,
                            "instruction": instruction,
                            "input": input_text,
                            "output": output,
                            "original_instruction": instruction,
                            "original_input": input_text,
                            "original_output": output,
                            "status": "valid" if valid else "invalid",
                            "error": (None if valid else missing_error),
                            "split": "train",
                        }
                    )
            if on_progress:
                on_progress(item_index + 1, total_items)
    finally:
        if owns_client:
            http_client.close()
    assignments = stable_split_assignments(
        [str(result["id"]) for result in results],
        split,
        seed=task_id,
    )
    for result, assignment in zip(results, assignments, strict=True):
        result["split"] = assignment
    return results


__all__ = ["ModelGenerationError", "chat_completions_url", "generate_model_records"]