"""
Evaluation runner — executes model evaluation as a subprocess job.

Usage:
    python -m compute.engines.llama_factory.eval_runner --config <config_json_path>

The config JSON is written by the compute API before spawning this subprocess.
Results are written to ``output_dir/eval_results.json`` and progress is printed
to stdout (captured as job logs).
"""

from __future__ import annotations

import json
import math
import re
import sys
import time
from difflib import SequenceMatcher
from pathlib import Path
from typing import Any


def _load_dataset(path: str) -> list[dict[str, Any]]:
    """Load a JSON or JSONL dataset file (jsonl-compatible).

    Content-sniffs instead of trusting the extension so jsonl files with a BOM,
    a single JSON array on one line, or mislabeled extensions all load correctly.

    Supports common field names used across the platform:
    * ``instruction`` + ``input`` + ``output``  (Alpaca-style)
    * ``question`` + ``answer``
    * ``messages``  (ShareGPT-style – the last assistant message is treated as reference)
    """
    file_path = Path(path)
    text = file_path.read_text(encoding="utf-8-sig", errors="replace").strip()
    if not text:
        return []
    try:
        value = json.loads(text)
    except json.JSONDecodeError:
        value = None
    if isinstance(value, list):
        return [item for item in value if isinstance(item, dict)]
    if isinstance(value, dict):
        return [value]

    samples: list[dict[str, Any]] = []
    for line in text.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(obj, dict):
            samples.append(obj)
    return samples


def _sample_question(sample: dict[str, Any]) -> str:
    """Extract the user-facing question / instruction from a sample."""
    if sample.get("instruction"):
        text = sample["instruction"]
        if sample.get("input"):
            text += "\n" + sample["input"]
        return text
    if sample.get("question"):
        return sample["question"]
    # ShareGPT-style: use the last user message as question
    messages = sample.get("messages") or []
    user_msgs = [m["content"] for m in messages if m.get("role") == "user"]
    return user_msgs[-1] if user_msgs else ""


def _sample_reference(sample: dict[str, Any]) -> str:
    """Extract the reference answer from a sample."""
    if sample.get("output"):
        return sample["output"]
    if sample.get("answer"):
        return sample["answer"]
    messages = sample.get("messages") or []
    assistant_msgs = [m["content"] for m in messages if m.get("role") == "assistant"]
    return assistant_msgs[-1] if assistant_msgs else ""


# ---------------------------------------------------------------------------
# Basic metrics
# ---------------------------------------------------------------------------

def _compute_bleu(references: list[str], predictions: list[str], ngram: int = 4) -> dict[str, Any]:
    """Compute BLEU score via sacrebleu (corpus-level)."""
    try:
        from sacrebleu.metrics import BLEU
    except ImportError:
        return {"enabled": False, "error": "sacrebleu not installed", "score": 0}
    bleu = BLEU(max_ngram_order=ngram)
    # sacrebleu expects list-of-strings; we have one reference per prediction
    score = bleu.corpus_score(predictions, [references])
    return {
        "enabled": True,
        "score": round(score.score, 2),
        "bleu": round(score.score, 2),
    }


def _compute_rouge(references: list[str], predictions: list[str], methods: list[str] | None = None) -> dict[str, Any]:
    """Compute ROUGE scores via rouge-score."""
    try:
        from rouge_score import rouge_scorer
    except ImportError:
        return {"enabled": False, "error": "rouge-score not installed", "score": 0}
    methods = methods or ["rouge1", "rouge2", "rougeL"]
    # Normalize: map "rouge_1"/"rouge1" → "rouge1", "rouge_l"/"rougeL" → "rougeL"
    _rouge_aliases = {"rouge_1": "rouge1", "rouge_2": "rouge2", "rouge_l": "rougeL"}
    methods = [_rouge_aliases.get(m, m.replace("_", "")) for m in methods]
    scorer = rouge_scorer.RougeScorer(methods, use_stemmer=True)
    totals: dict[str, float] = {}
    n = max(len(predictions), 1)
    for ref, pred in zip(references, predictions):
        result = scorer.score(ref, pred)
        for key in methods:
            totals[key] = totals.get(key, 0) + result[key].fmeasure
    avg = {k: round(v / n, 4) for k, v in totals.items()}
    return {"enabled": True, "score": round(avg.get("rougeL", avg.get("rouge1", 0)) * 100, 2), **avg}


def _compute_cosine(references: list[str], predictions: list[str]) -> dict[str, Any]:
    """Compute average cosine similarity via sklearn."""
    try:
        from sklearn.feature_extraction.text import TfidfVectorizer
        from sklearn.metrics.pairwise import cosine_similarity
    except ImportError:
        return {"enabled": False, "error": "scikit-learn not installed", "score": 0}
    try:
        vectorizer = TfidfVectorizer()
        tfidf = vectorizer.fit_transform(references + predictions)
        n = len(references)
        ref_vec = tfidf[:n]
        pred_vec = tfidf[n:]
        sims = cosine_similarity(ref_vec, pred_vec).diagonal()
        return {"enabled": True, "score": round(float(sims.mean()) * 100, 2)}
    except ValueError:
        return {"enabled": True, "score": 0, "error": "insufficient text for vectorization"}


def _normalize_text(value: str) -> str:
    return re.sub(r"\s+", " ", str(value or "").strip().lower())


def _compute_exact_match(references: list[str], predictions: list[str]) -> dict[str, Any]:
    total = len(predictions)
    if not total:
        return {"enabled": True, "score": 0, "matched": 0, "total": 0}
    matched = sum(
        1
        for ref, pred in zip(references, predictions)
        if _normalize_text(ref) == _normalize_text(pred)
    )
    return {"enabled": True, "score": round(matched / total * 100, 2), "matched": matched, "total": total}


def _compute_text_similarity(references: list[str], predictions: list[str]) -> dict[str, Any]:
    if not predictions:
        return {"enabled": True, "score": 0}
    scores = [
        SequenceMatcher(None, _normalize_text(ref), _normalize_text(pred)).ratio()
        for ref, pred in zip(references, predictions)
    ]
    return {"enabled": True, "score": round(sum(scores) / max(len(scores), 1) * 100, 2)}


# ---------------------------------------------------------------------------
# LLM Judge
# ---------------------------------------------------------------------------

def _judge_sample(
    question: str,
    reference: str,
    prediction: str,
    config: dict[str, Any],
) -> dict[str, Any]:
    """Call an OpenAI-compatible LLM to judge a single sample.

    Returns a dict with keys:
        score, max_score, passed, judgement, evaluation_reason, error_type
    """
    api_url = (config.get("api_url") or "").strip().rstrip("/")
    api_key = (config.get("api_key") or "").strip()
    eval_model = (config.get("eval_model") or "").strip()
    # 优先使用模型记录里配置的真实 API 模型名（如 deepseek-chat），
    # 否则回退到平台内部模型名
    api_model = (config.get("api_model") or "").strip() or eval_model
    eval_prompt = (config.get("eval_prompt") or "").strip()
    score_min = float(config.get("score_min", 0))
    score_max = float(config.get("score_max", 5))
    pass_threshold = float(config.get("pass_threshold", 3))

    if not api_url or not eval_model:
        return {"score": 0, "max_score": score_max, "passed": False, "judgement": "未配置",
                "evaluation_reason": "未配置评测模型", "error_type": "其他"}

    system_msg = (
        eval_prompt
        or "你是一个专业的评测专家。请根据参考答-案对被测模型的输出进行评分。"
    )
    user_msg = (
        f"## 问题\n{question}\n\n"
        f"## 参考答案\n{reference}\n\n"
        f"## 模型输出\n{prediction}\n\n"
        f"请给出 {score_min}-{score_max} 分的评分，并说明理由。"
    )

    try:
        import urllib.request
        import urllib.error

        body = json.dumps({
            "model": api_model,
            "messages": [
                {"role": "system", "content": system_msg},
                {"role": "user", "content": user_msg},
            ],
            "temperature": 0.3,
            "max_tokens": 512,
        }).encode("utf-8")

        req = urllib.request.Request(
            f"{api_url}/v1/chat/completions",
            data=body,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}",
            },
        )
        resp = urllib.request.urlopen(req, timeout=120)
        data = json.loads(resp.read().decode("utf-8"))
        reply = data["choices"][0]["message"]["content"]
    except Exception as exc:
        return {"score": 0, "max_score": score_max, "passed": False,
                "judgement": "错误", "evaluation_reason": f"评测模型调用失败: {exc}",
                "error_type": "其他"}

    # Parse score from reply — look for patterns like "4分" or "Score: 4"
    score = 0
    import re
    score_patterns = [
        r'(?:得分|分数|评分|score)[^\d]*(\d+(?:\.\d+)?)',
        r'(\d+(?:\.\d+)?)\s*分',
        r'(\d+(?:\.\d+)?)\s*/\s*\d+',
    ]
    for pat in score_patterns:
        m = re.search(pat, reply, re.IGNORECASE)
        if m:
            try:
                score = float(m.group(1))
            except ValueError:
                continue
            break
    score = max(score_min, min(score_max, score))
    passed = score >= pass_threshold

    # Determine judgement label
    if score >= pass_threshold + 1:
        judgement = "正确"
    elif score >= pass_threshold:
        judgement = "部分正确"
    else:
        judgement = "错误"

    # Guess error type from reply
    reply_lower = reply.lower()
    if any(w in reply_lower for w in ["幻觉", "hallucination", "编造"]):
        error_type = "幻觉"
    elif any(w in reply_lower for w in ["不完整", "incomplete", "遗漏"]):
        error_type = "不完整"
    elif any(w in reply_lower for w in ["格式", "format"]):
        error_type = "格式偏差"
    elif any(w in reply_lower for w in ["混淆", "confusion", "错误"]):
        error_type = "混淆"
    else:
        error_type = "其他"

    return {
        "score": score,
        "max_score": score_max,
        "passed": passed,
        "judgement": judgement,
        "evaluation_reason": reply[:2000],
        "error_type": error_type,
    }


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

def run_eval(config: dict[str, Any]) -> dict[str, Any]:
    """Execute a full evaluation run.  Returns the result dict (also written to file)."""
    model_path = config["model_name_or_path"]
    adapter_path = config.get("adapter_name_or_path", "")
    template = config.get("template", "qwen")
    dataset_path = config["dataset_path"]
    output_dir = Path(config["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    basic_cfg = config.get("basic_metrics", {})
    dimension_cfg = config.get("dimension", {}) or {}
    output_precision = int(basic_cfg.get("output_precision", 2))

    # ---- 1. Load dataset ----
    print(f"[eval] loading dataset: {dataset_path}")
    raw_samples = _load_dataset(dataset_path)
    print(f"[eval] loaded {len(raw_samples)} samples")

    # ---- 2. Load model ----
    print(f"[eval] loading model: {model_path}")
    from compute.engines.llama_factory.inference import InferenceSession
    session = InferenceSession()
    session.load(
        model_name_or_path=model_path,
        adapter_name_or_path=adapter_path,
        template=template,
        infer_backend=config.get("infer_backend", "huggingface"),
        infer_dtype=config.get("infer_dtype", "auto"),
    )
    # load() 为异步加载（立即返回 loading），必须等待后台线程完成后再进行推理
    load_result = session.wait_until_loaded(timeout=float(config.get("load_timeout", 1800)))
    if not load_result.get("loaded"):
        raise RuntimeError(f"model load failed: {load_result.get('error', 'unknown')}")
    print(f"[eval] model loaded OK")

    # ---- 3. Run inference on each sample ----
    samples: list[dict[str, Any]] = []
    predictions: list[str] = []
    references: list[str] = []
    questions: list[str] = []

    total = len(raw_samples)
    judge_enabled = bool(dimension_cfg.get("eval_model") and dimension_cfg.get("api_url"))
    print(f"[eval] starting inference on {total} samples, judge={'enabled' if judge_enabled else 'disabled'}")

    for idx, raw in enumerate(raw_samples, start=1):
        question = _sample_question(raw)
        reference = _sample_reference(raw)
        if not question:
            print(f"[eval] sample {idx}/{total}: skipped (no question)")
            continue

        # Inference
        chat_msgs = [{"role": "user", "content": question}]
        result = session.chat(
            chat_msgs,
            temperature=float(config.get("temperature", 0.1)),
            top_p=float(config.get("top_p", 0.95)),
            max_new_tokens=int(config.get("max_new_tokens", 512)),
            do_sample=False,
        )
        prediction = result.get("response", "") if not result.get("error") else f"[ERROR] {result['error']}"

        predictions.append(prediction)
        references.append(reference)
        questions.append(question)

        # LLM Judge
        judge_result: dict[str, Any] = {}
        if judge_enabled:
            judge_result = _judge_sample(question, reference, prediction, dimension_cfg)

        samples.append({
            "index": idx,
            "input": question,
            "reference_answer": reference,
            "model_output": prediction,
            "score": judge_result.get("score"),
            "max_score": judge_result.get("max_score", dimension_cfg.get("score_max", 5)),
            "passed": judge_result.get("passed"),
            "judgement": judge_result.get("judgement"),
            "evaluation_reason": judge_result.get("evaluation_reason", ""),
            "error_type": judge_result.get("error_type"),
            "dimension_scores": [
                {"name": "judge_score", "score": judge_result.get("score", 0),
                 "max_score": judge_result.get("max_score", dimension_cfg.get("score_max", 5))},
            ] if judge_result else [],
            "status": "completed",
        })

        progress_pct = int(idx / max(total, 1) * 100)
        print(f"[eval] sample {idx}/{total} ({progress_pct}%) done")

    # ---- 4. Compute basic metrics ----
    print(f"[eval] computing basic metrics on {len(predictions)} predictions")
    metrics_result: dict[str, Any] = {}

    bleu_cfg = basic_cfg.get("bleu", {})
    if bleu_cfg.get("enabled"):
        metrics_result["bleu"] = _compute_bleu(references, predictions, int(bleu_cfg.get("ngram", 4)))

    rouge_cfg = basic_cfg.get("rouge", {})
    if rouge_cfg.get("enabled"):
        metrics_result["rouge"] = _compute_rouge(references, predictions, rouge_cfg.get("methods"))

    cosine_cfg = basic_cfg.get("cosine", {})
    if cosine_cfg.get("enabled"):
        metrics_result["cosine"] = _compute_cosine(references, predictions)
    metrics_result["exact_match"] = _compute_exact_match(references, predictions)
    metrics_result["text_similarity"] = _compute_text_similarity(references, predictions)

    # ---- 5. Summarise ----
    completed = len(samples)
    if judge_enabled:
        scored = [s for s in samples if s.get("score") is not None]
        passed_count = len([s for s in scored if s.get("passed")])
        avg_score = round(sum(s["score"] for s in scored) / max(len(scored), 1), output_precision)
        max_score = dimension_cfg.get("score_max", 5)
        overall_score = round(avg_score / max_score * 100, output_precision)
        overall_score_max = 100
        dimension_summary = [{
            "name": "综合评分",
            "score": overall_score,
            "max_score": 100,
            "pass_rate": round(passed_count / max(completed, 1) * 100, 1),
        }]
        overall_evaluation = f"评测完成：{completed} 样本，{passed_count} 通过，平均 {avg_score}/{max_score} 分"
    else:
        passed_count = 0
        enabled_scores = [
            float(item.get("score") or 0)
            for item in metrics_result.values()
            if isinstance(item, dict) and item.get("enabled", True) and item.get("score") is not None
        ]
        overall_score = round(sum(enabled_scores) / len(enabled_scores), output_precision) if enabled_scores else 0
        overall_score_max = 100
        dimension_summary = [
            {
                "name": name,
                "score": float(item.get("score") or 0),
                "max_score": 100,
                "pass_rate": float(item.get("score") or 0),
            }
            for name, item in metrics_result.items()
            if isinstance(item, dict) and item.get("enabled", True) and item.get("score") is not None
        ]
        overall_evaluation = f"评测完成：{completed} 样本（未配置 LLM 评委）"

    result = {
        "overall_score": overall_score,
        "overall_score_max": overall_score_max,
        "overall_evaluation": overall_evaluation,
        "improvement_suggestions": [],
        "dimension_summary": dimension_summary,
        "samples": samples,
        "sample_count": total,
        "completed_count": completed,
        "passed_count": passed_count,
        "basic_metrics": metrics_result,
    }

    # ---- 6. Write results ----
    result_path = output_dir / "eval_results.json"
    result_path.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"[eval] results written to {result_path}")
    return result


def main() -> None:
    import argparse
    parser = argparse.ArgumentParser(description="YG-FT Evaluation Runner")
    parser.add_argument("--config", required=True, help="Path to eval config JSON file")
    args = parser.parse_args()

    config_path = Path(args.config)
    if not config_path.exists():
        print(f"FATAL: config file not found: {args.config}", file=sys.stderr)
        sys.exit(1)

    config = json.loads(config_path.read_text(encoding="utf-8"))
    start = time.time()
    try:
        run_eval(config)
        elapsed = time.time() - start
        print(f"[eval] DONE in {elapsed:.1f}s")
    except Exception as exc:
        print(f"[eval] FAILED: {exc}", file=sys.stderr)
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()