"""
Common utilities for OOLONG benchmark test runners.

Based on the OOLONG paper:
- OOLONG-synth: Synthetic ICL-based aggregation tasks
- OOLONG-real: Real D&D transcript aggregation tasks
"""

import ast
import json
import logging
import re
from collections.abc import Callable, Sequence
from datetime import date, datetime
from pathlib import Path
from typing import Any, cast

import dateutil.parser
import pyarrow.parquet as pq
import tiktoken
from typing_extensions import TypedDict

logger = logging.getLogger(__name__)


class SimpleDataset:
    """Simple dataset class that mimics HuggingFace dataset API."""

    data: list[dict[str, Any]]
    column_names: list[str]

    def __init__(self, data: list[dict[str, Any]]):
        """Initialize dataset with list of examples.

        Args:
            data: List of example dictionaries
        """
        self.data = data
        self.column_names = list(data[0].keys()) if data else []

    def __len__(self) -> int:
        """Return number of examples."""
        return len(self.data)

    def __getitem__(self, idx: int) -> dict[str, Any]:
        """Get example by index."""
        return self.data[idx]

    def filter(self, function: Callable[[dict[str, Any]], bool]) -> "SimpleDataset":
        """Filter dataset using a function.

        Args:
            function: Filter function

        Returns:
            Filtered dataset
        """
        filtered_data = [item for item in self.data if function(item)]
        return SimpleDataset(filtered_data)

    def select(self, indices: Sequence[int]) -> "SimpleDataset":
        """Select examples by indices.

        Args:
            indices: List of indices to select

        Returns:
            Dataset with selected examples
        """
        selected_data = [self.data[i] for i in indices]
        return SimpleDataset(selected_data)


class BaseQueryResult(TypedDict):
    """Base type definition for query execution results."""

    question: str
    expected_answer: str
    actual_response: str
    score: float
    context_length_tokens: int


class BaseTestResult(TypedDict):
    """Base type definition for test execution results."""

    question_id: str
    task_group: str
    dataset: str
    passed: bool
    score: float
    error: str | None
    start_time: float
    end_time: float
    duration_seconds: float
    output_lines: list[str]


def format_duration(total_seconds: float) -> str:
    """Format a duration in seconds into a human-readable string.

    If the duration is at least one minute, this returns a string in the
    form "XmYYs" with zero-padded seconds. Otherwise, it returns the
    duration in seconds with two decimal places, e.g., "12.34s".

    Args:
        total_seconds: The duration in seconds.

    Returns:
        A formatted duration string.
    """
    minutes = int(total_seconds // 60)
    if minutes > 0:
        seconds_rounded = int(round(total_seconds - minutes * 60))
        if seconds_rounded == 60:
            minutes += 1
            seconds_rounded = 0
        return f"{minutes}m{seconds_rounded:02d}s"
    return f"{total_seconds:.2f}s"


def calculate_context_length(text: str) -> int:
    """Calculate token count for context text using tiktoken.

    Args:
        text: Context text to count tokens for

    Returns:
        Number of tokens
    """
    try:
        tokenizer = tiktoken.get_encoding("o200k_base")
        return len(
            tokenizer.encode(
                text,
                disallowed_special=(tokenizer.special_tokens_set - {"<|endoftext|>"}),
            )
        )
    except Exception:
        # Fallback to character-based estimate
        return len(text) // 4


def load_oolong_synth_dataset(
    split: str = "test",
    data_dir: str | Path | None = None,
    max_context_len: int | None = None,
    min_context_len: int | None = None,
    max_examples: int | None = None,
    context_window_id: str | None = None,
) -> SimpleDataset:
    """Load the OOLONG-synth dataset from filesystem.

    Args:
        split: Dataset split to load (default: "test")
        data_dir: Path to the oolong-synth dataset directory (must contain a 'data' subdirectory)
        max_context_len: Maximum context length in tokens
        min_context_len: Minimum context length in tokens (strict >, upstream behavior)
        max_examples: Maximum number of examples to return
        context_window_id: Specific context window ID to filter to

    Returns:
        SimpleDataset object

    Raises:
        ValueError: If data_dir is not provided
        FileNotFoundError: If no parquet files found for the split
    """
    if data_dir is None:
        raise ValueError(
            "data_dir parameter is required. Please provide the path to the oolong-synth dataset."
        )

    dataset_path = Path(data_dir) / "data"
    if not dataset_path.exists():
        raise FileNotFoundError(
            f"Expected synth dataset directory at {dataset_path} (layout: data/*.parquet)"
        )

    # Find all parquet files for the given split
    parquet_files = sorted(dataset_path.glob(f"{split}-*.parquet"))

    if not parquet_files:
        raise FileNotFoundError(f"No {split} parquet files found in {dataset_path}")

    # Stream rows from parquet shards in batches so we can filter and stop early
    # without materializing the full split in memory first.
    all_data: list[dict[str, Any]] = []
    for parquet_file in parquet_files:
        parquet = pq.ParquetFile(parquet_file)
        for batch in parquet.iter_batches():  # pyright: ignore[reportUnknownVariableType]
            rows = cast(
                list[dict[str, Any]],
                batch.to_pylist(),  # pyright: ignore[reportUnknownMemberType]
            )
            for row in rows:
                if (
                    context_window_id is not None
                    and str(row.get("context_window_id")) != context_window_id
                ):
                    continue

                context_len = row.get("context_len")
                if not isinstance(context_len, int):
                    context_len = calculate_context_length(
                        str(row.get("context_window_text", ""))
                    )

                if max_context_len is not None and context_len > max_context_len:
                    continue

                # Keep strict greater-than for consistency with existing behavior.
                if min_context_len is not None and context_len <= min_context_len:
                    continue

                all_data.append(row)
                if (
                    max_examples is not None
                    and max_examples > 0
                    and len(all_data) >= max_examples
                ):
                    return SimpleDataset(all_data)

    return SimpleDataset(all_data)


def load_oolong_real_dataset(
    split: str = "test",
    data_dir: str | Path | None = None,
    max_context_len: int | None = None,
    min_context_len: int | None = None,
    max_examples: int | None = None,
    context_window_id: str | None = None,
) -> SimpleDataset:
    """Load the OOLONG-real dataset from filesystem.

    Args:
        split: Dataset split to load (default: "test")
        data_dir: Path to the oolong-real dataset directory (must contain a 'dnd' subdirectory)
        max_context_len: Maximum context length in tokens
        min_context_len: Minimum context length in tokens (strict >, upstream behavior)
        max_examples: Maximum number of examples to return
        context_window_id: Specific context window ID to filter to

    Returns:
        SimpleDataset object

    Raises:
        ValueError: If data_dir is not provided
        FileNotFoundError: If JSONL file not found
    """
    if data_dir is None:
        raise ValueError(
            "data_dir parameter is required. Please provide the path to the oolong-real dataset."
        )

    dataset_path = Path(data_dir) / "dnd"
    if not dataset_path.exists():
        raise FileNotFoundError(
            f"Expected real dataset directory at {dataset_path} (layout: dnd/*.jsonl)"
        )
    jsonl_file = dataset_path / f"{split}.jsonl"

    if not jsonl_file.exists():
        raise FileNotFoundError(f"JSONL file not found: {jsonl_file}")

    # Load JSONL file
    data: list[dict[str, Any]] = []
    with open(jsonl_file) as f:
        for line in f:
            if line.strip():
                row_raw = json.loads(line)
                if not isinstance(row_raw, dict):
                    continue
                row = cast(dict[str, Any], row_raw)

                if (
                    context_window_id is not None
                    and str(row.get("context_window_id")) != context_window_id
                ):
                    continue

                context_len = row.get("context_len")
                if not isinstance(context_len, int):
                    context_len = calculate_context_length(
                        str(row.get("context_window_text", ""))
                    )

                if max_context_len is not None and context_len > max_context_len:
                    continue

                # Keep strict greater-than for consistency with existing behavior.
                if min_context_len is not None and context_len <= min_context_len:
                    continue

                data.append(row)
                if (
                    max_examples is not None
                    and max_examples > 0
                    and len(data) >= max_examples
                ):
                    return SimpleDataset(data)

    return SimpleDataset(data)


def parse_synth_context_messages(context_text: str) -> list[dict[str, Any]]:
    """Parse OOLONG-synth context text into individual messages.

    Context format:
    Date: YYYY-MM-DD || User: user_XYZ || Instance: <text> [label]

    Args:
        context_text: Raw context window text

    Returns:
        List of message dictionaries with content and metadata
    """
    messages: list[dict[str, Any]] = []

    # Split by lines and parse each entry
    lines = context_text.strip().split("\n")
    for line in lines:
        if not line.strip():
            continue

        # Parse: Date: ... || User: ... || Instance: ... || Label: ...
        try:
            parts = line.split(" || ")
            if len(parts) < 3:
                logger.warning(f"Skipping malformed line: {line}")
                continue

            date_part = parts[0].replace("Date: ", "").strip()
            user_part = parts[1].replace("User: ", "").strip()
            instance_part = parts[2].replace("Instance: ", "").strip()

            # Check if there's a 4th part with label
            label = None
            if len(parts) >= 4:
                label_part = parts[3].replace("Label: ", "").strip()
                label = label_part if label_part else None

            # Include label in content so deriver can observe it.
            content = f"{instance_part} [Label: {label}]" if label else instance_part

            msg: dict[str, Any] = {
                "content": content,
                "metadata": {
                    "date": date_part,
                    "user_id": user_part,
                    "label": label,
                },
            }
            messages.append(msg)
        except Exception as e:
            logger.warning(f"Error parsing line: {line}. Error: {e}")
            continue

    return messages


def parse_real_context_messages(context_text: str) -> list[dict[str, Any]]:
    """Parse OOLONG-real D&D transcript into individual messages.

    Context format:
    Speaker: dialogue text
    [multiple lines]

    Args:
        context_text: Raw D&D transcript text

    Returns:
        List of message dictionaries with content and metadata
    """
    messages: list[dict[str, Any]] = []

    # Split by speaker turns (format: "SPEAKER: text")
    lines = context_text.strip().split("\n")
    current_speaker = None
    current_content = []

    for line in lines:
        if not line.strip():
            continue

        # Check if this is a new speaker turn (must start with speaker label)
        # Speaker labels are typically single words or use underscores/hyphens (no spaces)
        speaker_match = re.match(r"^\s*([A-Za-z0-9_-]+):", line)
        if speaker_match:
            # Save previous message if exists
            if current_speaker and current_content:
                # Include speaker in content for deriver visibility
                content_text = " ".join(current_content)
                content = f"[Speaker: {current_speaker}] {content_text}"
                prev_msg: dict[str, Any] = {
                    "content": content,
                    "metadata": {
                        "speaker": current_speaker,
                    },
                }
                messages.append(prev_msg)

            # Parse new speaker from regex match
            current_speaker = speaker_match.group(1).strip()
            # Extract content after the colon
            content_after_colon = line[speaker_match.end() :].strip()
            current_content = [content_after_colon] if content_after_colon else []
        else:
            # Continuation of current speaker's dialogue
            current_content.append(line.strip())

    # Save last message
    if current_speaker and current_content:
        content_text = " ".join(current_content)
        content = f"[Speaker: {current_speaker}] {content_text}"
        last_msg: dict[str, Any] = {
            "content": content,
            "metadata": {
                "speaker": current_speaker,
            },
        }
        messages.append(last_msg)

    return messages


def parse_synth_answer(answer_str: str) -> Any:
    """Parse OOLONG-synth answer string.

    Answers can be:
    - Strings (labels, comparisons, dates)
    - Numbers (counts)
    - Dates (datetime objects)

    Args:
        answer_str: Raw answer string from dataset

    Returns:
        Parsed answer value
    """
    # Handle datetime answers
    if "datetime" in answer_str:
        try:
            # Format: [datetime.date(2023, 5, 15)]
            match = re.search(r"datetime\.date\((\d+),\s*(\d+),\s*(\d+)\)", answer_str)
            if match:
                year, month, day = map(int, match.groups())
                return datetime(year, month, day).date()
        except Exception as e:
            logger.warning(f"Error parsing datetime answer: {answer_str}. Error: {e}")
            return answer_str

    # Try literal eval for lists/primitives
    try:
        parsed = ast.literal_eval(answer_str)
        # If it's a list with one element, return that element
        if isinstance(parsed, list) and len(parsed) == 1:  # pyright: ignore[reportUnknownArgumentType]
            return parsed[0]  # pyright: ignore[reportUnknownVariableType]
        return parsed  # pyright: ignore[reportUnknownVariableType]
    except (ValueError, SyntaxError):
        # Return as-is if can't parse
        return answer_str


def parse_real_answer(answer_str: str) -> int | str | list[str]:
    """Parse OOLONG-real answer string.

    Answers can be:
    - Integers (counts)
    - Strings (spell names, roll types)
    - Lists (comma-separated spells)

    Args:
        answer_str: Raw answer string from dataset

    Returns:
        Parsed answer value
    """
    # Try to convert to int first
    try:
        return int(answer_str)
    except ValueError:
        pass

    # Check if it contains commas (list case)
    if "," in answer_str:
        return [item.strip() for item in answer_str.split(",") if item.strip()]

    # Otherwise return as string
    return answer_str


def score_synth_response(
    gold_answer: Any, model_answer_str: str, answer_type: str
) -> float:
    """Score a response for OOLONG-synth following the paper's scoring rubric.

    Scoring:
    - Exact match: 1.0
    - Numeric: 0.75^|y-ŷ| (partial credit)
    - Other: 0.0

    Args:
        gold_answer: Expected answer
        model_answer_str: Model's response string
        answer_type: Type of answer (ANSWER_TYPE.NUMERIC, ANSWER_TYPE.DATE, etc.)

    Returns:
        Score between 0.0 and 1.0
    """
    # Try to extract answer from response
    # Look for common patterns at start of line or end of response
    model_answer = model_answer_str.strip()

    # Try to extract from common formats (use MULTILINE to match at line start, and $ to find at line end)
    # Look for patterns at the start of a line (final answer format)
    for pattern in [
        r"^[Aa]nswer:\s*(.+)$",  # "Answer: X" on its own line
        r"^[Ll]abel:\s*(.+)$",  # "Label: X" on its own line
        r"\n[Ll]abel:\s*(.+)$",  # "Label: X" after a newline (final answer)
        r"[Aa]nswer:\s*(.+)$",  # "Answer: X" at end of response
        r"[Uu]ser:\s*(.+)$",  # "User: X" at end
        r"[Dd]ate:\s*(.+)$",  # "Date: X" at end
    ]:
        match = re.search(pattern, model_answer, re.MULTILINE)
        if match:
            model_answer = match.group(1).strip()
            break

    # Remove formatting like ** or []
    model_answer = re.sub(r"[\*\[\]]", "", model_answer)

    # If still long, try last significant token
    if len(model_answer) > 50:
        model_answer = model_answer.split()[-1]

    # Exact string match
    if str(model_answer).lower() == str(gold_answer).lower():
        return 1.0

    # Check for comparison answers
    if (
        "more common" in model_answer.lower()
        and "more common" in str(gold_answer).lower()
    ):
        return 1.0
    if (
        "less common" in model_answer.lower()
        and "less common" in str(gold_answer).lower()
    ):
        return 1.0
    if (
        "same frequency" in model_answer.lower()
        and "same frequency" in str(gold_answer).lower()
    ):
        return 1.0

    # Numeric partial credit
    if "NUMERIC" in answer_type.upper():
        try:
            model_num = float(re.sub(r"[^\d.-]", "", model_answer))
            gold_num = float(gold_answer)
            return 0.75 ** abs(gold_num - model_num)
        except (ValueError, TypeError):
            return 0.0

    # Date matching
    if "DATE" in answer_type.upper():
        try:
            model_date = dateutil.parser.parse(model_answer)
            if isinstance(gold_answer, datetime):
                return 1.0 if model_date.date() == gold_answer.date() else 0.0
            elif isinstance(gold_answer, date):
                return 1.0 if model_date.date() == gold_answer else 0.0
            return 0.0
        except (ValueError, TypeError):
            return 0.0

    return 0.0


def score_real_response(
    gold_answer: int | str | list[str], model_answer_str: str
) -> float:
    """Score a response for OOLONG-real following the paper's scoring rubric.

    Scoring:
    - Integer: 0.75^|y-ŷ| (partial credit)
    - String: exact match (1.0 or 0.0)
    - List: set overlap / |gold| (Jaccard-style)

    Args:
        gold_answer: Expected answer
        model_answer_str: Model's response string

    Returns:
        Score between 0.0 and 1.0
    """
    # Extract answer from \boxed{} format if present
    match = re.search(r"\\boxed\{\\text\{([^}]*)\}\}", model_answer_str) or re.search(
        r"\\boxed[\{]+([^}]*)[\}]+", model_answer_str
    )

    if match:
        model_answer_str = match.group(1)

    # Parse model answer
    try:
        model_answer: int | str | list[str] = int(model_answer_str)
    except ValueError:
        if "," in model_answer_str:
            model_answer = [
                item.strip() for item in model_answer_str.split(",") if item.strip()
            ]
        else:
            model_answer = model_answer_str.strip()

    # Score based on type
    if isinstance(gold_answer, int) and isinstance(model_answer, int):
        return 0.75 ** abs(gold_answer - model_answer)
    elif isinstance(gold_answer, str) and isinstance(model_answer, str):
        return 1.0 if gold_answer.lower() == model_answer.lower() else 0.0
    elif isinstance(gold_answer, list) and isinstance(model_answer, list):
        overlap = set(gold_answer) & set(model_answer)
        return len(overlap) / len(gold_answer) if gold_answer else 0.0
    else:
        return 0.0


def filter_dataset(
    dataset: SimpleDataset,
    max_context_len: int | None = None,
    min_context_len: int | None = None,
    max_examples: int | None = None,
    context_window_id: str | None = None,
) -> SimpleDataset:
    """Filter dataset by context length and example count.

    Args:
        dataset: HuggingFace dataset
        max_context_len: Maximum context length in tokens
        min_context_len: Minimum context length in tokens
        max_examples: Maximum number of examples to return
        context_window_id: Specific context window ID to filter to

    Returns:
        Filtered dataset
    """
    if context_window_id is not None:
        dataset = dataset.filter(lambda x: x["context_window_id"] == context_window_id)

    if max_context_len is not None:
        dataset = dataset.filter(
            lambda x: (
                x.get(
                    "context_len",
                    calculate_context_length(str(x.get("context_window_text", ""))),
                )
                <= max_context_len
            )
        )

    if min_context_len is not None:
        dataset = dataset.filter(
            lambda x: (
                x.get(
                    "context_len",
                    calculate_context_length(str(x.get("context_window_text", ""))),
                )
                > min_context_len
            )
        )

    if max_examples is not None and max_examples > 0:
        # Get a slice of the dataset
        dataset = dataset.select(range(min(max_examples, len(dataset))))

    return dataset


def calculate_task_statistics(
    results: Sequence[Any],
) -> dict[str, dict[str, int | float]]:
    """Calculate score statistics grouped by task type.

    Args:
        results: List of test results

    Returns:
        Dictionary mapping task type to statistics
    """
    task_stats: dict[str, dict[str, int | float]] = {}

    for result in results:
        task_key = result.get("task_group", "unknown")
        if task_key not in task_stats:
            task_stats[task_key] = {"total": 0, "total_score": 0.0, "perfect_scores": 0}

        task_stats[task_key]["total"] += 1
        score = result.get("score", 0.0)
        task_stats[task_key]["total_score"] += score
        if score >= 0.99:  # Consider >= 0.99 as perfect to account for floating point
            task_stats[task_key]["perfect_scores"] += 1

    # Calculate averages
    for task_key in task_stats:
        stats = task_stats[task_key]
        total = stats["total"]
        stats["average_score"] = stats["total_score"] / total if total > 0 else 0.0
        stats["perfect_score_rate"] = (
            (stats["perfect_scores"] / total) * 100 if total > 0 else 0.0
        )

    return task_stats


def calculate_timing_statistics(
    results: Sequence[Any], total_elapsed_seconds: float
) -> dict[str, Any]:
    """Calculate timing statistics from test results.

    Args:
        results: List of test results
        total_elapsed_seconds: Total elapsed time for all tests

    Returns:
        Dictionary of timing statistics
    """
    durations = [r["duration_seconds"] for r in results]
    return {
        "total_duration_seconds": total_elapsed_seconds,
        "individual_test_durations": {
            "min_seconds": min(durations) if durations else 0,
            "max_seconds": max(durations) if durations else 0,
            "mean_seconds": sum(durations) / len(durations) if durations else 0,
            "median_seconds": sorted(durations)[len(durations) // 2]
            if durations
            else 0,
        },
    }


def write_json_summary(summary: dict[str, Any], output_file: Path) -> None:
    """Write a JSON summary to a file.

    Args:
        summary: Summary dictionary to write
        output_file: Path to output file
    """
    output_file.parent.mkdir(parents=True, exist_ok=True)
    with open(output_file, "w") as f:
        json.dump(summary, f, indent=2, default=str)
    print(f"\nJSON summary written to: {output_file}")