import json
from typing import Any

from backend.config import get_prompts_dir, get_data_dir, get_output_dir
from backend.evaluation import (
    evaluate_correctness,
    extract_expected_from_sample,
)
from backend.extract_schema import EXTRACT_FIELDS, normalize_extract_result
from backend.llm_client import call_llm


TASK_PROMPTS = {
    "classify": "classify.txt",
    "extract": "extract.txt",
    "match": "match.txt",
}

TASK_DATA = {
    "classify": "classify_samples.json",
    "extract": "extract_samples.json",
    "match": "match_samples.json",
}

INPUT_FIELDS = {"text", "text_a", "text_b"}


def validate_task(task: str) -> None:
    if task not in TASK_PROMPTS:
        allowed = "、".join(TASK_PROMPTS)
        raise ValueError(f"未知任务：{task}。可选任务：{allowed}")


def read_prompt(task: str) -> str:
    validate_task(task)
    return (get_prompts_dir() / TASK_PROMPTS[task]).read_text(encoding="utf-8")


def read_samples(
    task: str,
    difficulty: str | None = None,
) -> list[dict[str, Any]]:
    validate_task(task)
    data_path = get_data_dir() / TASK_DATA[task]
    samples = json.loads(data_path.read_text(encoding="utf-8"))

    if not difficulty:
        return samples

    normalized = difficulty.strip().lower()
    allowed = {"easy", "medium", "hard"}
    if normalized not in allowed:
        raise ValueError(f"未知难度：{difficulty}。可选难度：easy、medium、hard")

    return [sample for sample in samples if sample.get("difficulty") == normalized]


def build_prompt(task: str, payload: dict[str, Any]) -> str:
    template = read_prompt(task)

    if task in {"classify", "extract"}:
        text = str(payload.get("text", "")).strip()
        if not text:
            raise ValueError("text 不能为空")
        return template.replace("{{text}}", text)

    if task == "match":
        text_a = str(payload.get("text_a", "")).strip()
        text_b = str(payload.get("text_b", "")).strip()
        if not text_a or not text_b:
            raise ValueError("text_a 和 text_b 都不能为空")
        return template.replace("{{text_a}}", text_a).replace("{{text_b}}", text_b)

    raise ValueError(f"暂不支持任务：{task}")


def parse_json_output(raw_output: str) -> dict[str, Any]:
    text = raw_output.strip()
    if text.startswith("```"):
        text = text.strip("`")
        if text.lower().startswith("json"):
            text = text[4:].strip()

    try:
        return {"result": json.loads(text), "raw_output": raw_output}
    except json.JSONDecodeError as exc:
        return {
            "result": None,
            "raw_output": raw_output,
            "error": f"模型返回内容不是合法 JSON：{exc}",
        }


def build_result_record(
    task: str,
    payload: dict[str, Any],
    parsed: dict[str, Any],
    *,
    sample_id: str | None = None,
    difficulty: str | None = None,
    expected: dict[str, Any] | None = None,
    include_raw_on_success: bool = False,
) -> dict[str, Any]:
    record: dict[str, Any] = {
        "task": task,
        "input": payload,
    }

    if sample_id is not None:
        record["id"] = sample_id
    if difficulty is not None:
        record["difficulty"] = difficulty
    if expected is not None:
        record["expected"] = expected

    if parsed.get("result") is not None:
        result = parsed["result"]
        if task == "extract":
            try:
                parsed["result"] = normalize_extract_result(result)
            except ValueError as exc:
                parsed = {
                    "result": None,
                    "raw_output": parsed.get("raw_output"),
                    "error": str(exc),
                }

    if parsed.get("result") is not None:
        record["result"] = parsed["result"]
        record["status"] = "success"
        record["is_correct"] = evaluate_correctness(task, expected, parsed["result"])
        if include_raw_on_success:
            record["raw_output"] = parsed.get("raw_output")
        return record

    record["result"] = None
    record["status"] = "failed"
    record["is_correct"] = False if expected is not None else None
    record["error"] = parsed.get("error", "打标失败")
    record["raw_output"] = parsed.get("raw_output")
    return record


def label_one(
    task: str,
    payload: dict[str, Any],
    *,
    sample: dict[str, Any] | None = None,
    include_raw_on_success: bool = False,
) -> dict[str, Any]:
    prompt = build_prompt(task, payload)
    raw_output = call_llm(prompt)
    parsed = parse_json_output(raw_output)
    expected = extract_expected_from_sample(task, sample) if sample else None

    return build_result_record(
        task,
        payload,
        parsed,
        sample_id=sample.get("id") if sample else None,
        difficulty=sample.get("difficulty") if sample else None,
        expected=expected,
        include_raw_on_success=include_raw_on_success,
    )


def label_samples(
    task: str,
    difficulty: str | None = None,
    *,
    include_raw_on_success: bool = False,
) -> list[dict[str, Any]]:
    samples = read_samples(task, difficulty)
    results = []

    for sample in samples:
        payload = {
            key: value for key, value in sample.items() if key in INPUT_FIELDS
        }
        results.append(
            label_one(
                task,
                payload,
                sample=sample,
                include_raw_on_success=include_raw_on_success,
            )
        )

    return results


def save_results(task: str, results: list[dict[str, Any]]) -> str:
    output_dir = get_output_dir()
    output_dir.mkdir(exist_ok=True)
    output_path = output_dir / f"{task}_results.json"
    output_path.write_text(
        json.dumps(results, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    return str(output_path)


def get_public_response(record: dict[str, Any]) -> dict[str, Any]:
    response: dict[str, Any] = {
        "task": record.get("task"),
        "input": record.get("input"),
        "status": record.get("status"),
    }

    if record.get("id") is not None:
        response["id"] = record["id"]
    if record.get("difficulty") is not None:
        response["difficulty"] = record["difficulty"]
    if record.get("expected") is not None:
        response["expected"] = record["expected"]
    if record.get("result") is not None:
        response["result"] = record["result"]
    if record.get("is_correct") is not None:
        response["is_correct"] = record["is_correct"]
    if record.get("status") != "success":
        response["error"] = record.get("error")
        response["raw_output"] = record.get("raw_output")

    return response