"""
babilong_eval.harness — Evaluation runner for BABILong benchmark.

Orchestrates the full evaluation pipeline:
  1. Load BABILong dataset from HuggingFace
  2. For each (task, length, sample):
     a. Apply preprocessing condition (full/hcp/random)
     b. Format prompt per BABILong template
     c. Call LLM with greedy decoding
     d. Score via label matching
  3. Aggregate results

LLM Integration:
  Same pattern as ifeval_benchmark.py: OpenAI-compatible API with
  OPENAI_API_KEY or GITHUB_TOKEN fallback.

BABILong Protocol:
  - temperature=0, max_tokens=20 (greedy, short answers)
  - Scoring via compare_answers() (label matching, not exact)
  - Per-task prompt templates (qa1–qa5)
"""

from __future__ import annotations

import os
import sys
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

from babilong_eval.prompts import format_prompt
from babilong_eval.metrics import compare_answers
from babilong_eval.preprocessor import full_context, hcp_filtered, random_filtered, bm25_filtered, embedding_filtered
from babilong_eval.report import BabilongResult, SampleResult


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

@dataclass
class BabilongConfig:
    """Configuration for BABILong evaluation.

    Attributes:
        tasks: Which QA tasks to evaluate.
        lengths: Which context lengths to test (as BABILong split names).
        n_samples: Number of samples per (task, length) pair.
        retain_ratios: Fractions to keep for filtered conditions.
        conditions: Which preprocessing conditions to run.
        model: LLM model name.
        delay: Seconds between API calls (rate limiting).
        random_seed: Seed for random_filtered reproducibility.
    """
    tasks: Tuple[str, ...] = ("qa1", "qa2", "qa3", "qa4", "qa5")
    lengths: Tuple[str, ...] = ("4k", "8k", "16k", "32k")
    n_samples: int = 50
    retain_ratios: Tuple[float, ...] = (0.1, 0.25, 0.5)
    conditions: Tuple[str, ...] = ("full", "hcp", "random", "bm25", "embedding")
    model: str = "gpt-4o-mini"
    delay: float = 1.0
    random_seed: int = 42


# ---------------------------------------------------------------------------
# LLM Client (matches ifeval_benchmark.py pattern)
# ---------------------------------------------------------------------------

_CLIENT = None
_MODEL = "gpt-4o-mini"


def get_client():
    """Create OpenAI client. Prefers OPENAI_API_KEY, falls back to GitHub Models."""
    from openai import OpenAI

    openai_key = os.environ.get("OPENAI_API_KEY")
    if openai_key:
        print("Using OpenAI API (direct)")
        return OpenAI(api_key=openai_key, max_retries=0, timeout=60.0)

    token = os.environ.get("GITHUB_TOKEN")
    if token:
        print("Using GitHub Models API")
        return OpenAI(
            base_url="https://models.inference.ai.azure.com",
            api_key=token,
        )

    print("ERROR: Set OPENAI_API_KEY or GITHUB_TOKEN environment variable")
    sys.exit(1)


def call_llm(prompt: str, model: str = "gpt-4o-mini") -> str:
    """Call LLM with BABILong-appropriate settings.

    BABILong protocol: greedy decoding, short output.

    Args:
        prompt: The formatted prompt string.
        model: Model name.

    Returns:
        The LLM's response text.
    """
    global _CLIENT, _MODEL
    if _CLIENT is None:
        _CLIENT = get_client()
    _MODEL = model

    import time as _time
    max_retries = 5
    for attempt in range(max_retries):
        try:
            response = _CLIENT.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0,       # Greedy decoding — BABILong protocol
                max_tokens=20,       # Short answers (1-2 words expected)
            )
            return response.choices[0].message.content.strip()
        except KeyboardInterrupt:
            raise
        except Exception as e:
            err_str = str(e)
            if "429" in err_str or "rate" in err_str.lower():
                wait = 2 ** attempt * 5  # 5, 10, 20, 40, 80s
                print(f"  Rate limited, waiting {wait}s (attempt {attempt+1}/{max_retries})...")
                _time.sleep(wait)
                continue
            print(f"  API error: {e}")
            return ""
    print(f"  Max retries exceeded")
    return ""


# ---------------------------------------------------------------------------
# Dataset loading
# ---------------------------------------------------------------------------

def load_babilong_split(task: str, length: str, n_samples: int = 50) -> List[Dict[str, str]]:
    """Load a BABILong dataset split from HuggingFace.

    Args:
        task: Task name (qa1–qa5).
        length: Context length (4k, 8k, 16k, 32k, etc.).
        n_samples: Number of samples to load.

    Returns:
        List of dicts with keys: input, target, question.
    """
    from datasets import load_dataset

    dataset = load_dataset("RMT-team/babilong", length, split=task)

    samples = []
    for i, example in enumerate(dataset):
        if i >= n_samples:
            break
        samples.append({
            "input": example["input"],
            "target": example["target"],
            "question": example["question"],
        })

    return samples


# ---------------------------------------------------------------------------
# Single sample evaluation
# ---------------------------------------------------------------------------

def run_single_sample(
    sample: Dict[str, str],
    task: str,
    condition: str,
    retain_ratio: float = 0.25,
    model: str = "gpt-4o-mini",
    random_seed: int = 42,
) -> SampleResult:
    """Evaluate a single BABILong sample under one condition.

    Args:
        sample: Dict with input, target, question.
        task: Task name (qa1–qa5).
        condition: "full", "hcp", or "random".
        retain_ratio: For filtered conditions.
        model: LLM model name.
        random_seed: For random_filtered reproducibility.

    Returns:
        SampleResult with output, target, and correctness.
    """
    context = sample["input"]
    question = sample["question"]
    target = sample["target"]

    # Apply preprocessing condition
    if condition == "hcp":
        processed_context = hcp_filtered(context, question, task, retain_ratio)
    elif condition == "random":
        processed_context = random_filtered(context, retain_ratio, seed=random_seed)
    elif condition == "bm25":
        processed_context = bm25_filtered(context, question, retain_ratio)
    elif condition == "embedding":
        processed_context = embedding_filtered(context, question, retain_ratio)
    else:  # "full"
        processed_context = full_context(context)

    # Format prompt per BABILong template
    prompt = format_prompt(task, processed_context, question)

    # Call LLM
    output = call_llm(prompt, model=model)

    # Score
    correct = compare_answers(output, target, question, task)

    return SampleResult(
        task=task,
        condition=condition,
        retain_ratio=retain_ratio if condition != "full" else 1.0,
        output=output,
        target=target,
        question=question,
        correct=correct,
        context_len=len(context),
        filtered_len=len(processed_context),
    )


# ---------------------------------------------------------------------------
# Batch evaluation
# ---------------------------------------------------------------------------

def run_task_length(
    task: str,
    length: str,
    config: BabilongConfig,
) -> List[SampleResult]:
    """Run evaluation for one (task, length) pair across all conditions.

    Args:
        task: Task name (qa1–qa5).
        length: Context length (4k, 8k, etc.).
        config: Evaluation configuration.

    Returns:
        List of SampleResults for all conditions and samples.
    """
    print(f"\n{'='*60}")
    print(f"Task: {task} | Length: {length} | Samples: {config.n_samples}")
    print(f"{'='*60}")

    samples = load_babilong_split(task, length, config.n_samples)
    if not samples:
        print(f"  WARNING: No samples loaded for {task}/{length}")
        return []

    results: List[SampleResult] = []

    for condition in config.conditions:
        ratios = config.retain_ratios if condition != "full" else (1.0,)

        for ratio in ratios:
            correct_count = 0
            condition_label = f"{condition}" if condition == "full" else f"{condition}@{ratio}"
            print(f"\n  Condition: {condition_label}")

            for i, sample in enumerate(samples):
                try:
                    result = run_single_sample(
                        sample=sample,
                        task=task,
                        condition=condition,
                        retain_ratio=ratio,
                        model=config.model,
                        random_seed=config.random_seed + i,
                    )
                except KeyboardInterrupt:
                    raise  # allow clean Ctrl+C exit
                except Exception as e:
                    print(f"    ERROR sample {i}: {e}")
                    result = SampleResult(
                        task=task,
                        condition=condition,
                        retain_ratio=ratio if condition != "full" else 1.0,
                        output=f"ERROR: {e}",
                        target=sample.get("target", ""),
                        question=sample.get("question", ""),
                        correct=False,
                        context_len=len(sample.get("input", "")),
                        filtered_len=0,
                    )
                results.append(result)
                correct_count += int(result.correct)

                if (i + 1) % 10 == 0:
                    acc = correct_count / (i + 1)
                    print(f"    [{i+1}/{len(samples)}] accuracy so far: {acc:.1%}")

                time.sleep(config.delay)

            acc = correct_count / len(samples) if samples else 0
            print(f"  -> {condition_label}: {acc:.1%} ({correct_count}/{len(samples)})")

    return results


def run_full_eval(config: BabilongConfig) -> BabilongResult:
    """Run the complete BABILong evaluation.

    Iterates over all (task, length) pairs and all conditions.

    Args:
        config: Evaluation configuration.

    Returns:
        BabilongResult with all sample results.
    """
    all_results: List[SampleResult] = []

    total_calls = (
        len(config.tasks) * len(config.lengths) * config.n_samples
        * (
            ("full" in config.conditions)
            + sum(
                len(config.retain_ratios)
                for c in config.conditions if c != "full"
            )
        )
    )
    est_cost = total_calls * 0.005  # Rough estimate: $0.005/call for gpt-4o-mini

    print(f"\nBABILong Evaluation")
    print(f"  Tasks:      {', '.join(config.tasks)}")
    print(f"  Lengths:    {', '.join(config.lengths)}")
    print(f"  Conditions: {', '.join(config.conditions)}")
    print(f"  Samples:    {config.n_samples} per (task, length)")
    print(f"  Est. API calls: ~{total_calls}")
    print(f"  Est. cost:      ~${est_cost:.2f}")
    print(f"  Model:      {config.model}")

    for task in config.tasks:
        for length in config.lengths:
            task_results = run_task_length(task, length, config)
            all_results.extend(task_results)

    return BabilongResult(
        config=config,
        samples=all_results,
    )
