"""
Common benchmark runner framework based on the LoComo implementation.

This module provides a unified interface for running benchmarks with the same
optimizations as the working LoComo benchmark:
- Batch ingestion for speed
- Parallel question processing with semaphores
- Parallel LLM judging with rate limiting
- Progress tracking with Rich
- Comprehensive metrics collection
"""

import json
import asyncio
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from typing import List, Dict, Any, Optional, Tuple
from pathlib import Path
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
from rich.table import Table
from rich import box
import pydantic

from memora import TemporalSemanticMemory
from openai import AsyncOpenAI

console = Console()


class BenchmarkDataset(ABC):
    """Abstract base class for benchmark datasets."""

    @abstractmethod
    def load(self, path: Path, max_items: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        Load dataset from file.

        Returns:
            List of dataset items
        """
        pass

    @abstractmethod
    def get_item_id(self, item: Dict) -> str:
        """Get unique identifier for an item."""
        pass

    @abstractmethod
    def prepare_sessions_for_ingestion(self, item: Dict) -> List[Dict[str, Any]]:
        """
        Prepare conversation sessions for batch ingestion.

        Returns:
            List of session dicts with keys: 'content', 'context', 'event_date'
        """
        pass

    @abstractmethod
    def get_qa_pairs(self, item: Dict) -> List[Dict[str, Any]]:
        """
        Extract QA pairs from an item.

        Returns:
            List of QA dicts with keys: 'question', 'answer', 'category' (optional)
        """
        pass


class LLMAnswerGenerator(ABC):
    """Abstract base class for LLM-based answer generation."""

    @abstractmethod
    async def generate_answer(
        self,
        question: str,
        memories: List[Dict[str, Any]]
    ) -> Tuple[str, str]:
        """
        Generate answer from retrieved memories.

        Returns:
            Tuple of (answer, reasoning)
        """
        pass


class LLMAnswerEvaluator(ABC):
    """Abstract base class for LLM-based answer evaluation."""

    @abstractmethod
    async def judge_answer(
        self,
        question: str,
        correct_answer: str,
        predicted_answer: str,
        semaphore: asyncio.Semaphore
    ) -> Tuple[bool, str]:
        """
        Evaluate predicted answer against correct answer.

        Args:
            question: The question
            correct_answer: Gold/correct answer
            predicted_answer: Predicted answer
            semaphore: Semaphore for rate limiting

        Returns:
            Tuple of (is_correct, reasoning)
        """
        pass


class BenchmarkRunner:
    """
    Common benchmark runner using the proven LoComo approach.

    Optimizations:
    - Batch ingestion (put_batch_async)
    - Parallel question processing with rate limiting
    - Parallel LLM judging with rate limiting
    - Progress tracking
    """

    def __init__(
        self,
        dataset: BenchmarkDataset,
        answer_generator: LLMAnswerGenerator,
        answer_evaluator: LLMAnswerEvaluator,
        memory: Optional[TemporalSemanticMemory] = None
    ):
        """
        Initialize benchmark runner.

        Args:
            dataset: Dataset implementation
            answer_generator: Answer generator implementation
            answer_evaluator: Answer evaluator implementation
            memory: Memory system instance (creates new if None)
        """
        self.dataset = dataset
        self.answer_generator = answer_generator
        self.answer_evaluator = answer_evaluator
        self.memory = memory or TemporalSemanticMemory()

    async def ingest_conversation(
        self,
        item: Dict[str, Any],
        agent_id: str
    ) -> int:
        """
        Ingest conversation into memory using batch ingestion.

        Uses put_batch_async for maximum efficiency.

        Returns:
            Number of sessions ingested
        """
        batch_contents = self.dataset.prepare_sessions_for_ingestion(item)

        if batch_contents:
            await self.memory.put_batch_async(
                agent_id=agent_id,
                contents=batch_contents
            )

        return len(batch_contents)

    async def answer_question(
        self,
        agent_id: str,
        question: str,
        thinking_budget: int = 500,
        top_k: int = 20,
        weight_activation: float = 0.30,
        weight_semantic: float = 0.30,
        weight_recency: float = 0.25,
        weight_frequency: float = 0.15,
    ) -> Tuple[str, str, List[Dict]]:
        """
        Answer a question using memory retrieval.

        Returns:
            Tuple of (answer, reasoning, retrieved_memories)
        """
        # Search memory
        results, _ = await self.memory.search_async(
            agent_id=agent_id,
            query=question,
            thinking_budget=thinking_budget,
            top_k=top_k,
            weight_activation=weight_activation,
            weight_semantic=weight_semantic,
            weight_recency=weight_recency,
            weight_frequency=weight_frequency,
        )

        if not results:
            return "I don't have enough information to answer that question.", "No relevant memories found.", []

        # Generate answer using LLM
        answer, reasoning = await self.answer_generator.generate_answer(question, results)

        return answer, reasoning, results

    async def evaluate_qa_task(
        self,
        agent_id: str,
        qa_pairs: List[Dict],
        item_id: str,
        thinking_budget: int,
        top_k: int,
        max_questions: Optional[int] = None,
        semaphore: asyncio.Semaphore = None,
        weight_activation: float = 0.30,
        weight_semantic: float = 0.30,
        weight_recency: float = 0.25,
        weight_frequency: float = 0.15,
    ) -> List[Dict]:
        """
        Evaluate QA task with parallel question processing.

        Args:
            semaphore: Semaphore to limit concurrent question processing

        Returns:
            List of QA results
        """
        # Filter out questions without answers (category 5)
        qa_pairs = [pair for pair in qa_pairs if pair.get('answer')]
        questions_to_eval = qa_pairs[:max_questions] if max_questions else qa_pairs

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            console=console
        ) as progress:
            task = progress.add_task(
                f"[cyan]Evaluating QA for {item_id} - {len(questions_to_eval)} questions",
                total=len(questions_to_eval)
            )

            # Create tasks for all questions
            async def process_question(qa):
                async with semaphore:
                    question = qa['question']
                    correct_answer = qa['answer']
                    category = qa.get('category', 0)

                    # Get predicted answer, reasoning, and retrieved memories
                    predicted_answer, reasoning, retrieved_memories = await self.answer_question(
                        agent_id, question, thinking_budget, top_k,
                        weight_activation, weight_semantic, weight_recency, weight_frequency
                    )

                    return {
                        'question': question,
                        'correct_answer': correct_answer,
                        'predicted_answer': predicted_answer,
                        'reasoning': reasoning,
                        'category': category,
                        'retrieved_memories': retrieved_memories
                    }

            question_tasks = [process_question(qa) for qa in questions_to_eval]

            # Use as_completed to update progress as results come in
            results = []
            for coro in asyncio.as_completed(question_tasks):
                result = await coro
                results.append(result)
                progress.update(task, advance=1)

        return results

    async def calculate_metrics(self, results: List[Dict], eval_semaphore_size: int = 8) -> Dict:
        """
        Calculate evaluation metrics using parallel LLM-as-judge.

        Args:
            results: QA results to evaluate
            eval_semaphore_size: Max concurrent LLM judge requests

        Returns:
            Dict with evaluation metrics
        """
        total = len(results)

        # Semaphore to limit concurrent requests
        semaphore = asyncio.Semaphore(eval_semaphore_size)

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            console=console
        ) as progress:
            task = progress.add_task(
                f"[yellow]Judging answers with LLM (parallel, max {eval_semaphore_size})...",
                total=total
            )

            # Create all judgment tasks
            async def judge_single(result):
                is_correct, eval_reasoning = await self.answer_evaluator.judge_answer(
                    result['question'],
                    result['correct_answer'],
                    result['predicted_answer'],
                    semaphore
                )
                result['is_correct'] = is_correct
                result['correctness_reasoning'] = eval_reasoning
                return result

            judgment_tasks = [judge_single(result) for result in results]

            # Process in parallel with progress updates
            judged_results = []
            for coro in asyncio.as_completed(judgment_tasks):
                judged_result = await coro
                judged_results.append(judged_result)
                progress.update(task, advance=1)

        # Calculate stats
        correct = sum(1 for r in judged_results if r.get('is_correct', False))
        category_stats = {}

        for result in judged_results:
            category = result.get('category', 'unknown')
            if category not in category_stats:
                category_stats[category] = {'correct': 0, 'total': 0}
            category_stats[category]['total'] += 1
            if result.get('is_correct', False):
                category_stats[category]['correct'] += 1

        accuracy = (correct / total * 100) if total > 0 else 0

        return {
            'accuracy': accuracy,
            'correct': correct,
            'total': total,
            'category_stats': category_stats,
            'detailed_results': judged_results
        }

    async def process_single_item(
        self,
        item: Dict,
        agent_id: str,
        i: int,
        total_items: int,
        thinking_budget: int,
        top_k: int,
        max_questions_per_item: Optional[int],
        skip_ingestion: bool,
        question_semaphore: asyncio.Semaphore,
        eval_semaphore_size: int = 8,
        weight_activation: float = 0.30,
        weight_semantic: float = 0.30,
        weight_recency: float = 0.25,
        weight_frequency: float = 0.15,
    ) -> Dict:
        """
        Process a single item (ingest + evaluate).

        Returns:
            Result dict with metrics
        """
        item_id = self.dataset.get_item_id(item)

        console.print(f"\n[bold blue]Item {i}/{total_items}[/bold blue] (ID: {item_id})")

        if not skip_ingestion:
            # Clear previous agent data only on first item
            if i == 1:
                console.print("  [1] Clearing previous agent data...")
                await self.memory.delete_agent(agent_id)
                console.print(f"      [green]✓[/green] Cleared '{agent_id}' agent data")

            # Ingest conversation
            console.print("  [2] Ingesting conversation (batch mode)...")
            num_sessions = await self.ingest_conversation(item, agent_id)
            console.print(f"      [green]✓[/green] Ingested {num_sessions} sessions")
        else:
            num_sessions = -1

        # Evaluate QA
        qa_pairs = self.dataset.get_qa_pairs(item)
        console.print(f"  [3] Evaluating {len(qa_pairs)} QA pairs (parallel)...")
        qa_results = await self.evaluate_qa_task(
            agent_id,
            qa_pairs,
            item_id,
            thinking_budget,
            top_k,
            max_questions_per_item,
            question_semaphore,
            weight_activation,
            weight_semantic,
            weight_recency,
            weight_frequency
        )

        # Calculate metrics
        console.print("  [4] Calculating metrics...")
        metrics = await self.calculate_metrics(qa_results, eval_semaphore_size)

        console.print(f"      [green]✓[/green] Accuracy: {metrics['accuracy']:.2f}% ({metrics['correct']}/{metrics['total']})")

        return {
            'item_id': item_id,
            'metrics': metrics,
            'num_sessions': num_sessions
        }

    async def run(
        self,
        dataset_path: Path,
        agent_id: str,
        max_items: Optional[int] = None,
        max_questions_per_item: Optional[int] = None,
        thinking_budget: int = 500,
        top_k: int = 20,
        skip_ingestion: bool = False,
        max_concurrent_questions: int = 16,
        eval_semaphore_size: int = 8,
        clear_agent_per_item: bool = False,
        weight_activation: float = 0.30,
        weight_semantic: float = 0.30,
        weight_recency: float = 0.25,
        weight_frequency: float = 0.15,
    ) -> Dict[str, Any]:
        """
        Run the full benchmark evaluation.

        Args:
            dataset_path: Path to dataset file
            agent_id: Agent ID to use
            max_items: Maximum number of items to evaluate
            max_questions_per_item: Maximum questions per item
            thinking_budget: Thinking budget for search
            top_k: Number of memories to retrieve
            skip_ingestion: Skip ingestion and use existing data
            max_concurrent_questions: Max concurrent question processing
            eval_semaphore_size: Max concurrent LLM judge requests
            clear_agent_per_item: Clear agent data before each item (for isolation)
            weight_activation: Weight for activation score in final ranking (default: 0.30)
            weight_semantic: Weight for semantic similarity in final ranking (default: 0.30)
            weight_recency: Weight for recency score in final ranking (default: 0.25)
            weight_frequency: Weight for frequency score in final ranking (default: 0.15)

        Returns:
            Dict with complete benchmark results
        """
        console.print(f"\n[bold cyan]Benchmark Evaluation[/bold cyan]")
        console.print("=" * 80)

        # Load dataset
        console.print(f"\n[1] Loading dataset from {dataset_path}...")
        items = self.dataset.load(dataset_path, max_items)
        console.print(f"    [green]✓[/green] Loaded {len(items)} items")

        # Initialize memory system
        console.print(f"\n[2] Initializing memory system...")
        console.print(f"    [green]✓[/green] Memory system initialized")

        # Create semaphore for question processing
        question_semaphore = asyncio.Semaphore(max_concurrent_questions)

        # Process items
        all_results = []

        for i, item in enumerate(items, 1):
            # Clear agent per item if requested (for isolation in benchmarks like LongMemEval)
            if clear_agent_per_item and i > 1 and not skip_ingestion:
                await self.memory.delete_agent(agent_id)

            result = await self.process_single_item(
                item, agent_id, i, len(items),
                thinking_budget, top_k, max_questions_per_item,
                skip_ingestion, question_semaphore, eval_semaphore_size,
                weight_activation, weight_semantic, weight_recency, weight_frequency
            )
            all_results.append(result)

        # Calculate overall metrics
        total_correct = sum(r['metrics']['correct'] for r in all_results)
        total_questions = sum(r['metrics']['total'] for r in all_results)
        overall_accuracy = (total_correct / total_questions * 100) if total_questions > 0 else 0

        return {
            'overall_accuracy': overall_accuracy,
            'total_correct': total_correct,
            'total_questions': total_questions,
            'num_items': len(items),
            'item_results': all_results
        }

    def display_results(self, results: Dict[str, Any]):
        """Display benchmark results in a formatted table."""
        console.print("\n[bold green]✓ Benchmark Complete![/bold green]\n")

        # Display results table
        table = Table(title="Benchmark Results", box=box.ROUNDED)
        table.add_column("Item ID", style="cyan")
        table.add_column("Sessions", justify="right", style="yellow")
        table.add_column("Questions", justify="right", style="blue")
        table.add_column("Correct", justify="right", style="green")
        table.add_column("Accuracy", justify="right", style="magenta")

        for result in results['item_results']:
            metrics = result['metrics']
            table.add_row(
                result['item_id'],
                str(result['num_sessions']),
                str(metrics['total']),
                str(metrics['correct']),
                f"{metrics['accuracy']:.1f}%"
            )

        table.add_row(
            "[bold]OVERALL[/bold]",
            "-",
            f"[bold]{results['total_questions']}[/bold]",
            f"[bold]{results['total_correct']}[/bold]",
            f"[bold]{results['overall_accuracy']:.1f}%[/bold]"
        )

        console.print(table)

    def save_results(self, results: Dict[str, Any], output_path: Path):
        """Save results to JSON file."""
        with open(output_path, 'w') as f:
            json.dump(results, f, indent=2, default=str)
        console.print(f"\n[green]✓[/green] Results saved to {output_path}")