"""Dataset CLI commands."""

from __future__ import annotations

import json
import logging
from typing import Any

logger = logging.getLogger(__name__)


def run_dataset_cmd(group: str, command: str, args: Any) -> None:
    """Handle dataset commands."""
    if command == "create":
        _create(args)
    elif command == "list":
        _list(args)
    elif command == "inspect":
        _inspect(args)
    elif command == "compare":
        _compare(args)
    elif command == "regression":
        _regression(args)
    elif command == "export":
        _export(args)
    else:
        print(f"\n  Unknown command: {command}\n")


def _create(args: Any) -> None:
    """Create a dataset."""
    from .api import BanditoClient
    from .retrieval import get_retrieval

    client = BanditoClient()

    # First, fetch traces matching the criteria
    retrieval = get_retrieval(args.project)
    traces = retrieval.traces(limit=args.limit)

    # Filter by scores if specified
    evals = retrieval.evals()
    filtered_traces = []

    for trace in traces:
        trace_id = trace.get("trace_id")
        eval_data = evals.get(trace_id, {})

        # Filter by human score
        if hasattr(args, "human_score_eq") and args.human_score_eq is not None:
            if eval_data.get("human_score") != args.human_score_eq:
                continue

        # Filter by judge score
        if hasattr(args, "judge_score_eq") and args.judge_score_eq is not None:
            if eval_data.get("judge_score") != args.judge_score_eq:
                continue

        # Filter by judge score < threshold
        if hasattr(args, "judge_score_lt") and args.judge_score_lt is not None:
            score = eval_data.get("judge_score")
            if score is None or score >= args.judge_score_lt:
                continue

        filtered_traces.append(trace)

    print(f"\n  Found {len(filtered_traces)} traces matching criteria")
    print(f"  Creating dataset '{args.name}'...")

    # Create the dataset
    dataset = client.create_dataset(
        name=args.name,
        project=args.project,
        description=args.description,
    )

    print(f"  Created dataset {dataset['id']}")

    # Add items to the dataset
    failed = 0
    for trace in filtered_traces:
        try:
            client.add_dataset_item(
                dataset_id=dataset["id"],
                trace_id=trace["trace_id"],
                source="manual",
                metadata={"reason": "filtered"},
            )
        except Exception as e:
            failed += 1

    if failed:
        print(f"  Added {len(filtered_traces) - failed} items ({failed} failed)")
        print(f"  Warning: {failed} items failed to add. Dataset may be incomplete.")
    else:
        print(f"  Added {len(filtered_traces)} items to dataset")


def _list(args: Any) -> None:
    """List datasets for a project."""
    from .api import BanditoClient

    client = BanditoClient()
    datasets = client.list_datasets(project=args.project)

    if not datasets:
        print(f"\n  No datasets found for project '{args.project}'\n")
        return

    print(f"\n  Datasets for {args.project}:\n")
    print(f"  {'Name':<30} {'Items':<10} {'Created':<20}")
    print(f"  {'-'*30} {'-'*10} {'-'*20}")

    for ds in datasets:
        created = ds.get("created_at", "")
        if created:
            created = created[:19].replace("T", " ")
        print(f"  {ds['name']:<30} {ds.get('item_count', 0):<10} {created:<20}")


def _inspect(args: Any) -> None:
    """Inspect a dataset."""
    from .api import BanditoClient

    client = BanditoClient()

    # Resolve dataset ID from name
    datasets = client.list_datasets(project=args.project)
    dataset = None
    for ds in datasets:
        if ds["name"] == args.name or str(ds["id"]) == args.name:
            dataset = ds
            break

    if dataset is None:
        print(f"\n  Dataset '{args.name}' not found\n")
        return

    # Get full dataset with items
    full = client.get_dataset(dataset["id"])
    items = full.get("items", [])[:args.limit]

    print(f"\n  Dataset: {full['name']} (ID: {full['id']})")
    print(f"  Project: {full['project']}")
    print(f"  Items: {len(full.get('items', []))}")
    if full.get("description"):
        print(f"  Description: {full['description']}")

    if items:
        print(f"\n  First {len(items)} items:\n")
        print(f"  {'Trace ID':<40} {'Source':<12} {'Metadata'}")
        print(f"  {'-'*40} {'-'*12} {'-'*30}")
        for item in items:
            meta = json.dumps(item.get("metadata", {}))[:30]
            print(f"  {item['trace_id']:<40} {item['source']:<12} {meta}")


def _compare(args: Any) -> None:
    """Compare dataset vs production."""
    from .api import BanditoClient
    from .retrieval import get_retrieval

    client = BanditoClient()

    # Get dataset
    datasets = client.list_datasets(project=args.project)
    dataset = None
    for ds in datasets:
        if ds["name"] == args.name or str(ds["id"]) == args.name:
            dataset = ds
            break

    if dataset is None:
        print(f"\n  Dataset '{args.name}' not found\n")
        return

    full = client.get_dataset(dataset["id"])
    dataset_item_count = len(full.get("items", []))

    # Get production stats
    retrieval = get_retrieval(args.project)
    production_traces = retrieval.traces(limit=10000)
    production_evals = retrieval.evals()

    # Calculate stats for dataset
    dataset_eval_scores = []
    for item in full.get("items", []):
        eval_data = production_evals.get(item["trace_id"], {})
        score = eval_data.get("judge_score")
        if score is not None:
            if score not in (0, 1):
                logger.warning(f"Skipping invalid judge_score {score} for {item['trace_id']}")
                continue
            dataset_eval_scores.append(score)

    # Calculate stats for production
    prod_eval_scores = []
    for trace in production_traces:
        eval_data = production_evals.get(trace.get("trace_id"), {})
        score = eval_data.get("judge_score")
        if score is not None:
            if score not in (0, 1):
                logger.warning(f"Skipping invalid judge_score {score} for {trace.get('trace_id')}")
                continue
            prod_eval_scores.append(score)

    # Calculate pass rates
    dataset_pass_rate = sum(dataset_eval_scores) / len(dataset_eval_scores) * 100 if dataset_eval_scores else 0
    prod_pass_rate = sum(prod_eval_scores) / len(prod_eval_scores) * 100 if prod_eval_scores else 0

    print(f"\n  Dataset: {full['name']} ({dataset_item_count} traces)")
    print(f"  Production: {len(production_traces)} traces (last 14 days)")
    print()
    print(f"  {'Metric':<20} {'Dataset':<15} {'Production':<15} {'Diff':<10}")
    print(f"  {'-'*20} {'-'*15} {'-'*15} {'-'*10}")
    print(f"  {'Judge pass rate':<20} {dataset_pass_rate:>6.1f}%      {prod_pass_rate:>6.1f}%      {prod_pass_rate - dataset_pass_rate:+.1f}%")


def _regression(args: Any) -> None:
    """Run diagnostics on a dataset."""
    from .api import BanditoClient
    from .diagnostics_impl import run_diagnostics
    from .retrieval import get_retrieval

    client = BanditoClient()

    # Get dataset
    datasets = client.list_datasets(project=args.project)
    dataset = None
    for ds in datasets:
        if ds["name"] == args.name or str(ds["id"]) == args.name:
            dataset = ds
            break

    if dataset is None:
        print(f"\n  Dataset '{args.name}' not found\n")
        return

    full = client.get_dataset(dataset["id"])
    item_count = len(full.get("items", []))

    print(f"\n  Running diagnostics on dataset {full['name']} ({item_count} traces)...")

    # Get traces for the dataset items
    retrieval = get_retrieval(args.project)
    traces = []
    for item in full.get("items", []):
        trace = retrieval.get_trace(item["trace_id"])
        if trace:
            traces.append(trace)

    if not traces:
        print(f"  No traces found for dataset items\n")
        return

    # Run diagnostics
    evals = retrieval.evals()
    result = run_diagnostics(project=args.project, traces=traces, evaluations=evals)

    print(f"\n  Success rate:    {result.success_rate:.0%}")
    if result.judge_pass_rate:
        print(f"  Judge pass rate: {result.judge_pass_rate:.0%}")
    else:
        print(f"  Judge pass rate: N/A")
    print(f"  Eval coverage:   {result.eval_coverage_pct:.0f}%")

    if result.attention:
        print(f"\n  Attention:")
        for item in result.attention:
            status_icon = "x" if item.status == "critical" else "!"
            print(f"    {status_icon} {item.label}: {item.detail}")


def _export(args: Any) -> None:
    """Export dataset."""
    from .api import BanditoClient
    from .retrieval import get_retrieval

    client = BanditoClient()

    # Get dataset
    datasets = client.list_datasets(project=args.project)
    dataset = None
    for ds in datasets:
        if ds["name"] == args.name or str(ds["id"]) == args.name:
            dataset = ds
            break

    if dataset is None:
        print(f"\n  Dataset '{args.name}' not found\n")
        return

    full = client.get_dataset(dataset["id"])
    items = full.get("items", [])

    # Get traces
    retrieval = get_retrieval(args.project)

    if args.format == "jsonl":
        _export_jsonl(items, retrieval, args.output)
    elif args.format == "pytest":
        _export_pytest(items, retrieval, args.output, full["name"])
    elif args.format == "pairs":
        _export_pairs(items, retrieval, args.output)


def _export_jsonl(items: list[dict], retrieval: Any, output: str) -> None:
    """Export to JSONL format."""
    with open(output, "w") as f:
        for item in items:
            trace = retrieval.get_trace(item["trace_id"])
            if trace:
                record = {
                    "trace_id": item["trace_id"],
                    "input": trace.get("input"),
                    "output": trace.get("output"),
                    "metadata": item.get("metadata"),
                    "source": item.get("source"),
                }
                f.write(json.dumps(record) + "\n")

    print(f"\n  Exported {len(items)} items to {output}")


def _export_pytest(items: list[dict], retrieval: Any, output: str, dataset_name: str) -> None:
    """Export to pytest format."""
    cases = []
    for item in items:
        trace = retrieval.get_trace(item["trace_id"])
        if trace:
            cases.append({
                "trace_id": item["trace_id"],
                "input": trace.get("input"),
                "output": trace.get("output"),
                "metadata": item.get("metadata"),
            })

    with open(output, "w") as f:
        f.write(f"# Generated by Bandito - dataset: {dataset_name}\n")
        f.write("import pytest\n\n")
        f.write("CASES = ")
        f.write(json.dumps(cases, indent=2))
        f.write("\n\n")
        f.write("@pytest.mark.parametrize(\"case\", CASES)\n")
        f.write("def test_regression(case):\n")
        f.write("    \"\"\"Regression test for dataset.\"\"\"\n")
        f.write("    # Run your agent on the input\n")
        f.write("    # assert quality_check(result, case)\n")
        f.write("    pass\n")

    print(f"\n  Exported {len(cases)} cases to {output}")


def _export_pairs(items: list[dict], retrieval: Any, output: str) -> None:
    """Export to prompt-response pairs format."""
    with open(output, "w") as f:
        for item in items:
            trace = retrieval.get_trace(item["trace_id"])
            if trace:
                f.write(f"## INPUT\n{trace.get('input', '')}\n\n")
                f.write(f"## OUTPUT\n{trace.get('output', '')}\n\n")
                f.write("---\n\n")

    print(f"\n  Exported {len(items)} pairs to {output}")
