from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from fastapi import APIRouter, Header, HTTPException, Query
from fastapi.responses import StreamingResponse

from prompt_auto_lab.api.models import BadCase, BadCaseFilters
from prompt_auto_lab.api.services import user_service

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


def _get_braintrust_output_dir(user_id: str, session_id: str) -> Path:
    """
    Get Braintrust output directory using user-scoped session directory.
    """
    session_dir = user_service.get_user_session_dir(user_id, session_id)
    return session_dir / "braintrust_outputs"


def _load_bad_cases(user_id: str, session_id: str, experiment_id: str | None = None) -> list[dict[str, Any]]:
    """Load bad cases from JSONL files."""
    output_dir = _get_braintrust_output_dir(user_id, session_id)
    bad_cases = []

    if not output_dir.exists():
        return bad_cases

    # If experiment_id specified, look in that directory
    if experiment_id:
        exp_dir = output_dir / f"experiments_{experiment_id}"
        filtered_file = exp_dir / "event_details_filtered.jsonl"
        print(f"Loading bad cases from {filtered_file}")
        if filtered_file.exists():
            with open(filtered_file, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        case = json.loads(line)
                        case["experiment_id"] = experiment_id
                        bad_cases.append(case)
    else:
        # Load from all experiment directories
        for exp_dir in output_dir.glob("experiments_*"):
            if exp_dir.is_dir():
                exp_id = exp_dir.name.replace("experiments_", "")
                filtered_file = exp_dir / "event_details_filtered.jsonl"
                if filtered_file.exists():
                    with open(filtered_file, "r", encoding="utf-8") as f:
                        for line in f:
                            line = line.strip()
                            if line:
                                case = json.loads(line)
                                case["experiment_id"] = exp_id
                                bad_cases.append(case)

    return bad_cases


def _filter_bad_cases(
    cases: list[dict[str, Any]], filters: BadCaseFilters
) -> list[dict[str, Any]]:
    """Apply filters to bad cases."""
    filtered = cases

    if filters.experiment_id:
        filtered = [c for c in filtered if c.get("experiment_id") == filters.experiment_id]

    if filters.min_score is not None:
        filtered = [c for c in filtered if c.get("average_score", 0) >= filters.min_score]

    if filters.max_score is not None:
        filtered = [c for c in filtered if c.get("average_score", 1) <= filters.max_score]

    if filters.dimensions:
        # Filter cases that have low scores in specified dimensions
        def has_low_dimension(case: dict[str, Any]) -> bool:
            scores = case.get("merged_scores", {})
            for dim in filters.dimensions:  # type: ignore
                if dim in scores:
                    return True
            return False
        filtered = [c for c in filtered if has_low_dimension(c)]

    if filters.keyword:
        keyword_lower = filters.keyword.lower()
        def matches_keyword(case: dict[str, Any]) -> bool:
            response = str(case.get("test_model_response", "")).lower()
            reasoning = str(case.get("judge_reasoning", "")).lower()
            return keyword_lower in response or keyword_lower in reasoning
        filtered = [c for c in filtered if matches_keyword(c)]

    return filtered


def _enrich_bad_case(case: dict[str, Any]) -> dict[str, Any]:
    """Enrich bad case with computed fields."""
    scores = case.get("merged_scores", {})

    # Compute low_dimensions (scores below 0.5 threshold)
    threshold = 0.5
    low_dims = [k for k, v in scores.items() if isinstance(v, (int, float)) and v < threshold]

    return {
        "root_span_id": case.get("root_span_id", ""),
        "experiment_id": case.get("experiment_id", ""),
        "test_model_response": case.get("test_model_response", ""),
        "judge_response": case.get("judge_response", ""),
        "judge_reasoning": case.get("judge_reasoning", ""),
        "merged_scores": scores,
        "average_score": case.get("average_score", 0),
        "low_dimensions": low_dims,
    }


@router.get("")
async def list_bad_cases(
    experiment_id: str | None = Query(None, description="Filter by experiment ID"),
    min_score: float | None = Query(None, ge=0, le=1, description="Minimum average score"),
    max_score: float | None = Query(None, ge=0, le=1, description="Maximum average score"),
    dimensions: str | None = Query(None, description="Comma-separated dimension names"),
    keyword: str | None = Query(None, description="Search keyword"),
    limit: int = Query(100, ge=1, le=1000, description="Maximum results"),
    offset: int = Query(0, ge=0, description="Results offset"),
    session_id: str | None = Query(None, description="Session ID"),
    x_user_id: str = Header(None),
) -> dict[str, Any]:
    """List bad cases with optional filtering."""
    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")

    # Load all cases
    print(f"Listing bad cases for experiment_id={experiment_id}, session_id={session_id}, user_id={x_user_id}")
    cases = _load_bad_cases(x_user_id, session_id, experiment_id)

    # Build filters
    filters = BadCaseFilters(
        experiment_id=experiment_id,
        min_score=min_score,
        max_score=max_score,
        dimensions=dimensions.split(",") if dimensions else None,
        keyword=keyword,
    )

    # Apply filters
    filtered = _filter_bad_cases(cases, filters)

    # Enrich with computed fields
    enriched = [_enrich_bad_case(c) for c in filtered]

    # Pagination
    total = len(enriched)
    paginated = enriched[offset:offset + limit]

    # Compute dimension statistics
    dimension_counts: dict[str, int] = {}
    for case in enriched:
        for dim in case.get("low_dimensions", []):
            dimension_counts[dim] = dimension_counts.get(dim, 0) + 1

    return {
        "cases": paginated,
        "total": total,
        "limit": limit,
        "offset": offset,
        "dimension_stats": dimension_counts,
    }


@router.get("/{experiment_id}")
async def get_bad_cases_by_experiment(
    experiment_id: str,
    session_id: str | None = Query(None),
    x_user_id: str = Header(None),
) -> list[dict[str, Any]]:
    """Get all bad cases for a specific experiment."""
    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")
    cases = _load_bad_cases(x_user_id, session_id, experiment_id)
    return [_enrich_bad_case(c) for c in cases]


@router.get("/export")
async def export_bad_cases(
    experiment_id: str | None = Query(None),
    format: str = Query("json", regex="^(json|csv)$"),
    min_score: float | None = Query(None),
    max_score: float | None = Query(None),
    session_id: str | None = Query(None),
    x_user_id: str = Header(None),
) -> StreamingResponse:
    """Export bad cases as JSON or CSV."""
    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")
    cases = _load_bad_cases(x_user_id, session_id, experiment_id)

    filters = BadCaseFilters(
        experiment_id=experiment_id,
        min_score=min_score,
        max_score=max_score,
    )
    filtered = _filter_bad_cases(cases, filters)
    enriched = [_enrich_bad_case(c) for c in filtered]

    if format == "csv":
        # Generate CSV
        import csv
        import io

        output = io.StringIO()
        if enriched:
            writer = csv.DictWriter(output, fieldnames=enriched[0].keys())
            writer.writeheader()
            for case in enriched:
                # Flatten nested fields for CSV
                flat_case = {
                    **case,
                    "merged_scores": json.dumps(case.get("merged_scores", {})),
                    "low_dimensions": ",".join(case.get("low_dimensions", [])),
                }
                writer.writerow(flat_case)

        return StreamingResponse(
            iter([output.getvalue()]),
            media_type="text/csv",
            headers={"Content-Disposition": "attachment; filename=bad_cases.csv"},
        )
    else:
        # Generate JSON
        return StreamingResponse(
            iter([json.dumps(enriched, indent=2)]),
            media_type="application/json",
            headers={"Content-Disposition": "attachment; filename=bad_cases.json"},
        )


@router.get("/stats/dimensions")
async def get_dimension_stats(
    experiment_id: str | None = None,
    session_id: str | None = Query(None),
    x_user_id: str = Header(None),
) -> dict[str, Any]:
    """Get aggregated dimension failure statistics."""
    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")
    cases = _load_bad_cases(x_user_id, session_id, experiment_id)
    enriched = [_enrich_bad_case(c) for c in cases]

    # Count failures per dimension
    dimension_counts: dict[str, int] = {}
    dimension_avg_scores: dict[str, list[float]] = {}

    for case in enriched:
        scores = case.get("merged_scores", {})
        for dim, score in scores.items():
            if isinstance(score, (int, float)):
                if dim not in dimension_avg_scores:
                    dimension_avg_scores[dim] = []
                dimension_avg_scores[dim].append(score)

                if score < 0.5:
                    dimension_counts[dim] = dimension_counts.get(dim, 0) + 1

    # Compute averages
    dimension_averages = {
        dim: sum(scores) / len(scores)
        for dim, scores in dimension_avg_scores.items()
        if scores
    }

    # Sort by failure count
    sorted_failures = sorted(
        dimension_counts.items(),
        key=lambda x: x[1],
        reverse=True
    )

    return {
        "total_cases": len(enriched),
        "failure_counts": dict(sorted_failures),
        "average_scores": dimension_averages,
        "top_failing_dimensions": [d[0] for d in sorted_failures[:5]],
    }
