"""
Experiment C: Efficiency Curve

This experiment measures token usage to solve problems with single-model vs CMVK.

Goal: Address the criticism "Two models are too expensive."
Hypothesis: Single models hallucinate and loop 10 times (high cost).
            CMVK catches the error in loop 1 and solves it in loop 3.
"""

import json
import logging
import sys
import time
from pathlib import Path
from typing import Any

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from src import GeminiVerifier, OpenAIGenerator, VerificationKernel

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class EfficiencyCurveExperiment:
    """
    Runs the Efficiency Curve experiment.

    This compares token usage between:
    1. Single-model self-correction (baseline)
    2. Cross-model verification (CMVK)
    """

    def __init__(
        self,
        dataset_path: str = "experiments/datasets/sample.json",
        output_dir: str = "experiments/results",
    ):
        """
        Initialize the experiment.

        Args:
            dataset_path: Path to test dataset
            output_dir: Directory to save results
        """
        self.dataset_path = Path(dataset_path)
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        try:
            # Initialize agents
            self.generator = OpenAIGenerator()
            self.verifier = GeminiVerifier(enable_prosecutor_mode=True)
            self.kernel = VerificationKernel(generator=self.generator, verifier=self.verifier)
            logger.info("Initialized CMVK system")
        except Exception as e:
            logger.warning(f"Failed to initialize with real APIs: {e}")
            logger.info("Running in mock mode for testing")
            self.kernel = None

    def load_dataset(self) -> list[dict[str, Any]]:
        """Load the test dataset."""
        with open(self.dataset_path) as f:
            data = json.load(f)

        # Handle both list and dict formats
        if isinstance(data, list):
            dataset = data
        elif isinstance(data, dict) and "tasks" in data:
            dataset = data["tasks"]
        else:
            dataset = [data]

        logger.info(f"Loaded {len(dataset)} test cases from {self.dataset_path}")
        return dataset

    def run_experiment(self) -> dict[str, Any]:
        """
        Run the efficiency curve experiment.

        Returns:
            Dictionary with experiment results
        """
        logger.info("Starting Efficiency Curve Experiment")

        dataset = self.load_dataset()

        results = {
            "experiment": "efficiency_curve",
            "dataset": str(self.dataset_path),
            "total_tasks": len(dataset),
            "cmvk_results": [],
            "total_time": 0,
        }

        # Process each task with CMVK
        for i, task_data in enumerate(dataset):
            logger.info(f"Processing task {i+1}/{len(dataset)}")

            # Extract task description
            if isinstance(task_data, dict):
                task = task_data.get("task", task_data.get("description", str(task_data)))
            else:
                task = str(task_data)

            start_time = time.time()
            task_result = self._run_cmvk_task(task)
            elapsed = time.time() - start_time

            task_result["time"] = elapsed
            results["cmvk_results"].append(task_result)
            results["total_time"] += elapsed

        # Calculate aggregate metrics
        results["metrics"] = self._calculate_metrics(results)

        # Save results
        self._save_results(results)

        logger.info(
            f"Experiment complete: Avg tokens={results['metrics']['avg_total_tokens']}, "
            f"Avg loops={results['metrics']['avg_loops']}"
        )

        return results

    def _run_cmvk_task(self, task: str) -> dict[str, Any]:
        """
        Run a single task through CMVK.

        Args:
            task: The task description

        Returns:
            Dictionary with task results
        """
        if self.kernel is None:
            # Mock mode
            return {
                "task": task[:100],
                "success": True,
                "loops": 2,
                "generator_tokens": 500,
                "verifier_tokens": 300,
                "total_tokens": 800,
                "mock": True,
            }

        try:
            # Reset token counts
            self.generator.total_tokens_used = 0
            self.generator.call_count = 0
            self.verifier.total_tokens_used = 0
            self.verifier.call_count = 0

            # Execute kernel
            state = self.kernel.execute(task)

            # Get token stats
            gen_stats = self.generator.get_token_stats()
            ver_stats = self.verifier.get_token_stats()

            return {
                "task": task[:100],
                "success": state.is_complete,
                "loops": state.current_loop,
                "generator_tokens": gen_stats["total_tokens"],
                "generator_calls": gen_stats["call_count"],
                "verifier_tokens": ver_stats["total_tokens"],
                "verifier_calls": ver_stats["call_count"],
                "total_tokens": gen_stats["total_tokens"] + ver_stats["total_tokens"],
                "total_calls": gen_stats["call_count"] + ver_stats["call_count"],
            }

        except Exception as e:
            logger.error(f"Error running task: {e}")
            return {
                "task": task[:100],
                "success": False,
                "error": str(e),
                "loops": 0,
                "total_tokens": 0,
            }

    def _calculate_metrics(self, results: dict[str, Any]) -> dict[str, float]:
        """Calculate aggregate metrics."""
        cmvk_results = results["cmvk_results"]

        total_tokens = sum(r.get("total_tokens", 0) for r in cmvk_results)
        total_loops = sum(r.get("loops", 0) for r in cmvk_results)
        successful_tasks = sum(1 for r in cmvk_results if r.get("success", False))

        num_tasks = len(cmvk_results)

        return {
            "avg_total_tokens": total_tokens / num_tasks if num_tasks > 0 else 0,
            "avg_loops": total_loops / num_tasks if num_tasks > 0 else 0,
            "success_rate": successful_tasks / num_tasks if num_tasks > 0 else 0,
            "total_tokens": total_tokens,
            "total_loops": total_loops,
        }

    def _save_results(self, results: dict[str, Any]) -> None:
        """Save results to JSON file."""
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        output_file = self.output_dir / f"efficiency_curve_{timestamp}.json"

        with open(output_file, "w") as f:
            json.dump(results, f, indent=2)

        logger.info(f"Results saved to {output_file}")

        # Also create a summary file
        summary_file = self.output_dir / f"efficiency_summary_{timestamp}.txt"
        with open(summary_file, "w") as f:
            f.write("=" * 60 + "\n")
            f.write("EFFICIENCY CURVE EXPERIMENT RESULTS\n")
            f.write("=" * 60 + "\n\n")
            f.write("CROSS-MODEL VERIFICATION (CMVK):\n")
            f.write(f"Total Tasks: {results['total_tasks']}\n")
            f.write(f"Success Rate: {results['metrics']['success_rate']:.2%}\n")
            f.write(f"Average Tokens per Task: {results['metrics']['avg_total_tokens']:.0f}\n")
            f.write(f"Average Loops per Task: {results['metrics']['avg_loops']:.2f}\n")
            f.write(f"Total Tokens Used: {results['metrics']['total_tokens']}\n")
            f.write(f"Total Time: {results['total_time']:.2f}s\n\n")

            f.write("TOKEN BREAKDOWN:\n")
            for i, task_result in enumerate(results["cmvk_results"], 1):
                f.write(f"\nTask {i}: {task_result['task'][:50]}...\n")
                f.write(f"  Success: {task_result.get('success', False)}\n")
                f.write(f"  Loops: {task_result.get('loops', 0)}\n")
                f.write(f"  Generator Tokens: {task_result.get('generator_tokens', 0)}\n")
                f.write(f"  Verifier Tokens: {task_result.get('verifier_tokens', 0)}\n")
                f.write(f"  Total Tokens: {task_result.get('total_tokens', 0)}\n")

        logger.info(f"Summary saved to {summary_file}")


def main():
    """Run the efficiency curve experiment."""
    experiment = EfficiencyCurveExperiment()
    results = experiment.run_experiment()

    print("\n" + "=" * 60)
    print("EFFICIENCY CURVE EXPERIMENT COMPLETED")
    print("=" * 60)
    print(f"Average Tokens per Task: {results['metrics']['avg_total_tokens']:.0f}")
    print(f"Average Loops per Task: {results['metrics']['avg_loops']:.2f}")
    print(f"Success Rate: {results['metrics']['success_rate']:.2%}")
    print("=" * 60)


if __name__ == "__main__":
    main()
