from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from fastapi import APIRouter, Header, HTTPException

from prompt_auto_lab.api.models import MetricsSnapshot, ParetoFrontResponse
from prompt_auto_lab.api.services import user_service
from prompt_auto_lab.core.pareto_utils import select_pareto_candidates

router = APIRouter(prefix="/api/versions", tags=["versions"])


def _get_versions_file(user_id: str, session_id: str) -> Path:
    """
    Get path to versions storage file.
    Uses the user-scoped session directory under the configured data store:
    {store_dir}/users/{user_id}/sessions/{session_id}/version.jsonl.
    Note: The orchestrator writes to 'version.jsonl' (singular).
    """
    session_dir = user_service.get_user_session_dir(user_id, session_id)
    return session_dir / "version.jsonl"


def _load_versions(user_id: str, session_id: str) -> list[dict[str, Any]]:
    """Load all versions from JSONL file."""
    versions_file = _get_versions_file(user_id, session_id)
    if not versions_file.exists():
        return []

    versions = []
    try:
        with open(versions_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    versions.append(json.loads(line))
    except Exception:
        pass

    return versions


def _load_latest_version(user_id: str, session_id: str) -> dict[str, Any] | None:
    """Load only the last line from JSONL file (cumulative state)."""
    versions_file = _get_versions_file(user_id, session_id)
    if not versions_file.exists():
        return None

    last_line = None
    try:
        with open(versions_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    last_line = line
        if last_line:
            return json.loads(last_line)
    except Exception:
        pass

    return None


def _resolve_version_id(raw_version: Any, iteration: dict[str, Any] | None, fallback_idx: int, used: set[int | str]) -> int | str:
    """
    Resolve a unique version identifier for a record.
    Priority:
      1) explicit version if present and unique
      2) highest iteration index + 1 (latest iteration count)
      3) fallback to running index
    """
    if raw_version is not None and raw_version not in used:
        return raw_version

    if iteration:
        try:
            max_iter = max(int(k) for k in iteration.keys())
            candidate = max_iter + 1
            if candidate not in used:
                return candidate
        except Exception:
            pass

    return fallback_idx


def _iter_blocks(version: dict[str, Any]) -> list[tuple[int | str, dict[str, Any]]]:
    """Extract iteration blocks sorted by numeric key when possible.

    The "root" iteration (baseline) is always placed first, followed by
    numeric iterations in ascending order.
    """
    iteration = version.get("iteration") or {}
    if not isinstance(iteration, dict):
        return []

    def sort_key(k: str) -> tuple[int, str]:
        # "root" should come first (baseline iteration)
        if k == "root":
            return (-1, k)
        try:
            return (int(k), k)
        except Exception:
            return (10**9, k)

    items = sorted(iteration.items(), key=lambda kv: sort_key(kv[0]))
    result: list[tuple[int | str, dict[str, Any]]] = []
    for k, v in items:
        try:
            k_norm: int | str = int(k)
        except Exception:
            k_norm = k
        result.append((k_norm, v))
    return result


def _normalize_dimensions(raw_dims: dict[str, Any] | None) -> dict[str, float]:
    """
    Normalize dimension_metrics values to simple floats.

    Handles cases where dimension values might be objects like:
    {"name": "BackstoryJudge", "score": 0.42, "diff": -0.04, ...}
    instead of simple floats like 0.42.
    """
    if not raw_dims or not isinstance(raw_dims, dict):
        return {}

    result: dict[str, float] = {}
    for key, value in raw_dims.items():
        if isinstance(value, (int, float)):
            result[key] = float(value)
        elif isinstance(value, dict):
            # Extract score from object-style dimension metrics
            score = value.get("score")
            if isinstance(score, (int, float)):
                result[key] = float(score)
            # Fallback: if no score, skip this dimension
    return result


def _extract_metrics_fields(version: dict[str, Any], iter_block: dict[str, Any] | None = None) -> dict[str, Any]:
    """
    Extract metrics from the best candidate in an iteration block.
    Selects the candidate with highest overall_score and returns its metrics.
    """
    overall_score: float | None = None
    dimensions: dict[str, float] = {}
    experiment_id: str = ""
    experiment_link: str = ""
    commit_id: str = ""

    # If iteration block provided, use it; otherwise fallback to latest iteration
    if iter_block is None:
        iteration = version.get("iteration") or {}
        if iteration:
            try:
                last_key = sorted(iteration.keys(), key=lambda k: int(k) if k != "root" else -1)[-1]
                iter_block = iteration.get(last_key)
            except Exception:
                iter_block = next(iter(iteration.values()), None)

    if isinstance(iter_block, dict):
        candidates = iter_block.get("candidates")
        if isinstance(candidates, dict) and candidates:
            # Find the best candidate by overall_score
            best_candidate: dict[str, Any] | None = None
            best_score: float | None = None
            for cand in candidates.values():
                if not isinstance(cand, dict):
                    continue
                cand_score = cand.get("overall_score")
                if isinstance(cand_score, (int, float)):
                    if best_score is None or cand_score > best_score:
                        best_score = cand_score
                        best_candidate = cand

            if best_candidate:
                overall_score = best_candidate.get("overall_score")
                raw_dims = best_candidate.get("dimension_metrics") or best_candidate.get("metrics") or {}
                dimensions = _normalize_dimensions(raw_dims)
                experiment_id = best_candidate.get("experiment_id") or ""
                experiment_link = best_candidate.get("experiment_link") or ""
                commit_id = best_candidate.get("commit_id") or ""

        # Fallback to iteration-level fields if no candidates found
        if overall_score is None:
            overall_score = iter_block.get("overall_score")
        if not dimensions:
            metrics_block = iter_block.get("metrics") or {}
            if metrics_block:
                raw_fallback = metrics_block.get("dimensions") or metrics_block
                dimensions = _normalize_dimensions(raw_fallback)
        if not dimensions:
            known_dims = ["completeness", "clarity_readability", "engagement", "factual_accuracy", "groundedness"]
            direct_dims = {k: float(iter_block[k]) for k in known_dims if k in iter_block and isinstance(iter_block.get(k), (int, float))}
            if direct_dims:
                dimensions = direct_dims
        if not experiment_id:
            experiment_id = iter_block.get("experiment_id") or ""
        if not experiment_link:
            experiment_link = iter_block.get("experiment_link") or ""
        if not commit_id:
            commit_id = iter_block.get("commit_id") or ""

    return {
        "overall_score": overall_score if overall_score is not None else 0.0,
        "dimensions": dimensions or {},
        "experiment_id": experiment_id,
        "experiment_link": experiment_link,
        "commit_id": commit_id,
    }


@router.get("")
async def list_versions(
    session_id: str | None = None,
    x_user_id: str = Header(None),
) -> list[dict[str, Any]]:
    """List all version records."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")
    if not session_id:
        raise HTTPException(status_code=400, detail="session_id is required")
    return _load_versions(x_user_id, session_id)


@router.get("/history/summary")
async def get_version_history(
    session_id: str | None = None,
    x_user_id: str = Header(None),
) -> dict[str, Any]:
    """Get version history summary with best/current version.

    Reads only the last line from version.jsonl (cumulative state).
    For each iteration, selects the best candidate by overall_score.
    """
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")
    if not session_id:
        raise HTTPException(status_code=400, detail="session_id is required")

    latest = _load_latest_version(x_user_id, session_id)

    if not latest:
        return {
            "versions": [],
            "best_version": None,
            "current_version": None,
            "total_versions": 0,
        }

    # Extract iteration blocks and build snapshots
    iter_blocks = _iter_blocks(latest)
    snapshots = []

    if iter_blocks:
        for display_idx, (_, iter_block) in enumerate(iter_blocks, start=1):
            metrics = _extract_metrics_fields(latest, iter_block)
            snapshots.append({
                "version": display_idx,
                "timestamp": latest.get("timestamp", ""),
                "overall_score": metrics["overall_score"],
                "dimensions": metrics["dimensions"],
                "experiment_id": metrics["experiment_id"],
                "experiment_link": metrics["experiment_link"],
                "commit_id": metrics["commit_id"],
            })
    else:
        # Fallback for records without iteration structure
        metrics = _extract_metrics_fields(latest, None)
        snapshots.append({
            "version": 1,
            "timestamp": latest.get("timestamp", ""),
            "overall_score": metrics["overall_score"],
            "dimensions": metrics["dimensions"],
            "experiment_id": metrics["experiment_id"],
            "experiment_link": metrics["experiment_link"],
            "commit_id": metrics["commit_id"],
        })

    # Determine best/current based on snapshots
    best_version_id = None
    current_version_id = None
    if snapshots:
        best_snapshot = max(snapshots, key=lambda s: s.get("overall_score", 0))
        best_version_id = best_snapshot.get("version")
        current_version_id = len(snapshots)

    return {
        "versions": snapshots,
        "best_version": best_version_id,
        "current_version": current_version_id,
        "total_versions": len(snapshots),
    }


def _extract_best_candidate(node_data: dict[str, Any]) -> dict[str, Any]:
    """Extract the best candidate from an iteration node.

    Looks for candidates in node_data.candidates and returns the one with highest overall_score.
    Falls back to node-level fields if no candidates found.
    """
    candidates = node_data.get("candidates")
    if isinstance(candidates, dict) and candidates:
        best_candidate: dict[str, Any] | None = None
        best_score: float | None = None
        best_candidate_id: str | None = None
        for cand_id, cand in candidates.items():
            if not isinstance(cand, dict):
                continue
            cand_score = cand.get("overall_score")
            if isinstance(cand_score, (int, float)):
                if best_score is None or cand_score > best_score:
                    best_score = cand_score
                    best_candidate = cand
                    best_candidate_id = cand_id

        if best_candidate:
            return {
                "overall_score": best_candidate.get("overall_score"),
                "experiment_id": best_candidate.get("experiment_id") or "",
                "experiment_link": best_candidate.get("experiment_link") or "",
                "commit_id": best_candidate.get("commit_id") or "",
                "branch": best_candidate.get("branch") or "",
                "candidate_id": best_candidate_id or "",
            }

    # Fallback to node-level fields
    return {
        "overall_score": node_data.get("overall_score"),
        "experiment_id": node_data.get("experiment_id") or "",
        "experiment_link": node_data.get("experiment_link") or "",
        "commit_id": node_data.get("commit_id") or "",
        "branch": node_data.get("analysis_branch") or "",
        "candidate_id": "",
    }


@router.get("/iteration-tree")
async def get_iteration_tree(
    session_id: str | None = None,
    x_user_id: str = Header(None),
) -> dict[str, Any]:
    """Get iteration tree structure for visualization.

    Returns tree nodes with parent-child relationships for the iteration graph.
    """
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")
    if not session_id:
        raise HTTPException(status_code=400, detail="session_id is required")

    latest = _load_latest_version(x_user_id, session_id)

    if not latest:
        return {
            "nodes": [],
            "best_node": None,
        }

    iteration = latest.get("iteration") or {}
    if not isinstance(iteration, dict):
        return {
            "nodes": [],
            "best_node": None,
        }

    nodes = []
    best_node = None
    best_score = None

    for node_id, node_data in iteration.items():
        if not isinstance(node_data, dict):
            continue

        # Extract best candidate info from this iteration node
        candidate_info = _extract_best_candidate(node_data)
        overall_score = candidate_info["overall_score"]
        # Use analysis_source_node for tree relationship (which node was analyzed to optimize from)
        analysis_source_node = node_data.get("analysis_source_node")

        node = {
            "id": node_id,
            "parent_id": analysis_source_node,  # Use analysis_source_node as parent
            "overall_score": overall_score if overall_score is not None else 0.0,
            "experiment_id": candidate_info["experiment_id"],
            "experiment_link": candidate_info["experiment_link"],
            "commit_id": candidate_info["commit_id"],
            "branch": candidate_info["branch"],  # Use candidate's branch
            "candidate_id": candidate_info["candidate_id"],
            "optimization_summary": node_data.get("optimization_summary") or "",
            "if_optimized": node_data.get("if_optimized"),
        }
        nodes.append(node)

        # Track best node
        if overall_score is not None:
            if best_score is None or overall_score > best_score:
                best_score = overall_score
                best_node = node_id

    return {
        "nodes": nodes,
        "best_node": best_node,
    }


@router.get("/compare/{v1}/{v2}")
async def compare_versions(
    v1: int,
    v2: int,
    session_id: str | None = None,
    x_user_id: str = Header(None),
) -> dict[str, Any]:
    """Compare two versions."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")
    if not session_id:
        raise HTTPException(status_code=400, detail="session_id is required")
    versions = _load_versions(x_user_id, session_id)

    version1 = None
    version2 = None

    for v in versions:
        if v.get("version") == v1:
            version1 = v
        if v.get("version") == v2:
            version2 = v

    if version1 is None:
        raise HTTPException(status_code=404, detail=f"Version {v1} not found")
    if version2 is None:
        raise HTTPException(status_code=404, detail=f"Version {v2} not found")

    # Calculate score difference (normalize if embedded in iteration)
    score1 = _extract_metrics_fields(version1)["overall_score"]
    score2 = _extract_metrics_fields(version2)["overall_score"]

    return {
        "version1": version1,
        "version2": version2,
        "score_diff": score2 - score1,
        "improved": score2 > score1,
    }


@router.get("/{version}")
async def get_version(
    version: int,
    session_id: str | None = None,
    x_user_id: str = Header(None),
) -> dict[str, Any]:
    """Get specific version details."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")
    if not session_id:
        raise HTTPException(status_code=400, detail="session_id is required")
    versions = _load_versions(x_user_id, session_id)

    for v in versions:
        if v.get("version") == version:
            return v

    raise HTTPException(status_code=404, detail=f"Version {version} not found")


@router.get("/sessions/{session_id}/iterations/{iteration_idx}/pareto")
async def get_pareto_analysis(
    session_id: str,
    iteration_idx: str,
    x_user_id: str = Header(alias="X-User-Id"),
) -> ParetoFrontResponse:
    """Get Pareto front analysis for an iteration's candidates."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    # Load the latest state snapshot
    latest = _load_latest_version(x_user_id, session_id)

    if not latest:
        return ParetoFrontResponse(pareto_front=[], primary=None, dominated=[])

    # Extract the specified iteration
    iteration = latest.get("iteration") or {}
    if not isinstance(iteration, dict):
        return ParetoFrontResponse(pareto_front=[], primary=None, dominated=[])

    iter_block = iteration.get(iteration_idx)
    if not iter_block or not isinstance(iter_block, dict):
        raise HTTPException(
            status_code=404,
            detail=f"Iteration {iteration_idx} not found in session {session_id}",
        )

    # Extract candidates from the iteration
    candidates = iter_block.get("candidates")
    if not candidates or not isinstance(candidates, dict):
        return ParetoFrontResponse(pareto_front=[], primary=None, dominated=[])

    # Perform Pareto analysis
    result = select_pareto_candidates(candidates)

    return ParetoFrontResponse(
        pareto_front=result["pareto_front"],
        primary=result["primary"],
        dominated=result["dominated"],
    )
