"""Sentrial Experiments - Run prompt A/B tests locally with results synced to Sentrial."""

import os
import logging
from dataclasses import dataclass
from typing import Any, Callable, Optional, List
from concurrent.futures import ThreadPoolExecutor, as_completed

from .client import SentrialClient

_logger = logging.getLogger("sentrial.experiment")


@dataclass
class ExperimentVariant:
    """A prompt variant to test."""
    name: str
    system_prompt: str
    description: Optional[str] = None


@dataclass
class ExperimentTestCase:
    """A test case (input) to run against each variant."""
    session_id: str
    user_input: str


@dataclass
class ExperimentRunResult:
    """Result of a single run."""
    variant_name: str
    test_case_session_id: str
    result_session_id: Optional[str]
    success: bool
    error_message: Optional[str] = None


class ExperimentRunTracker:
    """
    Context manager for tracking a single experiment run.

    Usage:
        with experiment.track_run("variant_a", "session_123") as tracker:
            tracker.session_id  # Use this for your agent
            # Run your agent...
            tracker.set_result_session_id(session_id)
    """

    def __init__(
        self,
        experiment: "Experiment",
        variant_name: str,
        base_session_id: str,
    ):
        self.experiment = experiment
        self.variant_name = variant_name
        self.base_session_id = base_session_id
        self.run_id: Optional[str] = None
        self.result_session_id: Optional[str] = None
        self._success = True
        self._error_message: Optional[str] = None

    def __enter__(self) -> "ExperimentRunTracker":
        """Start the run - creates a run record via API."""
        response = self.experiment._client._safe_request(
            "post",
            f"{self.experiment._client.api_url}/api/experiments/{self.experiment.experiment_id}/runs",
            json={
                "variantName": self.variant_name,
                "baseSessionId": self.base_session_id,
            },
        )
        if response and response.ok:
            data = response.json()
            self.run_id = data.get("run", {}).get("id")
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
        """Complete the run - updates status via API."""
        if exc_type is not None:
            self._success = False
            self._error_message = f"{exc_type.__name__}: {exc_val}" if exc_val else str(exc_type.__name__)

        if self.run_id:
            self.experiment._client._safe_request(
                "patch",
                f"{self.experiment._client.api_url}/api/experiments/{self.experiment.experiment_id}/runs/{self.run_id}",
                json={
                    "status": "completed" if self._success else "failed",
                    "resultSessionId": self.result_session_id,
                    "errorMessage": self._error_message,
                },
            )

        return False  # Don't suppress exceptions

    def set_result_session_id(self, session_id: str) -> None:
        """Set the session ID of the result session."""
        self.result_session_id = session_id

    def set_error(self, message: str) -> None:
        """Mark this run as failed."""
        self._success = False
        self._error_message = message


class Experiment:
    """
    Run experiments locally with results synced to Sentrial.

    Usage:
        exp = Experiment("exp_abc123")

        # Get config
        print(f"Testing {len(exp.variants)} variants on {len(exp.test_cases)} inputs")

        # Run with your agent function
        exp.run(my_agent_function, parallel=True)

        # Or run manually
        for variant in exp.variants:
            for test_case in exp.test_cases:
                with exp.track_run(variant.name, test_case.session_id) as tracker:
                    # Run your agent
                    session_id = run_agent(test_case.user_input, variant.system_prompt)
                    tracker.set_result_session_id(session_id)
    """

    def __init__(
        self,
        experiment_id: str,
        api_key: Optional[str] = None,
        api_url: Optional[str] = None,
    ):
        """
        Initialize an experiment.

        Args:
            experiment_id: The experiment ID from Sentrial dashboard
            api_key: API key (defaults to SENTRIAL_API_KEY env var)
            api_url: API URL (defaults to SENTRIAL_API_URL env var or production)
        """
        self.experiment_id = experiment_id
        self._client = SentrialClient(
            api_key=api_key,
            api_url=api_url,
            fail_silently=False,  # We want errors for experiments
        )

        # Lazy-loaded config
        self._config: Optional[dict] = None
        self._variants: Optional[List[ExperimentVariant]] = None
        self._test_cases: Optional[List[ExperimentTestCase]] = None

    def _fetch_config(self) -> dict:
        """Fetch experiment configuration from API."""
        if self._config is None:
            response = self._client._safe_request(
                "get",
                f"{self._client.api_url}/api/experiments/{self.experiment_id}/runs",
            )
            if response is None:
                raise RuntimeError(f"Failed to fetch experiment config for {self.experiment_id}")

            data = response.json()
            self._config = data.get("experiment", {})

            # Parse variants
            variants_data = self._config.get("variants", [])
            self._variants = [
                ExperimentVariant(
                    name=v.get("name", ""),
                    system_prompt=v.get("systemPrompt", ""),
                    description=v.get("description"),
                )
                for v in variants_data
            ]

            # Parse test cases
            test_cases_data = self._config.get("testCases", [])
            self._test_cases = [
                ExperimentTestCase(
                    session_id=tc.get("sessionId", ""),
                    user_input=tc.get("userInput", ""),
                )
                for tc in test_cases_data
                if tc.get("userInput")  # Only include test cases with inputs
            ]

        return self._config

    @property
    def name(self) -> str:
        """Get experiment name."""
        return self._fetch_config().get("name", "")

    @property
    def status(self) -> str:
        """Get experiment status."""
        return self._fetch_config().get("status", "")

    @property
    def variants(self) -> List[ExperimentVariant]:
        """Get list of variants to test."""
        self._fetch_config()
        return self._variants or []

    @property
    def test_cases(self) -> List[ExperimentTestCase]:
        """Get list of test cases (inputs) to run."""
        self._fetch_config()
        return self._test_cases or []

    def get_variant(self, name: str) -> Optional[ExperimentVariant]:
        """Get a specific variant by name."""
        for v in self.variants:
            if v.name == name:
                return v
        return None

    def track_run(self, variant_name: str, base_session_id: str) -> ExperimentRunTracker:
        """
        Create a run tracker for manual experiment runs.

        Args:
            variant_name: Name of the variant being tested
            base_session_id: Session ID of the original test case

        Returns:
            ExperimentRunTracker context manager
        """
        return ExperimentRunTracker(self, variant_name, base_session_id)

    def run(
        self,
        agent_fn: Callable[[ExperimentTestCase, ExperimentVariant, ExperimentRunTracker], Any],
        parallel: bool = False,
        max_workers: int = 4,
    ) -> List[ExperimentRunResult]:
        """
        Run the experiment with all variants and test cases.

        Args:
            agent_fn: Function that takes (test_case, variant, tracker) and runs your agent.
                      The function should:
                      1. Create a session using tracker.experiment._client or your own client
                      2. Run the agent with variant.system_prompt
                      3. Call tracker.set_result_session_id(session_id)
            parallel: Whether to run test cases in parallel (default: False)
            max_workers: Number of parallel workers (default: 4)

        Returns:
            List of ExperimentRunResult objects

        Example:
            def my_agent(test_case, variant, tracker):
                session_id = client.create_session(
                    name=f"Experiment: {test_case.session_id}",
                    agent_name="support-agent",
                    user_id="experiment"
                )
                tracker.set_result_session_id(session_id)

                # Run your agent with the variant's system prompt
                result = agent_executor.invoke({
                    "input": test_case.user_input,
                    "system_prompt": variant.system_prompt,
                })

                client.complete_session(session_id, success=True)

            experiment.run(my_agent, parallel=True)
        """
        results: List[ExperimentRunResult] = []
        total_runs = len(self.variants) * len(self.test_cases)
        completed = 0

        print(f"Running experiment: {self.name}")
        print(f"  Variants: {len(self.variants)}")
        print(f"  Test cases: {len(self.test_cases)}")
        print(f"  Total runs: {total_runs}")
        print()

        def run_single(variant: ExperimentVariant, test_case: ExperimentTestCase) -> ExperimentRunResult:
            """Run a single variant × test case combination."""
            try:
                with self.track_run(variant.name, test_case.session_id) as tracker:
                    agent_fn(test_case, variant, tracker)
                    return ExperimentRunResult(
                        variant_name=variant.name,
                        test_case_session_id=test_case.session_id,
                        result_session_id=tracker.result_session_id,
                        success=tracker._success,
                        error_message=tracker._error_message,
                    )
            except Exception as e:
                _logger.error(f"Run failed: {variant.name} x {test_case.session_id}: {e}")
                return ExperimentRunResult(
                    variant_name=variant.name,
                    test_case_session_id=test_case.session_id,
                    result_session_id=None,
                    success=False,
                    error_message=str(e),
                )

        if parallel:
            # Run in parallel
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                futures = {}
                for variant in self.variants:
                    for test_case in self.test_cases:
                        future = executor.submit(run_single, variant, test_case)
                        futures[future] = (variant.name, test_case.session_id)

                for future in as_completed(futures):
                    variant_name, session_id = futures[future]
                    result = future.result()
                    results.append(result)
                    completed += 1
                    status = "✓" if result.success else "✗"
                    print(f"  [{completed}/{total_runs}] {status} {variant_name} x {session_id[:8]}...")
        else:
            # Run sequentially
            for variant in self.variants:
                for test_case in self.test_cases:
                    result = run_single(variant, test_case)
                    results.append(result)
                    completed += 1
                    status = "✓" if result.success else "✗"
                    print(f"  [{completed}/{total_runs}] {status} {variant.name} x {test_case.session_id[:8]}...")

        print()
        print(f"Experiment complete!")
        print(f"  Successful: {sum(1 for r in results if r.success)}/{total_runs}")
        print(f"  Failed: {sum(1 for r in results if not r.success)}/{total_runs}")

        return results

    def reset(self) -> bool:
        """
        Reset all runs for this experiment (for re-running).

        Returns:
            True if successful, False otherwise
        """
        response = self._client._safe_request(
            "delete",
            f"{self._client.api_url}/api/experiments/{self.experiment_id}/runs",
        )
        if response and response.ok:
            # Clear cached config
            self._config = None
            self._variants = None
            self._test_cases = None
            return True
        return False

    def get_results(self) -> dict:
        """
        Fetch aggregated results from the API.

        Returns:
            Dict with experiment results including winner
        """
        response = self._client._safe_request(
            "get",
            f"{self._client.api_url}/api/experiments/{self.experiment_id}/results",
        )
        if response and response.ok:
            return response.json()
        return {}
