"""Speculative Fix Application with Racing (Phase 6).

Applies multiple fix candidates in parallel and races them:
1. Take top N fix candidates
2. Apply all in parallel temp sandboxes
3. Run verification tests on each
4. Return first one that passes all N runs
5. Cancel remaining workers when winner found
"""

from __future__ import annotations

import asyncio
import logging
import shutil
import tempfile
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Set

from autodebug.llm_first.iterative_refiner import (
    IterativeRefiner,
)
from autodebug.llm_first.models import (
    FailureContext,
    FixHypothesis,
    VerificationResult,
)

logger = logging.getLogger(__name__)

# Default configuration
DEFAULT_MAX_CANDIDATES = 5
DEFAULT_VERIFICATION_RUNS = 10
DEFAULT_RACE_TIMEOUT = 300  # 5 minutes


@dataclass
class RaceResult:
    """Result of racing multiple fix candidates.

    Attributes:
        winner: The winning fix hypothesis that passed all verification runs.
        winner_result: Verification result for the winner.
        all_results: Dict mapping fix IDs to their verification results.
        race_duration: Total time taken for the race in seconds.
        cancelled_count: Number of workers cancelled after winner found.
    """

    winner: Optional[FixHypothesis] = None
    winner_result: Optional[VerificationResult] = None
    partial_fix: Optional[FixHypothesis] = None
    partial_result: Optional[VerificationResult] = None
    all_results: Dict[str, VerificationResult] = field(default_factory=dict)
    race_duration: float = 0.0
    cancelled_count: int = 0

    @property
    def has_winner(self) -> bool:
        """Check if a winner was found."""
        return self.winner is not None and self.winner_result is not None

    @property
    def best_partial(self) -> Optional[tuple[FixHypothesis, VerificationResult]]:
        """Get the best partial result if no winner."""
        if self.has_winner:
            return None
        if self.partial_fix is None or self.partial_result is None:
            return None
        return self.partial_fix, self.partial_result


class SpeculativeFixer:
    """Apply multiple fix candidates in parallel and race them.

    Uses asyncio for concurrent execution with proper cancellation handling.
    Integrates with IterativeRefiner for test execution.

    Example:
        >>> fixer = SpeculativeFixer(refiner=refiner)
        >>> result = await fixer.apply_and_race(fixes, context)
        >>> if result.winner:
        ...     print(f"Winner: {result.winner.description}")
    """

    def __init__(
        self,
        refiner: Optional[IterativeRefiner] = None,
        max_candidates: int = DEFAULT_MAX_CANDIDATES,
        verification_runs: int = DEFAULT_VERIFICATION_RUNS,
        race_timeout: float = DEFAULT_RACE_TIMEOUT,
        parallel_verifications: int = 3,
        cleanup_on_completion: bool = True,
    ):
        """Initialize speculative fixer.

        Args:
            refiner: IterativeRefiner instance for test execution.
            max_candidates: Maximum number of fix candidates to race.
            verification_runs: Number of verification runs per fix.
            race_timeout: Maximum time for the race in seconds.
            parallel_verifications: Max parallel verification jobs.
            cleanup_on_completion: Whether to clean up temp dirs.
        """
        self.refiner = refiner or IterativeRefiner()
        self.max_candidates = max_candidates
        self.verification_runs = verification_runs
        self.race_timeout = race_timeout
        self.parallel_verifications = parallel_verifications
        self.cleanup_on_completion = cleanup_on_completion

        # Track active tasks for cancellation
        self._active_tasks: Dict[str, asyncio.Task] = {}
        self._cancelled_ids: Set[str] = set()
        self._winner_found: asyncio.Event = asyncio.Event()
        self._winner_id: Optional[str] = None
        self._lock: asyncio.Lock = asyncio.Lock()

        # Work directories for cleanup
        self._work_dirs: List[Path] = []

    async def apply_and_race(
        self,
        fixes: List[FixHypothesis],
        context: FailureContext,
    ) -> RaceResult:
        """Apply multiple fixes in parallel and race them.

        Takes top N fix candidates, applies each in a temp sandbox,
        runs verification tests, and returns the first one that passes
        all runs. Cancels remaining workers when winner found.

        Args:
            fixes: List of fix hypotheses to race.
            context: Failure context for test execution.

        Returns:
            RaceResult with winner (if found) and all results.
        """
        start_time = time.time()

        # Reset state for new race
        self._active_tasks.clear()
        self._cancelled_ids.clear()
        self._winner_found.clear()
        self._winner_id = None

        # Take top N candidates
        candidates = fixes[:self.max_candidates]

        if not candidates:
            return RaceResult(race_duration=0.0)

        # Ensure all fixes have IDs
        for fix in candidates:
            if not fix.id:
                fix.id = str(uuid.uuid4())

        logger.info(
            "Starting race with %d candidates (max %d verification runs each)",
            len(candidates),
            self.verification_runs,
        )

        all_results: Dict[str, VerificationResult] = {}
        winner: Optional[FixHypothesis] = None
        winner_result: Optional[VerificationResult] = None
        partial_fix: Optional[FixHypothesis] = None
        partial_result: Optional[VerificationResult] = None

        try:
            # Create verification tasks for each candidate
            semaphore = asyncio.Semaphore(self.parallel_verifications)

            async def run_with_semaphore(fix: FixHypothesis) -> tuple[str, VerificationResult]:
                async with semaphore:
                    return await self._verify_candidate(fix, context)

            # Start all tasks
            tasks = []
            for fix in candidates:
                task = asyncio.create_task(
                    run_with_semaphore(fix),
                    name=f"verify_{fix.id}",
                )
                self._active_tasks[fix.id] = task
                tasks.append(task)

            # Race with timeout and early winner detection
            try:
                done, pending = await asyncio.wait(
                    tasks,
                    timeout=self.race_timeout,
                    return_when=asyncio.FIRST_COMPLETED,
                )

                # Process completed tasks looking for winner
                while done or pending:
                    for task in done:
                        try:
                            fix_id, result = task.result()
                            all_results[fix_id] = result

                            # Check if this is a winner
                            if self._is_winner(result):
                                async with self._lock:
                                    if not self._winner_found.is_set():
                                        self._winner_id = fix_id
                                        self._winner_found.set()

                                        # Find the fix object
                                        for fix in candidates:
                                            if fix.id == fix_id:
                                                winner = fix
                                                winner_result = result
                                                break

                                        logger.info(
                                            "Winner found: %s (passed %d/%d runs)",
                                            fix_id,
                                            result.passed_runs,
                                            result.total_runs,
                                        )
                        except asyncio.CancelledError:
                            pass
                        except Exception as e:
                            logger.warning("Task failed: %s", e)

                    # If winner found, cancel remaining and break
                    if self._winner_found.is_set():
                        cancelled = await self.cancel_remaining(except_winner=self._winner_id or "")
                        break

                    # Wait for more completions if no winner yet
                    if pending:
                        done, pending = await asyncio.wait(
                            pending,
                            timeout=self.race_timeout - (time.time() - start_time),
                            return_when=asyncio.FIRST_COMPLETED,
                        )
                    else:
                        break

            except asyncio.TimeoutError:
                logger.warning("Race timeout after %.1f seconds", self.race_timeout)
                await self.cancel_remaining(except_winner="")

        except Exception as e:
            logger.exception("Race failed: %s", e)
            await self.cancel_remaining(except_winner="")

        finally:
            # Cleanup work directories
            if self.cleanup_on_completion:
                await self._cleanup_work_dirs()

        race_duration = time.time() - start_time

        # If no winner, find best partial result
        if not winner and all_results:
            best_fix, best_result = self._find_best_partial(candidates, all_results)
            if best_fix and best_result:
                partial_fix = best_fix
                partial_result = best_result
                logger.info(
                    "No winner found. Best partial: %s (passed %d/%d runs)",
                    best_fix.id,
                    best_result.passed_runs,
                    best_result.total_runs,
                )

        return RaceResult(
            winner=winner,
            winner_result=winner_result,
            partial_fix=partial_fix,
            partial_result=partial_result,
            all_results=all_results,
            race_duration=race_duration,
            cancelled_count=len(self._cancelled_ids),
        )

    async def verify_in_parallel(
        self,
        fix: FixHypothesis,
        n_workers: int,
    ) -> VerificationResult:
        """Verify a single fix in parallel using multiple workers.

        Runs the verification tests N times in parallel to check
        for determinism.

        Args:
            fix: Fix hypothesis to verify.
            n_workers: Number of parallel verification workers.

        Returns:
            VerificationResult with pass/fail statistics.
        """
        if not fix.id:
            fix.id = str(uuid.uuid4())

        # Create a temporary context with minimal info
        # In practice, caller should provide full context
        context = FailureContext(
            run_id=str(uuid.uuid4()),
            test_name="unknown",
            test_file="unknown",
            error_type="unknown",
            error_message="unknown",
            stack_trace="",
        )

        # Use IterativeRefiner for verification
        original_runs = self.refiner.verification_runs
        original_parallel = self.refiner.parallel_runs

        try:
            self.refiner.verification_runs = n_workers
            self.refiner.parallel_runs = min(n_workers, 4)

            result = await self.refiner._verify_determinism(fix, context)
            return result
        finally:
            self.refiner.verification_runs = original_runs
            self.refiner.parallel_runs = original_parallel

    async def cancel_remaining(self, except_winner: str) -> int:
        """Cancel all remaining verification tasks except the winner.

        Args:
            except_winner: Fix ID to exclude from cancellation.

        Returns:
            Number of tasks cancelled.
        """
        cancelled_count = 0

        async with self._lock:
            for fix_id, task in list(self._active_tasks.items()):
                if fix_id == except_winner:
                    continue

                if not task.done() and fix_id not in self._cancelled_ids:
                    task.cancel()
                    self._cancelled_ids.add(fix_id)
                    cancelled_count += 1
                    logger.debug("Cancelled task for fix: %s", fix_id)

        # Wait briefly for cancellations to complete
        if cancelled_count > 0:
            await asyncio.sleep(0.1)

        return cancelled_count

    async def _verify_candidate(
        self,
        fix: FixHypothesis,
        context: FailureContext,
    ) -> tuple[str, VerificationResult]:
        """Verify a single candidate fix.

        Creates a temp sandbox, applies the fix, and runs verification.

        Args:
            fix: Fix to verify.
            context: Failure context.

        Returns:
            Tuple of (fix_id, verification_result).
        """
        work_dir = None

        try:
            # Check if we've been cancelled
            if fix.id in self._cancelled_ids or self._winner_found.is_set():
                return fix.id, VerificationResult(
                    total_runs=0,
                    passed_runs=0,
                    failed_runs=0,
                    is_deterministic=False,
                    failure_patterns=["Cancelled"],
                    run_details=[],
                )

            # Create isolated work directory
            work_dir = Path(tempfile.mkdtemp(prefix=f"hikaflow_race_{fix.id[:8]}_"))
            self._work_dirs.append(work_dir)

            logger.debug("Verifying fix %s in %s", fix.id, work_dir)

            # Run verification using IterativeRefiner
            result = await self.refiner._verify_determinism(fix, context)

            return fix.id, result

        except asyncio.CancelledError:
            logger.debug("Verification cancelled for fix: %s", fix.id)
            return fix.id, VerificationResult(
                total_runs=0,
                passed_runs=0,
                failed_runs=0,
                is_deterministic=False,
                failure_patterns=["Cancelled"],
                run_details=[],
            )

        except Exception as e:
            logger.error("Verification failed for fix %s: %s", fix.id, e)
            return fix.id, VerificationResult(
                total_runs=self.verification_runs,
                passed_runs=0,
                failed_runs=self.verification_runs,
                is_deterministic=False,
                failure_patterns=[f"Error: {type(e).__name__}: {e}"],
                run_details=[],
            )

        finally:
            # Cleanup work dir if configured
            if self.cleanup_on_completion and work_dir and work_dir.exists():
                try:
                    shutil.rmtree(work_dir, ignore_errors=True)
                    if work_dir in self._work_dirs:
                        self._work_dirs.remove(work_dir)
                except Exception as e:
                    logger.warning("Failed to cleanup %s: %s", work_dir, e)

    def _is_winner(self, result: VerificationResult) -> bool:
        """Check if a verification result qualifies as a winner.

        A winner must:
        1. Have run the required number of verification runs
        2. Pass all runs (is_deterministic = True)

        Args:
            result: Verification result to check.

        Returns:
            True if this is a winning result.
        """
        return (
            result.is_deterministic and
            result.total_runs >= self.verification_runs and
            result.passed_runs == result.total_runs
        )

    def _find_best_partial(
        self,
        candidates: List[FixHypothesis],
        results: Dict[str, VerificationResult],
    ) -> tuple[Optional[FixHypothesis], Optional[VerificationResult]]:
        """Find the best partial result when no winner is found.

        Args:
            candidates: Original list of fix candidates.
            results: Dict of fix IDs to verification results.

        Returns:
            Tuple of (best_fix, best_result) or (None, None).
        """
        best_fix: Optional[FixHypothesis] = None
        best_result: Optional[VerificationResult] = None
        best_score = -1.0

        for fix in candidates:
            result = results.get(fix.id)
            if result and result.total_runs > 0:
                # Score based on pass rate and total runs
                pass_rate = result.passed_runs / result.total_runs
                completeness = result.total_runs / self.verification_runs
                score = pass_rate * 0.8 + completeness * 0.2

                if score > best_score:
                    best_score = score
                    best_fix = fix
                    best_result = result

        return best_fix, best_result

    async def _cleanup_work_dirs(self) -> None:
        """Clean up all work directories."""
        for work_dir in list(self._work_dirs):
            if work_dir.exists():
                try:
                    shutil.rmtree(work_dir, ignore_errors=True)
                except Exception as e:
                    logger.warning("Failed to cleanup %s: %s", work_dir, e)

        self._work_dirs.clear()

    async def __aenter__(self) -> SpeculativeFixer:
        """Async context manager entry."""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Async context manager exit with cleanup."""
        await self.cancel_remaining(except_winner="")
        await self._cleanup_work_dirs()


class AdaptiveRacer(SpeculativeFixer):
    """Adaptive racing strategy that adjusts based on early results.

    Extends SpeculativeFixer with:
    - Early termination when clearly better candidates emerge
    - Resource reallocation from failing candidates to promising ones
    - Progressive verification (start with fewer runs, increase for promising)
    """

    def __init__(
        self,
        refiner: Optional[IterativeRefiner] = None,
        max_candidates: int = DEFAULT_MAX_CANDIDATES,
        verification_runs: int = DEFAULT_VERIFICATION_RUNS,
        race_timeout: float = DEFAULT_RACE_TIMEOUT,
        initial_runs: int = 3,
        pass_threshold: float = 0.8,
    ):
        """Initialize adaptive racer.

        Args:
            refiner: IterativeRefiner instance.
            max_candidates: Maximum candidates to race.
            verification_runs: Full verification runs for winners.
            race_timeout: Maximum race time.
            initial_runs: Initial quick verification runs.
            pass_threshold: Minimum pass rate to continue.
        """
        super().__init__(
            refiner=refiner,
            max_candidates=max_candidates,
            verification_runs=verification_runs,
            race_timeout=race_timeout,
        )
        self.initial_runs = initial_runs
        self.pass_threshold = pass_threshold

    async def apply_and_race(
        self,
        fixes: List[FixHypothesis],
        context: FailureContext,
    ) -> RaceResult:
        """Race with adaptive strategy.

        1. Quick initial verification (3 runs)
        2. Filter out candidates with poor pass rate
        3. Full verification on remaining candidates
        4. Return winner or best partial

        Args:
            fixes: List of fix hypotheses.
            context: Failure context.

        Returns:
            RaceResult with winner or best partial.
        """
        start_time = time.time()

        # Reset state
        self._active_tasks.clear()
        self._cancelled_ids.clear()
        self._winner_found.clear()
        self._winner_id = None

        candidates = fixes[:self.max_candidates]
        if not candidates:
            return RaceResult(race_duration=0.0)

        # Ensure IDs
        for fix in candidates:
            if not fix.id:
                fix.id = str(uuid.uuid4())

        all_results: Dict[str, VerificationResult] = {}

        logger.info(
            "Starting adaptive race: %d candidates, %d initial runs, %.0f%% threshold",
            len(candidates),
            self.initial_runs,
            self.pass_threshold * 100,
        )

        try:
            # Phase 1: Quick initial verification
            initial_results = await self._quick_verify_all(candidates, context)

            # Filter candidates that pass threshold
            promising: List[FixHypothesis] = []
            for fix in candidates:
                result = initial_results.get(fix.id)
                if result and result.total_runs > 0:
                    pass_rate = result.passed_runs / result.total_runs
                    if pass_rate >= self.pass_threshold:
                        promising.append(fix)
                        logger.debug(
                            "Fix %s passed initial filter (%.0f%%)",
                            fix.id[:8],
                            pass_rate * 100,
                        )
                    else:
                        logger.debug(
                            "Fix %s eliminated (%.0f%% < %.0f%%)",
                            fix.id[:8],
                            pass_rate * 100,
                            self.pass_threshold * 100,
                        )

            # Phase 2: Full verification on promising candidates
            if promising:
                logger.info(
                    "Phase 2: Full verification on %d promising candidates",
                    len(promising),
                )

                # Use parent's race implementation for full verification
                full_result = await super().apply_and_race(promising, context)

                # Merge results
                all_results.update(initial_results)
                all_results.update(full_result.all_results)

                return RaceResult(
                    winner=full_result.winner,
                    winner_result=full_result.winner_result,
                    partial_fix=full_result.partial_fix,
                    partial_result=full_result.partial_result,
                    all_results=all_results,
                    race_duration=time.time() - start_time,
                    cancelled_count=full_result.cancelled_count + (
                        len(candidates) - len(promising)
                    ),
                )
            else:
                # No promising candidates, return best from initial
                logger.warning("No candidates passed initial filter")
                all_results.update(initial_results)

        except Exception as e:
            logger.exception("Adaptive race failed: %s", e)

        finally:
            if self.cleanup_on_completion:
                await self._cleanup_work_dirs()

        # Return best partial
        best_fix, best_result = self._find_best_partial(candidates, all_results)

        return RaceResult(
            winner=None,
            winner_result=None,
            partial_fix=best_fix,
            partial_result=best_result,
            all_results=all_results,
            race_duration=time.time() - start_time,
            cancelled_count=len(self._cancelled_ids),
        )

    async def _quick_verify_all(
        self,
        candidates: List[FixHypothesis],
        context: FailureContext,
    ) -> Dict[str, VerificationResult]:
        """Run quick initial verification on all candidates.

        Args:
            candidates: Fixes to verify.
            context: Failure context.

        Returns:
            Dict of fix IDs to initial verification results.
        """
        # Temporarily reduce verification runs
        original_runs = self.refiner.verification_runs
        self.refiner.verification_runs = self.initial_runs

        results: Dict[str, VerificationResult] = {}

        try:
            # Run all in parallel
            tasks = [
                self._verify_candidate(fix, context)
                for fix in candidates
            ]

            completed = await asyncio.gather(*tasks, return_exceptions=True)

            for item in completed:
                if isinstance(item, tuple):
                    fix_id, result = item
                    results[fix_id] = result
                elif isinstance(item, Exception):
                    logger.warning("Quick verify failed: %s", item)

        finally:
            self.refiner.verification_runs = original_runs

        return results
