import argparse
import copy
import json
import os
import re
from collections import Counter
from glob import glob
from typing import Any, Dict, List, Optional, Tuple


DEFAULT_TEMPLATE_PATH = "outputs/predictions/track_1_test.json"
DEFAULT_PRED_PATH = "outputs/predictions/predict_ckpt2660/generated_predictions.jsonl"
DEFAULT_OUTPUT_PATH = "outputs/submissions/answers/track_1_test.json"


def load_json(path: str) -> Any:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def load_jsonl(path: str) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    with open(path, "r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, start=1):
            raw = line.strip()
            if not raw:
                continue
            try:
                rows.append(json.loads(raw))
            except Exception:
                print(f"[WARN] skip invalid jsonl line: {line_no}")
    return rows


def normalize_level(level: Any) -> Optional[str]:
    if level is None:
        return None
    s = str(level).strip()
    if not s:
        return None

    m = {
        "poor": "Poor",
        "medium": "Medium",
        "good": "Good",
        "a": "A",
        "b": "B",
        "c": "C",
    }
    return m.get(s.lower(), s)


def score_to_submission_level(score: float, low_threshold: float = 5.0, high_threshold: float = 7.0) -> str:
    # 比赛映射：A=Poor(0~5), B=Medium(5~7), C=Good(7~10)
    if score < low_threshold:
        return "A"
    if score < high_threshold:
        return "B"
    return "C"


def to_submission_level(level: Any) -> Optional[str]:
    """
    Submission mapping required by Track-1 template:
    A -> Poor, B -> Medium, C -> Good
    """
    norm = normalize_level(level)
    if norm is None:
        return None
    m = {
        "Poor": "A",
        "Medium": "B",
        "Good": "C",
        "A": "A",
        "B": "B",
        "C": "C",
    }
    return m.get(norm)


def normalize_prediction_text(pred_row: Dict[str, Any]) -> str:
    # 兼容不同推理后端字段名。
    for k in ("predict", "response", "output", "text", "generation"):
        if k in pred_row:
            return str(pred_row.get(k, ""))
    return str(pred_row)


def extract_level_from_text(pred_text: str, crit_key: str) -> Optional[str]:
    pattern = rf'"{re.escape(crit_key)}"\s*:\s*\{{.*?"level"\s*:\s*"([^"]+)"'
    m = re.search(pattern, pred_text, flags=re.IGNORECASE | re.DOTALL)
    if not m:
        return None
    return to_submission_level(m.group(1))


def try_parse_predict_json(text: str) -> Optional[Dict[str, Any]]:
    text = str(text or "").strip()
    if not text:
        return None

    # 1) 直接解析
    try:
        obj = json.loads(text)
        if isinstance(obj, dict):
            return obj
    except Exception:
        pass

    # 2) 尝试提取从首个 { 到最后一个 } 的片段
    left = text.find("{")
    right = text.rfind("}")
    if left != -1 and right != -1 and left < right:
        snippet = text[left:right + 1]
        try:
            obj = json.loads(snippet)
            if isinstance(obj, dict):
                return obj
        except Exception:
            pass

    return None


def extract_answer(pred_obj: Optional[Dict[str, Any]], pred_text: str) -> Optional[str]:
    if pred_obj is not None:
        answer = str(pred_obj.get("answer", "")).strip().upper()
        if answer in {"A", "B", "C", "D"}:
            return answer

    m = re.search(r'"answer"\s*:\s*"([A-D])"', pred_text, re.IGNORECASE)
    if m:
        return m.group(1).upper()

    # 兼容 QA-only 推理：模型可能只输出单个字母（如 "C"）。
    raw = str(pred_text or "").strip().upper()
    if raw:
        # 情况1：整行仅有一个候选字母（允许尾随标点）。
        m = re.match(r"^\s*([A-D])(?:[\.\)\]:：]|\s)*$", raw)
        if m:
            return m.group(1)
        # 情况2：短文本中只出现唯一一个 A/B/C/D，且不包含常见 JSON/选项结构。
        if len(raw) <= 12 and ("{" not in raw) and ("\"" not in raw):
            hits = re.findall(r"[A-D]", raw)
            if len(hits) == 1:
                return hits[0]
    return None


def extract_total_score(pred_obj: Optional[Dict[str, Any]], pred_text: str) -> Optional[float]:
    val: Optional[float] = None
    if pred_obj is not None and "total_score" in pred_obj:
        try:
            val = float(pred_obj["total_score"])
        except Exception:
            val = None

    if val is None:
        m = re.search(r'"total_score"\s*:\s*([0-9]+(?:\.[0-9]+)?)', pred_text)
        if m:
            val = float(m.group(1))

    if val is None:
        return None

    return max(0.0, min(100.0, float(val)))


def merge_total_scores(total_votes: List[float], mode: str) -> Optional[int]:
    if not total_votes:
        return None

    vals = [float(v) for v in total_votes]
    if mode == "trim_mean" and len(vals) >= 3:
        vals = sorted(vals)[1:-1]
    merged = sum(vals) / len(vals)
    return max(0, min(100, int(round(merged))))


def parse_score_from_text(pred_text: str, crit_key: str) -> Optional[float]:
    pattern = rf'"{re.escape(crit_key)}"\s*:\s*\{{.*?"score"\s*:\s*([0-9]+(?:\.[0-9]+)?)'
    m = re.search(pattern, pred_text, flags=re.IGNORECASE | re.DOTALL)
    if not m:
        return None
    try:
        return float(m.group(1))
    except Exception:
        return None


def extract_one_criteria_level_and_score(
    crit_key: str,
    pred_obj: Optional[Dict[str, Any]],
    pred_text: str,
) -> Tuple[Optional[str], Optional[float]]:
    score: Optional[float] = None
    level: Optional[str] = None

    if isinstance(pred_obj, dict):
        src = pred_obj.get("criteria", {})
        if isinstance(src, dict) and crit_key in src:
            value = src.get(crit_key)
            if isinstance(value, dict):
                if "score" in value:
                    try:
                        score = float(value["score"])
                    except Exception:
                        score = None
                level = to_submission_level(value.get("level"))
            else:
                level = to_submission_level(value)

    if score is None:
        score = parse_score_from_text(pred_text, crit_key)

    if level is None:
        level = extract_level_from_text(pred_text, crit_key)

    if level not in {"A", "B", "C"}:
        level = None

    return level, score


def choose_majority_level(
    levels: List[str],
    scores: List[float],
    fallback_level: Optional[str],
    low_threshold: float,
    high_threshold: float,
) -> str:
    if levels:
        cnt = Counter(levels)
        top_n = max(cnt.values())
        top_levels = sorted([k for k, v in cnt.items() if v == top_n])
        if len(top_levels) == 1:
            return top_levels[0]

    # 平票时，用多次预测的均值 score 判定等级。
    if scores:
        mean_score = sum(scores) / len(scores)
        return score_to_submission_level(mean_score, low_threshold=low_threshold, high_threshold=high_threshold)

    if fallback_level in {"A", "B", "C"}:
        return fallback_level
    return "B"


def parse_weights(raw_weights: Optional[List[float]], n_models: int) -> List[float]:
    # 默认等权；若传入权重则要求和预测文件数一致。
    if raw_weights is None:
        return [1.0] * n_models
    if len(raw_weights) != n_models:
        raise ValueError(
            f"--weights length ({len(raw_weights)}) must equal number of prediction files ({n_models})."
        )
    for w in raw_weights:
        if w < 0:
            raise ValueError("weights must be non-negative.")
    # 全零没有意义，回退等权。
    if sum(raw_weights) == 0:
        return [1.0] * n_models
    return raw_weights


def load_thresholds(
    thresholds_json: str,
) -> Tuple[Dict[str, Dict[str, float]], Dict[str, float]]:
    """
    Accepts JSON in either format:
    1) {"criteria": {"Color Harmony": {"low": 4.9, "high": 7.1}}, "default": {"low":5,"high":7}}
    2) {"Color Harmony": {"low": 4.9, "high": 7.1}, ...}
    """
    default = {"low": 5.0, "high": 7.0}
    per_criteria: Dict[str, Dict[str, float]] = {}

    if not thresholds_json:
        return per_criteria, default
    if not os.path.exists(thresholds_json):
        print(f"[WARN] thresholds file not found: {thresholds_json}, fallback to default 5/7")
        return per_criteria, default

    obj = load_json(thresholds_json)
    if not isinstance(obj, dict):
        print(f"[WARN] invalid thresholds json format: {thresholds_json}, fallback to default 5/7")
        return per_criteria, default

    if "default" in obj and isinstance(obj.get("default"), dict):
        d = obj["default"]
        low = d.get("low", 5.0)
        high = d.get("high", 7.0)
        try:
            low_f = float(low)
            high_f = float(high)
            if low_f < high_f:
                default = {"low": low_f, "high": high_f}
        except Exception:
            pass

    src = obj.get("criteria") if isinstance(obj.get("criteria"), dict) else obj
    if isinstance(src, dict):
        for k, v in src.items():
            if not isinstance(v, dict):
                continue
            if "low" not in v or "high" not in v:
                continue
            try:
                low = float(v["low"])
                high = float(v["high"])
            except Exception:
                continue
            if low < high:
                per_criteria[str(k)] = {"low": low, "high": high}

    return per_criteria, default


def resolve_best_index(best_index: int, weights: List[float]) -> int:
    # best_index=-1 表示自动选择权重最高的模型作为平票时的优先模型。
    if best_index >= 0:
        if best_index >= len(weights):
            raise ValueError(f"--best_index out of range: {best_index}, num_models={len(weights)}")
        return best_index
    return max(range(len(weights)), key=lambda i: weights[i])


def extract_criteria_voting(
    template_item: Dict[str, Any],
    pred_objs: List[Optional[Dict[str, Any]]],
    pred_texts: List[str],
    per_criteria_thresholds: Dict[str, Dict[str, float]],
    default_thresholds: Dict[str, float],
) -> Dict[str, Dict[str, str]]:
    out: Dict[str, Dict[str, str]] = {}

    for crit_key, crit_val in template_item.get("criteria", {}).items():
        thresholds = per_criteria_thresholds.get(crit_key, default_thresholds)
        low = float(thresholds.get("low", 5.0))
        high = float(thresholds.get("high", 7.0))
        if not (low < high):
            low, high = 5.0, 7.0

        level_votes: List[str] = []
        score_votes: List[float] = []
        for obj, text in zip(pred_objs, pred_texts):
            level, score = extract_one_criteria_level_and_score(crit_key, obj, text)
            if level is not None:
                level_votes.append(level)
            if score is not None:
                score_votes.append(score)

        prev = str(crit_val.get("level", "")).strip() if isinstance(crit_val, dict) else ""
        final_level = choose_majority_level(
            level_votes,
            score_votes,
            prev,
            low_threshold=low,
            high_threshold=high,
        )
        out[crit_key] = {"level": final_level}

    return out


def pick_default_predictions_path() -> str:
    candidates = [
        "outputs/predictions/predict_ckpt2660/generated_predictions.jsonl",
        "outputs/predictions/generated_predictions.jsonl",
    ]
    # 自动兜底：在预测目录里找最近一次 generated_predictions.jsonl。
    dynamic = sorted(
        glob("outputs/predictions/**/generated_predictions.jsonl", recursive=True),
        key=lambda x: os.path.getmtime(x),
        reverse=True,
    )
    candidates = dynamic + candidates

    for p in candidates:
        if os.path.exists(p):
            return p
    return DEFAULT_PRED_PATH


def choose_weighted_answer(
    votes_by_model: List[Optional[str]],
    weights: List[float],
    best_index: int,
    tie_break_answer: Optional[str],
) -> str:
    """
    更合理的答案融合策略：
    1) 加权投票（按各变体可靠性权重）
    2) 若平票且提供 tie-break 结果，则优先用 tie-break
    3) 若仍平票，采用最佳变体(best_index)在平票选项中的答案
    4) 最后才做稳定兜底（字母序）
    """
    label_scores = {"A": 0.0, "B": 0.0, "C": 0.0, "D": 0.0}
    for i, ans in enumerate(votes_by_model):
        if ans in label_scores:
            label_scores[ans] += weights[i]

    max_score = max(label_scores.values())
    if max_score <= 0:
        return "A"

    tied = sorted([k for k, v in label_scores.items() if v == max_score])
    if len(tied) == 1:
        return tied[0]

    if tie_break_answer in tied:
        return tie_break_answer  # 用专门 tie-break 结果判平票

    best_vote = votes_by_model[best_index]
    if best_vote in tied:
        return str(best_vote)

    return tied[0]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--template_json", type=str, default=DEFAULT_TEMPLATE_PATH)
    parser.add_argument(
        "--predictions_jsonl",
        type=str,
        nargs="+",
        default=None,
        help="One or more generated_predictions.jsonl paths for voting.",
    )
    parser.add_argument(
        "--weights",
        type=float,
        nargs="+",
        default=None,
        help="Optional weights for prediction files (same length as --predictions_jsonl).",
    )
    parser.add_argument(
        "--best_index",
        type=int,
        default=-1,
        help="Best model index for tie fallback. -1 means auto argmax(weights).",
    )
    parser.add_argument(
        "--tie_break_jsonl",
        type=str,
        default="",
        help="Optional tie-break predictions file used only when weighted vote ties.",
    )
    parser.add_argument(
        "--thresholds_json",
        type=str,
        default="",
        help="Optional per-criterion thresholds json for score->A/B/C mapping.",
    )
    parser.add_argument(
        "--total_score_fusion",
        type=str,
        default="mean",
        choices=["mean", "trim_mean"],
        help="Fusion mode for total_score across multi-prompt predictions.",
    )
    parser.add_argument("--output_json", type=str, default=DEFAULT_OUTPUT_PATH)
    args = parser.parse_args()

    pred_paths = args.predictions_jsonl if args.predictions_jsonl else [pick_default_predictions_path()]
    pred_paths = [p for p in pred_paths if str(p).strip()]
    if not pred_paths:
        raise ValueError("No predictions_jsonl provided or discovered.")

    template_data = load_json(args.template_json)
    pred_sets = [load_jsonl(p) for p in pred_paths]
    weights = parse_weights(args.weights, len(pred_sets))
    best_index = resolve_best_index(args.best_index, weights)
    per_criteria_thresholds, default_thresholds = load_thresholds(args.thresholds_json)
    tie_break_rows: List[Dict[str, Any]] = []
    if args.tie_break_jsonl:
        tie_break_rows = load_jsonl(args.tie_break_jsonl)

    if not isinstance(template_data, list):
        raise ValueError("template_json must be a list.")

    print(f"[INFO] template items: {len(template_data)}")
    for p, rows in zip(pred_paths, pred_sets):
        print(f"[INFO] prediction rows: {len(rows)} ({p})")
    print(f"[INFO] answer weights: {weights}")
    print(f"[INFO] answer best_index: {best_index}")
    print(f"[INFO] criteria thresholds default: low={default_thresholds['low']}, high={default_thresholds['high']}")
    if per_criteria_thresholds:
        print(f"[INFO] criteria thresholds loaded: {len(per_criteria_thresholds)}")
    if args.tie_break_jsonl:
        print(f"[INFO] tie_break rows: {len(tie_break_rows)} ({args.tie_break_jsonl})")

    final_submission: List[Dict[str, Any]] = []
    parsed_ok = 0
    filled = 0

    for i, item in enumerate(template_data):
        new_item = copy.deepcopy(item)
        # 收集每次预测在第 i 条样本上的结果。
        row_texts: List[str] = []
        row_objs: List[Optional[Dict[str, Any]]] = []
        answer_votes_by_model: List[Optional[str]] = []
        total_votes: List[int] = []

        for rows in pred_sets:
            if i >= len(rows):
                answer_votes_by_model.append(None)
                continue
            pred_text = normalize_prediction_text(rows[i])
            pred_obj = try_parse_predict_json(pred_text)
            if pred_obj is not None:
                parsed_ok += 1
            row_texts.append(pred_text)
            row_objs.append(pred_obj)

            ans = extract_answer(pred_obj, pred_text)
            answer_votes_by_model.append(ans)

            ts = extract_total_score(pred_obj, pred_text)
            if ts is not None:
                total_votes.append(ts)

        if not row_texts:
            final_submission.append(new_item)
            continue

        new_item["criteria"] = extract_criteria_voting(
            new_item,
            row_objs,
            row_texts,
            per_criteria_thresholds=per_criteria_thresholds,
            default_thresholds=default_thresholds,
        )

        # total_score 用多次预测均值，减少单次抖动。
        merged_total = merge_total_scores(total_votes, args.total_score_fusion)
        if merged_total is not None:
            new_item["total_score"] = max(0, min(100, merged_total))

        tie_break_answer: Optional[str] = None
        if i < len(tie_break_rows):
            tb_text = normalize_prediction_text(tie_break_rows[i])
            tb_obj = try_parse_predict_json(tb_text)
            tie_break_answer = extract_answer(tb_obj, tb_text)

        new_item["answer"] = choose_weighted_answer(
            votes_by_model=answer_votes_by_model,
            weights=weights,
            best_index=best_index,
            tie_break_answer=tie_break_answer,
        )

        final_submission.append(new_item)
        filled += 1

    os.makedirs(os.path.dirname(args.output_json), exist_ok=True)
    with open(args.output_json, "w", encoding="utf-8") as f:
        json.dump(final_submission, f, ensure_ascii=False, indent=2)

    parsed_total = filled * len(pred_paths)
    print(f"[INFO] parsed predict json ok: {parsed_ok}/{parsed_total}")
    print(f"[INFO] saved submission: {args.output_json}")


if __name__ == "__main__":
    main()