#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Lazy-loaded shared data cache for data viewer tabs.
Loads data_viewer.jsonl once on first access, not at import time.
"""

from __future__ import annotations
import json
import pandas as pd
from pathlib import Path

BASE_DIR = Path(__file__).resolve().parent.parent

# Bucket mount point (HF Storage Bucket mounted at /data in Space runtime)
_BUCKET_DIR = Path("/data")
# Prefer bucket path if available, fallback to repo-local path
DATA_VIEWER_FILE = (
    _BUCKET_DIR / "data_viewer.jsonl"
    if (_BUCKET_DIR / "data_viewer.jsonl").exists()
    else BASE_DIR / "data" / "data_viewer.jsonl"
)
DATA_VIEWER_INDEX_FILE = BASE_DIR / "data" / "data_viewer_index.json"

# ICBCBench data viewer schema: dual-track scores + optional legacy DeepResearch Bench fields
_REQUIRED_COLS = [
    "model_name", "id", "prompt", "article",
    "overall_score", "track",  # track: "objective" or "subjective"
]
_OPTIONAL_SCORE_COLS = [
    "objective_score", "subjective_score",
    "expert_score", "citation_score", "source_quality_score",
    "confidence", "correct",
    # Legacy DeepResearch Bench fields (kept for backward compatibility)
    "comprehensiveness_score", "insight_score",
    "instruction_following_score", "readability_score",
]

_cache: pd.DataFrame | None = None
_index_cache: dict | None = None


def _normalize_record(item: dict) -> dict:
    """Ensure every record has the ICBCBench score fields, falling back to legacy fields."""
    record = dict(item)

    # Determine track
    if "track" not in record or not record["track"]:
        # Heuristic: if article is present and non-empty, treat as subjective
        article = record.get("article") or ""
        record["track"] = "subjective" if isinstance(article, str) and article.strip() else "objective"

    # Normalize scores to 0-100 scale if they look like ratios
    for key in ["overall_score", "objective_score", "subjective_score",
                "expert_score", "citation_score", "source_quality_score",
                "comprehensiveness_score", "insight_score",
                "instruction_following_score", "readability_score"]:
        val = record.get(key)
        if isinstance(val, (int, float)) and val is not None and val <= 1.0:
            record[key] = val * 100

    # If the new ICBCBench fields are missing but legacy fields exist, synthesize them.
    if record.get("subjective_score") is None:
        legacy = []
        for k in ["comprehensiveness_score", "insight_score",
                  "instruction_following_score", "readability_score"]:
            val = record.get(k)
            if val is None:
                continue
            try:
                legacy.append(float(val))
            except (TypeError, ValueError):
                pass
        if legacy:
            record["subjective_score"] = sum(legacy) / len(legacy)

    if record.get("overall_score") is None:
        candidates = []
        for k in ["objective_score", "subjective_score"]:
            val = record.get(k)
            if val is None:
                continue
            try:
                candidates.append(float(val))
            except (TypeError, ValueError):
                pass
        if candidates:
            record["overall_score"] = sum(candidates) / len(candidates)

    return record


def get_data() -> pd.DataFrame:
    global _cache
    if _cache is not None:
        return _cache

    records = []
    if DATA_VIEWER_FILE.exists():
        with DATA_VIEWER_FILE.open(encoding="utf-8") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                try:
                    records.append(_normalize_record(json.loads(line)))
                except json.JSONDecodeError:
                    continue

    df = pd.DataFrame(records)
    if df.empty or not all(c in df.columns for c in _REQUIRED_COLS):
        _cache = pd.DataFrame(columns=_REQUIRED_COLS + _OPTIONAL_SCORE_COLS)
    else:
        df["id"] = df["id"].astype(str)
        _cache = df
    return _cache


def get_index() -> dict:
    global _index_cache
    if _index_cache is not None:
        return _index_cache

    if DATA_VIEWER_INDEX_FILE.exists():
        try:
            _index_cache = json.loads(DATA_VIEWER_INDEX_FILE.read_text(encoding="utf-8"))
            return _index_cache
        except json.JSONDecodeError:
            pass

    models = set()
    tasks = {}
    if DATA_VIEWER_FILE.exists():
        with DATA_VIEWER_FILE.open(encoding="utf-8") as fh:
            for line in fh:
                if not line.strip():
                    continue
                try:
                    item = json.loads(line)
                except json.JSONDecodeError:
                    continue
                item = _normalize_record(item)
                model = item.get("model_name")
                item_id = str(item.get("id"))
                prompt = item.get("prompt") or ""
                if model:
                    models.add(model)
                if item_id and item_id not in tasks:
                    tasks[item_id] = prompt

    _index_cache = {
        "models": sorted(models),
        "tasks": [
            {"id": item_id, "prompt": tasks[item_id]}
            for item_id in sorted(tasks, key=lambda value: int(value) if value.isdigit() else value)
        ],
    }
    return _index_cache


def get_entry(model_name: str, item_id: str) -> dict | None:
    if not model_name or not item_id or not DATA_VIEWER_FILE.exists():
        return None

    item_id = str(item_id)
    index = get_index()
    location = index.get("lookup", {}).get(f"{model_name}\t{item_id}")
    if location:
        offset, length = location
        with DATA_VIEWER_FILE.open("rb") as fh:
            fh.seek(offset)
            line = fh.read(length).decode("utf-8")
        try:
            item = _normalize_record(json.loads(line))
            if item.get("model_name") == model_name and str(item.get("id")) == item_id:
                return item
        except json.JSONDecodeError:
            pass

    with DATA_VIEWER_FILE.open(encoding="utf-8") as fh:
        for line in fh:
            if not line.strip():
                continue
            try:
                item = _normalize_record(json.loads(line))
            except json.JSONDecodeError:
                continue
            if item.get("model_name") == model_name and str(item.get("id")) == item_id:
                return item
    return None


def get_entries_for_task(item_id: str, model_names: set[str]) -> dict[str, dict]:
    if not item_id or not model_names or not DATA_VIEWER_FILE.exists():
        return {}

    item_id = str(item_id)
    index = get_index()
    locations = {
        model: index.get("lookup", {}).get(f"{model}\t{item_id}")
        for model in model_names
    }
    locations = {model: loc for model, loc in locations.items() if loc}
    if locations:
        found = {}
        with DATA_VIEWER_FILE.open("rb") as fh:
            for model, (offset, length) in locations.items():
                fh.seek(offset)
                try:
                    item = _normalize_record(json.loads(fh.read(length).decode("utf-8")))
                    if item.get("model_name") == model and str(item.get("id")) == item_id:
                        found[model] = item
                except json.JSONDecodeError:
                    pass
        if len(found) == len(locations):
            return found

    found = {}
    with DATA_VIEWER_FILE.open(encoding="utf-8") as fh:
        for line in fh:
            if not line.strip():
                continue
            try:
                item = _normalize_record(json.loads(line))
            except json.JSONDecodeError:
                continue
            model = item.get("model_name")
            if str(item.get("id")) == item_id and model in model_names:
                found[model] = item
                if len(found) == len(model_names):
                    break
    return found