import time
from dataclasses import dataclass, field
from enum import Enum
from multiprocessing import Queue
from multiprocessing import forkserver as _forkserver
from multiprocessing.process import BaseProcess
from threading import RLock, Thread
from typing import Any, Dict, List, Optional
from uuid import uuid4

from abstra_internals.controllers.execution.executor_config import (
    ExecutorConfig,
)
from abstra_internals.controllers.execution.executor_metrics import MetricsCollector
from abstra_internals.controllers.execution.executor_process import (
    ExecuteRequest,
    ExecutorCommand,
    ExecutorResponse,
    RabbitMQParams,
    ShutdownRequest,
    WarmupRequest,
    executor_main,
)
from abstra_internals.logger import AbstraLogger
from abstra_internals.repositories.multiprocessing import MPContext
from abstra_internals.repositories.project.project import StageType, StageWithFile
from abstra_internals.utils.multiprocessing import safe_multiprocessing_queue


class ExecutorStatus(str, Enum):
    SPAWNING = "spawning"
    WARMING_UP = "warming_up"
    IDLE = "idle"
    BUSY = "busy"
    DEAD = "dead"


@dataclass
class ExecutorHandle:
    executor_id: str
    process: BaseProcess
    work_queue: Queue
    response_queue: Queue
    status: ExecutorStatus
    pid: Optional[int] = None
    is_form_reserved: bool = True
    executions_completed: int = 0
    current_execution_id: Optional[str] = None
    spawn_time: float = field(default_factory=time.time)
    warmup_start_time: Optional[float] = None
    warmup_complete_time: Optional[float] = None
    last_execution_start: Optional[float] = None

    def is_alive(self) -> bool:
        return self.process.is_alive()

    def age_seconds(self) -> float:
        return time.time() - self.spawn_time

    def is_ready(self) -> bool:
        return self.status == ExecutorStatus.IDLE and self.is_alive()

    def mark_busy(self, execution_id: str) -> None:
        self.status = ExecutorStatus.BUSY
        self.current_execution_id = execution_id
        self.last_execution_start = time.time()

    def mark_idle(self) -> None:
        self.status = ExecutorStatus.IDLE
        self.current_execution_id = None
        self.executions_completed += 1

    def mark_dead(self) -> None:
        self.status = ExecutorStatus.DEAD
        self.current_execution_id = None

    def kill(self) -> None:
        if self.is_alive():
            try:
                self.process.terminate()
                self.process.join(timeout=2.0)
                if self.process.is_alive():
                    self.process.kill()
                    self.process.join(timeout=1.0)
            except Exception as e:
                AbstraLogger.error(f"[ExecutorPool] Error killing executor: {e}")
        self.mark_dead()


class ExecutorPool:
    def __init__(
        self,
        config: ExecutorConfig,
        mp_context: MPContext,
        root_path: str,
        server_port: int,
        worker_id: str,
        parent_executions_queue: Optional[Queue] = None,
        verbose: bool = True,
    ):
        self.config = config
        self.mp_context = mp_context
        self.root_path = root_path
        self.server_port = server_port
        self.worker_id = worker_id
        self.parent_executions_queue = parent_executions_queue
        self.verbose = verbose

        self.executors: Dict[str, ExecutorHandle] = {}
        self.execution_to_executor: Dict[str, str] = {}
        self._pending_spawns: List[tuple] = []

        self.lock = RLock()

        self.shutdown_requested = False

        self.metrics = MetricsCollector()

    def start(self) -> None:
        if self.verbose:
            AbstraLogger.warning(
                f"[ExecutorPool] Starting {self.config.total_executors} executors "
                f"with warmup_parallelism={self.config.warmup_parallelism}"
            )

        try:
            _forkserver.set_forkserver_preload(
                ["abstra_internals.controllers.execution.executor_process"]
            )
            if self.verbose:
                AbstraLogger.warning(
                    "[ExecutorPool] Forkserver preload configured for executor_process"
                )
        except Exception as e:
            AbstraLogger.warning(
                f"[ExecutorPool] Could not configure forkserver preload (non-forkserver context?): {e}"
            )

        self._pending_spawns: List[tuple] = [
            (i, i < self.config.min_form_executors)
            for i in range(self.config.total_executors)
        ]

        self.monitor_thread = Thread(target=self._monitor_loop, daemon=True)
        self.monitor_thread.start()

    def shutdown(self) -> None:
        if self.verbose:
            AbstraLogger.warning("[ExecutorPool] Shutting down pool")
        self.shutdown_requested = True

        with self.lock:
            for executor in list(
                self.executors.values()
            ):  # Create snapshot to avoid dictionary modification during iteration
                try:
                    if executor.is_alive():
                        executor.work_queue.put(ShutdownRequest(), timeout=2.0)
                except Exception as e:
                    AbstraLogger.error(
                        f"[ExecutorPool] Error sending shutdown to executor: {e}"
                    )

        shutdown_start = time.time()
        for executor in list(
            self.executors.values()
        ):  # Create snapshot to avoid dictionary modification during iteration
            remaining_time = max(0, 5.0 - (time.time() - shutdown_start))
            executor.process.join(timeout=remaining_time)
            if executor.is_alive():
                executor.kill()

        if self.verbose:
            AbstraLogger.warning("[ExecutorPool] Pool shutdown complete")

    def execute(
        self,
        stage: StageWithFile,
        worker_id: str,
        execution_id: str,
        request,
        connection,
        rabbitmq_params: Optional[RabbitMQParams],
        user_jwt: Optional[str] = None,
    ) -> ExecutorResponse:
        max_wait = self.config.acquire_timeout_seconds
        start_time = time.time()
        retry_count = 0

        while True:
            executor = self._acquire_executor(stage.type_name, execution_id)

            if executor is not None:
                break

            retry_count += 1

            elapsed = time.time() - start_time
            if elapsed >= max_wait:
                with self.lock:
                    status_summary = {}
                    for e in self.executors.values():
                        status = e.status.value
                        status_summary[status] = status_summary.get(status, 0) + 1

                AbstraLogger.error(
                    f"[ExecutorPool] No executor available after {elapsed:.1f}s wait "
                    f"(execution_id={execution_id}, stage_type={stage.type_name}, "
                    f"executor_status={status_summary})"
                )
                raise RuntimeError(
                    f"No available executor for stage type {stage.type_name} after {elapsed:.1f}s wait. "
                    "All executors are busy or dead."
                )

            time.sleep(0.25)

        try:
            self.metrics.record_execution_started(
                execution_id=execution_id,
                stage_type=stage.type_name,
                executor_id=executor.executor_id,
                wait_time_ms=0.0,
            )

            if not executor.is_alive():
                AbstraLogger.error(
                    f"[ExecutorPool] Executor died before work submission "
                    f"(execution_id={execution_id}, executor_id={executor.executor_id}, pid={executor.pid})"
                )
                executor.mark_dead()
                self.metrics.record_executor_died(executor.executor_id)
                raise RuntimeError(
                    f"Executor {executor.executor_id} died before execution"
                )

            execute_request = ExecuteRequest(
                command=ExecutorCommand.EXECUTE,
                worker_id=worker_id,
                execution_id=execution_id,
                stage=stage,
                request=request,
                connection=connection,
                rabbitmq_params=rabbitmq_params,
                user_jwt=user_jwt,
            )

            try:
                executor.work_queue.put(execute_request, timeout=5.0)
            except Exception as e:
                AbstraLogger.error(
                    f"[ExecutorPool] Failed to send work request to executor "
                    f"(execution_id={execution_id}, executor_id={executor.executor_id}, error={e})"
                )

                executor.mark_dead()
                self.metrics.record_executor_died(executor.executor_id)
                raise RuntimeError(
                    f"Failed to send work to executor {executor.executor_id}: {e}"
                )

            timeout = self.config.execution_timeout_seconds

            response = None
            wait_start = time.time()

            while time.time() - wait_start < timeout:
                if not executor.is_alive():
                    AbstraLogger.error(
                        f"[ExecutorPool] Executor died while waiting for response "
                        f"(execution_id={execution_id}, executor_id={executor.executor_id}, pid={executor.pid})"
                    )
                    executor.mark_dead()
                    self.metrics.record_executor_died(executor.executor_id)
                    raise RuntimeError(
                        f"Executor {executor.executor_id} died unexpectedly during execution"
                    )

                try:
                    response = executor.response_queue.get(timeout=1.0)
                    break
                except Exception:  # Empty queue
                    continue

            if response is None:
                AbstraLogger.warning(
                    f"[ExecutorPool] Execution timed out, killing stuck executor "
                    f"(execution_id={execution_id}, executor_id={executor.executor_id}, pid={executor.pid})"
                )
                self._kill_and_replace_executor(executor, reason="timeout")
                raise TimeoutError(f"Execution timed out after {timeout}s")

            if response.execution_id != execution_id:
                raise RuntimeError(
                    f"Received response for wrong execution: expected={execution_id}, got={response.execution_id}"
                )

            execution_time_ms = (time.time() - start_time) * 1000
            self.metrics.record_execution_completed(
                execution_id=execution_id,
                success=response.success,
                execution_time_ms=execution_time_ms,
            )

            return response

        except Exception as e:
            AbstraLogger.error(
                f"[ExecutorPool] Error during execution "
                f"(execution_id={execution_id}): {e}"
            )
            raise

        finally:
            self._release_executor(executor, execution_id)

    def kill_execution(self, execution_id: str) -> bool:
        with self.lock:
            executor_id = self.execution_to_executor.get(execution_id)
            if executor_id is None:
                AbstraLogger.warning(
                    f"[ExecutorPool] Cannot kill execution {execution_id}: "
                    "not found in active executions"
                )
                return False

            executor = self.executors.get(executor_id)
            if executor is None or not executor.is_alive():
                AbstraLogger.warning(
                    f"[ExecutorPool] Cannot kill execution {execution_id}: "
                    "executor is dead"
                )
                self.execution_to_executor.pop(execution_id, None)
                return False

        self._kill_and_replace_executor(executor, reason="killed")

        return True

    def get_stats(self) -> Dict[str, Any]:
        with self.lock:
            executors = list(self.executors.values())
            active_executions_count = len(self.execution_to_executor)
            metrics_summary = self.metrics.get_summary()

        total = len(executors)
        idle = sum(1 for e in executors if e.status == ExecutorStatus.IDLE)
        busy = sum(1 for e in executors if e.status == ExecutorStatus.BUSY)
        warming_up = sum(1 for e in executors if e.status == ExecutorStatus.WARMING_UP)
        dead = sum(1 for e in executors if e.status == ExecutorStatus.DEAD)
        form_capable_idle = sum(
            1
            for e in executors
            if e.status == ExecutorStatus.IDLE and e.is_form_reserved
        )

        warming_times = []
        slow_warmups = 0
        current_time = time.time()

        for e in executors:
            if e.status == ExecutorStatus.WARMING_UP and e.warmup_start_time:
                warmup_elapsed = current_time - e.warmup_start_time
                warming_times.append(warmup_elapsed)
                if warmup_elapsed > 5.0:  # Consider > 5s as slow
                    slow_warmups += 1

        return {
            **metrics_summary,
            "total_executors": total,
            "idle": idle,
            "busy": busy,
            "warming_up": warming_up,
            "dead": dead,
            "form_capable_idle": form_capable_idle,
            "active_executions": active_executions_count,
            "slow_warmups": slow_warmups,
            "avg_warmup_time": sum(warming_times) / len(warming_times)
            if warming_times
            else 0,
            "max_warmup_time": max(warming_times) if warming_times else 0,
        }

    def can_start_loop(self) -> bool:
        with self.lock:
            executors = list(self.executors.values())

        idle_count = sum(1 for e in executors if e.status == ExecutorStatus.IDLE)

        return idle_count > 0

    def _spawn_executor(self, is_form_reserved: bool) -> ExecutorHandle:
        executor_id = str(uuid4())[:8]

        work_queue = safe_multiprocessing_queue(self.mp_context)
        response_queue = safe_multiprocessing_queue(self.mp_context)

        p = self.mp_context.Process(
            target=executor_main,
            name=f"Executor-{executor_id}",
            args=(
                work_queue,
                response_queue,
                self.root_path,
                self.server_port,
                self.parent_executions_queue,
                self.verbose,
            ),
        )

        p.start()

        executor = ExecutorHandle(
            executor_id=executor_id,
            process=p,
            work_queue=work_queue,
            response_queue=response_queue,
            status=ExecutorStatus.SPAWNING,
            pid=p.pid,
            is_form_reserved=is_form_reserved,
        )

        self.metrics.record_executor_spawned(
            executor_id=executor_id,
            spawn_time=time.time(),
            is_form_reserved=is_form_reserved,
        )

        return executor

    def _start_warmup(self, executor: ExecutorHandle) -> None:
        executor.status = ExecutorStatus.WARMING_UP
        executor.warmup_start_time = time.time()

        if not executor.is_alive():
            AbstraLogger.error(
                f"[ExecutorPool] Process died before warmup (executor_id={executor.executor_id})"
            )
            executor.mark_dead()
            self.metrics.record_executor_died(executor.executor_id)
            return

        try:
            executor.work_queue.put(WarmupRequest(), timeout=5.0)
        except Exception as e:
            AbstraLogger.error(
                f"[ExecutorPool] Failed to start warmup (executor_id={executor.executor_id}): {e}"
            )
            executor.mark_dead()
            self.metrics.record_executor_died(executor.executor_id)

    def _acquire_executor(
        self, stage_type: StageType, execution_id: str
    ) -> Optional[ExecutorHandle]:
        with self.lock:
            idle: List[ExecutorHandle] = []
            for e in self.executors.values():
                if e.status != ExecutorStatus.IDLE:
                    continue
                if not e.is_alive():
                    e.mark_dead()
                    self.metrics.record_executor_died(e.executor_id)
                    continue
                idle.append(e)

            if not idle:
                return None

            form_executors = [e for e in idle if e.is_form_reserved]
            non_form_executors = [e for e in idle if not e.is_form_reserved]

            if stage_type == "form":
                pool = form_executors + non_form_executors
            else:
                pool = non_form_executors

            if not pool:
                return None

            executor = pool[0]

            executor.mark_busy(execution_id)
            self.execution_to_executor[execution_id] = executor.executor_id

            return executor

    def _kill_and_replace_executor(
        self, executor: ExecutorHandle, reason: str = "stuck"
    ) -> None:
        with self.lock:
            self.executors.pop(executor.executor_id, None)
            for exec_id, exec_executor_id in list(self.execution_to_executor.items()):
                if exec_executor_id == executor.executor_id:
                    self.execution_to_executor.pop(exec_id, None)

        executor.mark_dead()
        executor.kill()
        self.metrics.record_executor_died(executor.executor_id)

        replacement = self._spawn_executor(is_form_reserved=executor.is_form_reserved)

        with self.lock:
            self.executors[replacement.executor_id] = replacement

        self._start_warmup(replacement)

        AbstraLogger.warning(
            f"[ExecutorPool] Replaced {reason} executor {executor.executor_id} "
            f"with new executor {replacement.executor_id}"
        )

    def _release_executor(self, executor: ExecutorHandle, execution_id: str) -> None:
        should_recycle = False

        with self.lock:
            self.execution_to_executor.pop(execution_id, None)

            if executor.status == ExecutorStatus.DEAD:
                return

            if not executor.is_alive():
                AbstraLogger.warning(
                    f"[ExecutorPool] Executor died during execution "
                    f"(executor_id={executor.executor_id}, pid={executor.pid})"
                )
                executor.mark_dead()
                self.metrics.record_executor_died(executor.executor_id)
                return

            if (
                executor.executions_completed + 1
                >= self.config.max_executions_per_executor
            ):
                should_recycle = True
                self.executors.pop(executor.executor_id, None)
            else:
                executor.mark_idle()

        if should_recycle:
            self._kill_and_replace_executor(executor, reason="recycled")

    def _monitor_loop(self) -> None:
        while not self.shutdown_requested:
            time.sleep(0.1)

            try:
                should_spawn = False
                spawn_info: Optional[tuple] = None

                with self.lock:
                    if self._pending_spawns:
                        warming_count = sum(
                            1
                            for e in self.executors.values()
                            if e.status == ExecutorStatus.WARMING_UP
                        )
                        if warming_count < self.config.warmup_parallelism:
                            spawn_info = self._pending_spawns.pop(0)
                            should_spawn = True

                if should_spawn and spawn_info is not None:
                    idx, is_form_reserved = spawn_info
                    executor = self._spawn_executor(is_form_reserved=is_form_reserved)

                    with self.lock:
                        self.executors[executor.executor_id] = executor

                    if self.verbose:
                        AbstraLogger.warning(
                            f"[ExecutorPool] Spawned executor {idx + 1}/{self.config.total_executors}"
                        )

                    self._start_warmup(executor)

                warming_executors: List[ExecutorHandle] = []
                timed_out_executors: List[ExecutorHandle] = []
                dead_executors: List[ExecutorHandle] = []

                with self.lock:
                    for executor in list(self.executors.values()):
                        if executor.status == ExecutorStatus.WARMING_UP:
                            if executor.current_execution_id is None:
                                warming_executors.append(executor)
                        elif executor.status == ExecutorStatus.DEAD:
                            dead_executors.append(executor)

                for executor in warming_executors:
                    try:
                        response = executor.response_queue.get_nowait()

                        if response.execution_id is not None:
                            executor.response_queue.put(response)
                            continue

                        with self.lock:
                            if response.success:
                                executor.status = ExecutorStatus.IDLE
                                executor.warmup_complete_time = time.time()
                                warmup_duration = executor.warmup_complete_time - (
                                    executor.warmup_start_time or 0
                                )

                                if self.verbose:
                                    AbstraLogger.warning(
                                        f"[ExecutorPool] Warmup completed in {warmup_duration:.3f}s"
                                    )

                                self.metrics.record_executor_warmup(
                                    executor_id=executor.executor_id,
                                    warmup_duration_ms=warmup_duration * 1000,
                                )
                            else:
                                AbstraLogger.error(
                                    f"[ExecutorPool] Warmup failed (executor_id={executor.executor_id})"
                                )
                                executor.mark_dead()
                                self.metrics.record_executor_died(executor.executor_id)
                    except Exception:
                        if executor.warmup_start_time:
                            warmup_elapsed = time.time() - executor.warmup_start_time

                            if warmup_elapsed > self.config.warmup_timeout_seconds:
                                AbstraLogger.error(
                                    f"[ExecutorPool] Warmup timeout (executor_id={executor.executor_id}, "
                                    f"elapsed={warmup_elapsed:.1f}s)"
                                )
                                timed_out_executors.append(executor)

                for executor in timed_out_executors:
                    executor.kill()

                for dead_executor in dead_executors:
                    with self.lock:
                        self.executors.pop(dead_executor.executor_id, None)

                    replacement = self._spawn_executor(
                        is_form_reserved=dead_executor.is_form_reserved
                    )

                    with self.lock:
                        self.executors[replacement.executor_id] = replacement

                    self._start_warmup(replacement)

            except Exception as e:
                AbstraLogger.error(f"[ExecutorPool] Error in monitor loop: {e}")
                AbstraLogger.capture_exception(e)
