"""Transform SWE-bench harness outputs into router-ready data."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Iterable, Mapping, MutableMapping

import numpy as np
import pandas as pd

from ..utils import get_embedding, loadjson, savepkl
from .schema import PASS_FAIL_METRIC, SWEbenchExample, SWEbenchResult
from .prompt_builder import build_task_description


def _embed_cached(text: str, cache: MutableMapping[str, list[float]]) -> list[float]:
    if text in cache:
        return cache[text]
    embedding = get_embedding([text])
    cache[text] = np.asarray(embedding[0], dtype=float).tolist()
    return cache[text]


def load_results_jsonl(
    path: Path, examples_by_task: Mapping[str, SWEbenchExample]
) -> list[SWEbenchResult]:
    """
    Load SWE-bench harness results from a JSONL file.

    Expected per-line fields:
    - task_id (must align with repo:instance_id)
    - llm_name
    - prompt
    - response
    - pass_rate (float 0..1)
    - latency_seconds (float)
    - cost_usd (float)
    """
    results: list[SWEbenchResult] = []
    raw = path.read_text(encoding="utf-8").splitlines()
    for lineno, line in enumerate(raw, start=1):
        if not line.strip():
            continue
        try:
            payload = json.loads(line)
        except json.JSONDecodeError as exc:
            raise ValueError(f"Invalid JSON on line {lineno} of {path}: {exc}") from exc

        task_id = payload["task_id"]
        if task_id not in examples_by_task:
            raise KeyError(f"SWE-bench task_id '{task_id}' missing from loaded examples.")
        example = examples_by_task[task_id]
        try:
            pass_rate = float(payload["pass_rate"])
            latency_seconds = float(payload["latency_seconds"])
            cost_usd = float(payload["cost_usd"])
        except KeyError as exc:
            raise KeyError(f"Required SWE-bench result field missing: {exc.args[0]}") from exc

        results.append(
            SWEbenchResult(
                example=example,
                llm_name=payload["llm_name"],
                prompt=payload["prompt"],
                response=str(payload.get("response", "")),
                pass_rate=pass_rate,
                latency_seconds=latency_seconds,
                cost_usd=cost_usd,
            )
        )
    return results


class SWEbenchRouterDataBuilder:
    """Build router data artifacts from SWE-bench harness outputs."""

    def __init__(
        self,
        llm_description_path: Path,
        router_output_path: Path,
        llm_embedding_path: Path,
    ) -> None:
        self.llm_description_path = llm_description_path
        self.router_output_path = router_output_path
        self.llm_embedding_path = llm_embedding_path

    def build(self, results: Iterable[SWEbenchResult]) -> Path:
        """
        Convert harness results into router CSV + LLM description embeddings.

        Args:
            results: iterable of SWEbenchResult from an external harness run.

        Returns:
            Path to the router CSV written on disk.
        """
        results_list = list(results)
        if not results_list:
            raise ValueError("No SWE-bench results provided; nothing to build.")

        query_cache: dict[str, list[float]] = {}
        desc_cache: dict[str, list[float]] = {}
        records: list[dict[str, object]] = []

        for result in results_list:
            example = result.example
            task_desc = build_task_description(example)
            query_embedding = _embed_cached(result.prompt, query_cache)
            task_desc_embedding = _embed_cached(task_desc, desc_cache)
            records.append(
                {
                    "task_id": example.task_id,
                    "task_description": task_desc,
                    "task_description_embedding": task_desc_embedding,
                    "query": result.prompt,
                    "query_embedding": query_embedding,
                    "ground_truth": "",  # The harness score is authoritative.
                    "metric": PASS_FAIL_METRIC,
                    "llm": result.llm_name,
                    "effect": result.pass_rate,
                    "cost_raw": result.cost_usd,
                    "latency_raw": result.latency_seconds,
                    "cost": result.cost_usd,
                    "latency": result.latency_seconds,
                }
            )

        df = pd.DataFrame.from_records(records)
        self._normalize_cost_latency(df)

        self.router_output_path.parent.mkdir(parents=True, exist_ok=True)
        df.to_csv(self.router_output_path, index=False)

        llm_description = loadjson(self.llm_description_path)
        llm_features = [llm_description[name]["feature"] for name in llm_description]
        llm_description_embedding = get_embedding(llm_features)
        savepkl(llm_description_embedding, self.llm_embedding_path)

        metadata_path = self.router_output_path.with_suffix(
            self.router_output_path.suffix + ".metadata.json"
        )
        metadata_path.write_text(
            json.dumps(
                {
                    "llm_names": list(llm_description.keys()),
                    "total_rows": int(len(df)),
                    # Uses unique queries across all LLMs; assumes every query is paired with each LLM.
                    "total_queries": int(df[["task_id", "query"]].drop_duplicates().shape[0]),
                    "task_query_counts": self._task_counts(df),
                },
                indent=2,
            ),
            encoding="utf-8",
        )

        return self.router_output_path

    @staticmethod
    def _normalize_cost_latency(frame: pd.DataFrame) -> None:
        def _minmax(series: pd.Series) -> pd.Series:
            if series.isna().all():
                return pd.Series(0.0, index=series.index)
            maximum = series.max()
            minimum = series.min()
            denom = maximum - minimum
            if denom == 0:
                return pd.Series(0.0, index=series.index)
            return (series - minimum) / denom

        frame["cost"] = frame.groupby("task_id")["cost_raw"].transform(_minmax)
        if "latency_raw" in frame.columns:
            latency_clipped = frame.groupby("task_id")["latency_raw"].transform(
                lambda s: s.clip(upper=s.quantile(0.99))
            )
            frame["latency"] = latency_clipped.groupby(frame["task_id"]).transform(_minmax)

    @staticmethod
    def _task_counts(frame: pd.DataFrame) -> dict[str, int]:
        unique_pairs = frame.groupby("task_id")[["query"]].nunique()
        return {task_id: int(row["query"]) for task_id, row in unique_pairs.iterrows()}
