"""
Job queue and worker pool for Tako VM.

Provides async job submission, execution, and cancellation.
"""

import asyncio
import logging
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Dict, Optional

from tako_vm.models import (
    DeadLetterEntry,
    ExecutionError,
    ExecutionRecord,
    sha256_content,
    sha256_json,
)
from tako_vm.server.correlation import get_correlation_id, set_correlation_id
from tako_vm.storage import ExecutionStorage

if TYPE_CHECKING:
    from tako_vm.execution.worker import CodeExecutor

logger = logging.getLogger(__name__)


@dataclass
class QueuedJob:
    """Job waiting in or being processed by the queue."""

    job_id: str
    """Unique job identifier."""

    job_data: Dict[str, Any]
    """Job data including code, input_data, etc."""

    client_ip: Optional[str]
    """Client IP address."""

    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    """When the job was queued."""

    future: Optional[asyncio.Future] = field(default=None)
    """Future that will hold the result (set by WorkerPool.submit)."""

    cancelled: bool = False
    """Whether this job has been cancelled."""


class WorkerPool:
    """
    Async worker pool for job execution.

    Uses asyncio.Queue for pending jobs and ThreadPoolExecutor
    for running Docker operations.
    """

    def __init__(
        self,
        executor: "CodeExecutor",
        storage: ExecutionStorage,
        max_workers: int = 4,
        max_queue_size: int = 100,
        queue_wait_timeout: float = 1.0,
    ):
        """
        Initialize worker pool.

        Args:
            executor: CodeExecutor instance for running jobs
            storage: ExecutionStorage for persisting records
            max_workers: Maximum concurrent workers
            max_queue_size: Maximum pending jobs in queue
            queue_wait_timeout: Timeout for queue wait operations in seconds
        """
        self.executor = executor
        self.storage = storage
        self.max_workers = max_workers
        self.max_queue_size = max_queue_size
        self.queue_wait_timeout = queue_wait_timeout

        self._queue: asyncio.Queue[QueuedJob] = asyncio.Queue(maxsize=max_queue_size)
        self._workers: list[asyncio.Task] = []
        self._active_jobs: Dict[str, QueuedJob] = {}
        self._running_jobs: Dict[str, QueuedJob] = {}
        self._jobs_lock: asyncio.Lock = asyncio.Lock()  # Protect job dictionary access
        self._thread_pool: Optional[ThreadPoolExecutor] = None
        self._shutdown = False
        self._started = False

    async def start(self) -> None:
        """Start worker tasks."""
        if self._started:
            return

        self._shutdown = False
        self._thread_pool = ThreadPoolExecutor(max_workers=self.max_workers)

        for i in range(self.max_workers):
            worker = asyncio.create_task(self._worker_loop(i), name=f"worker-{i}")
            self._workers.append(worker)

        self._started = True
        logger.info(f"Started worker pool with {self.max_workers} workers")

    async def stop(self, timeout: float = 30.0) -> None:
        """
        Graceful shutdown.

        Args:
            timeout: Maximum time to wait for workers to finish
        """
        if not self._started:
            return

        self._shutdown = True

        # Cancel all pending jobs
        while not self._queue.empty():
            try:
                job = self._queue.get_nowait()
                if job.future and not job.future.done():
                    job.future.cancel()
            except asyncio.QueueEmpty:
                break

        # Wait for workers with timeout
        if self._workers:
            done, pending = await asyncio.wait(
                self._workers, timeout=timeout, return_when=asyncio.ALL_COMPLETED
            )

            # Cancel any still running
            for task in pending:
                task.cancel()

        # Shutdown thread pool - wait for running jobs to complete
        if self._thread_pool:
            self._thread_pool.shutdown(wait=True, cancel_futures=True)
            self._thread_pool = None

        self._workers.clear()
        self._active_jobs.clear()
        self._running_jobs.clear()
        self._started = False

        logger.info("Worker pool stopped")

    async def submit(self, job_data: Dict[str, Any], client_ip: Optional[str] = None) -> str:
        """
        Submit job to queue.

        Args:
            job_data: Job data dictionary
            client_ip: Client IP address

        Returns:
            Job ID

        Raises:
            RuntimeError: If queue is full
        """
        job_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc)

        # Create future in async context where event loop is running
        loop = asyncio.get_running_loop()

        job = QueuedJob(
            job_id=job_id,
            job_data=job_data,
            client_ip=client_ip,
            future=loop.create_future(),
        )

        # Create preliminary "queued" record for idempotency tracking
        # This ensures concurrent requests with same idempotency_key see the record
        job_type_name = job_data.get("job_type") or "default"
        queued_record = ExecutionRecord(
            execution_id=job_id,
            status="queued",
            job_type=job_type_name,
            job_ref=f"{job_type_name}@latest",
            created_at=now,
            queued_at=now,
            code_hash=sha256_content(job_data.get("code", "")),
            input_hash=sha256_json(job_data.get("input_data", {})),
            client_ip=client_ip,
            idempotency_key=job_data.get("idempotency_key"),
            idempotency_fingerprint=job_data.get("idempotency_fingerprint"),
            parent_execution_id=job_data.get("parent_execution_id"),
            relationship=job_data.get("relationship"),
        )
        if self._queue.full():
            raise RuntimeError("Job queue is full, try again later")

        await self.storage.save_record(queued_record)

        try:
            self._queue.put_nowait(job)
        except asyncio.QueueFull as exc:
            queued_record.status = "failed"
            queued_record.error = ExecutionError(
                type="service_unavailable", message="Job queue is full, try again later"
            )
            await self.storage.save_record(queued_record)
            raise RuntimeError("Job queue is full, try again later") from exc

        async with self._jobs_lock:
            self._active_jobs[job_id] = job

        logger.info(f"Job {job_id} queued (queue size: {self._queue.qsize()})")

        return job_id

    async def wait_for_result(
        self, job_id: str, timeout: Optional[float] = None
    ) -> ExecutionRecord:
        """
        Wait for job completion.

        Args:
            job_id: Job ID to wait for
            timeout: Maximum time to wait (None = no timeout)

        Returns:
            ExecutionRecord with results

        Raises:
            KeyError: If job not found
            asyncio.TimeoutError: If timeout exceeded
            asyncio.CancelledError: If job was cancelled
        """
        # Get job reference with lock protection
        async with self._jobs_lock:
            job = self._active_jobs.get(job_id) or self._running_jobs.get(job_id)

        if not job:
            # Check if already completed in storage
            record = await self.storage.get_record(job_id)
            if record:
                return record
            raise KeyError(f"Job {job_id} not found")

        if not job.future:
            raise KeyError(f"Job {job_id} has no future (internal error)")

        if timeout:
            return await asyncio.wait_for(job.future, timeout=timeout)
        else:
            return await job.future

    async def get_job_status(self, job_id: str) -> Optional[dict]:
        """
        Get current job status.

        Args:
            job_id: Job ID to check

        Returns:
            Status dict or None if not found
        """
        # Check active/running jobs with lock protection
        async with self._jobs_lock:
            if job_id in self._running_jobs:
                return {
                    "job_id": job_id,
                    "status": "running",
                    "created_at": self._running_jobs[job_id].created_at.isoformat(),
                }

            if job_id in self._active_jobs:
                job = self._active_jobs[job_id]
                if job.future and job.future.done():
                    return {
                        "job_id": job_id,
                        "status": "completed",
                        "created_at": job.created_at.isoformat(),
                    }
                queue_position = self._estimate_queue_position_unlocked(job_id)
                return {
                    "job_id": job_id,
                    "status": "pending",
                    "created_at": job.created_at.isoformat(),
                    "queue_position": queue_position,
                }

        # Check storage for completed jobs (outside lock - storage has its own lock)
        record = await self.storage.get_record(job_id)
        if record:
            return {
                "job_id": job_id,
                "status": record.status,
                "created_at": record.created_at.isoformat(),
                "duration_ms": record.duration_ms,
            }

        return None

    async def cancel(self, job_id: str) -> bool:
        """
        Cancel a queued or running job.

        Args:
            job_id: Job ID to cancel

        Returns:
            True if job was found and cancelled
        """
        async with self._jobs_lock:
            # Check pending jobs
            job = self._active_jobs.get(job_id)
            if job and job.future and not job.future.done():
                job.cancelled = True
                job.future.cancel()
                logger.info(f"Job {job_id} cancelled (was pending)")
                return True

            # Check running jobs - can't cancel Docker execution directly
            # but we can mark it for cleanup
            job = self._running_jobs.get(job_id)
            if job:
                job.cancelled = True
                logger.info(f"Job {job_id} marked for cancellation (was running)")
                return True

        return False

    def _estimate_queue_position_unlocked(self, job_id: str) -> int:
        """Estimate position in queue (rough estimate). Must be called with _jobs_lock held."""
        position = 1
        for jid in self._active_jobs:
            if jid == job_id:
                break
            if jid not in self._running_jobs:
                position += 1
        return position

    async def get_stats(self) -> dict:
        """Get current pool statistics."""
        async with self._jobs_lock:
            running_count = len(self._running_jobs)
        return {
            "pending": self._queue.qsize(),
            "running": running_count,
            "max_workers": self.max_workers,
            "max_queue_size": self.max_queue_size,
        }

    @property
    def stats(self) -> dict:
        """Get current pool statistics (sync version for backward compatibility).

        DEPRECATED: This property doesn't use lock protection for running count,
        which can cause race conditions. Use get_stats() for async code paths.

        Returns:
            Dict with pending, running, max_workers, max_queue_size
        """
        import warnings

        warnings.warn(
            "WorkerPool.stats property is deprecated due to race condition. "
            "Use 'await worker_pool.get_stats()' instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return {
            "pending": self._queue.qsize(),
            "running": len(self._running_jobs),
            "max_workers": self.max_workers,
            "max_queue_size": self.max_queue_size,
        }

    async def _worker_loop(self, worker_id: int) -> None:
        """
        Worker coroutine that processes jobs.

        Args:
            worker_id: Worker identifier for logging
        """
        logger.info(f"Worker {worker_id} started")

        while not self._shutdown:
            # Initialize job to None for safe error recovery
            job: Optional[QueuedJob] = None
            try:
                # Wait for job with timeout so we can check shutdown flag
                try:
                    job = await asyncio.wait_for(self._queue.get(), timeout=self.queue_wait_timeout)
                except asyncio.TimeoutError:
                    continue

                # Skip if already cancelled
                if job.cancelled or (job.future and job.future.cancelled()):
                    async with self._jobs_lock:
                        self._active_jobs.pop(job.job_id, None)
                    continue

                # Move to running
                async with self._jobs_lock:
                    self._running_jobs[job.job_id] = job

                # Set correlation ID from job data for logging
                correlation_id = job.job_data.get("correlation_id")
                if correlation_id:
                    set_correlation_id(correlation_id)

                logger.info(f"Worker {worker_id} executing job {job.job_id}")

                try:
                    # Run execution in thread pool
                    record = await self._execute_job(job)

                    # Save to storage
                    await self.storage.save_record(record)

                    # Set result (with race condition protection)
                    if job.future:
                        try:
                            job.future.set_result(record)
                        except asyncio.InvalidStateError:
                            logger.debug(f"Future already done for job {job.job_id}")

                except asyncio.CancelledError:
                    # Job was cancelled
                    job_type_name = job.job_data.get("job_type") or "default"
                    record = ExecutionRecord(
                        execution_id=job.job_id,
                        status="cancelled",
                        job_type=job_type_name,
                        job_ref=f"{job_type_name}@latest",
                        created_at=job.created_at,
                        queued_at=job.created_at,
                        code_hash=sha256_content(job.job_data.get("code", "")),
                        input_hash=sha256_json(job.job_data.get("input_data", {})),
                        client_ip=job.client_ip,
                        # Propagate idempotency and lineage fields
                        idempotency_key=job.job_data.get("idempotency_key"),
                        idempotency_fingerprint=job.job_data.get("idempotency_fingerprint"),
                        parent_execution_id=job.job_data.get("parent_execution_id"),
                        relationship=job.job_data.get("relationship"),
                    )
                    await self.storage.save_record(record)

                    if job.future:
                        try:
                            job.future.set_result(record)
                        except asyncio.InvalidStateError:
                            logger.debug(f"Future already done for cancelled job {job.job_id}")

                except Exception as e:
                    logger.error(
                        f"Worker {worker_id} error on job {job.job_id}: {e}", exc_info=True
                    )

                    # Always create an ExecutionRecord for audit trail (P0 fix)
                    from tako_vm.security import sanitize_error

                    error_type = "internal_error"
                    error_msg = sanitize_error(str(e))
                    job_type_name = job.job_data.get("job_type") or "default"

                    record = ExecutionRecord(
                        execution_id=job.job_id,
                        status="failed",
                        job_type=job_type_name,
                        job_ref=f"{job_type_name}@latest",
                        created_at=job.created_at,
                        queued_at=job.created_at,
                        code_hash=sha256_content(job.job_data.get("code", "")),
                        input_hash=sha256_json(job.job_data.get("input_data", {})),
                        client_ip=job.client_ip,
                        error=ExecutionError(type=error_type, message=error_msg),
                        # Propagate idempotency and lineage fields
                        idempotency_key=job.job_data.get("idempotency_key"),
                        idempotency_fingerprint=job.job_data.get("idempotency_fingerprint"),
                        parent_execution_id=job.job_data.get("parent_execution_id"),
                        relationship=job.job_data.get("relationship"),
                    )
                    await self.storage.save_record(record)

                    # Add to dead letter queue for internal errors (P3 fix)
                    try:
                        dlq_entry = DeadLetterEntry(
                            job_id=job.job_id,
                            job_data=job.job_data,
                            error_type=error_type,
                            error_message=error_msg,
                            retry_count=0,
                            client_ip=job.client_ip,
                            correlation_id=get_correlation_id(),
                        )
                        await self.storage.add_to_dlq(dlq_entry)
                        logger.info(f"Job {job.job_id} added to dead letter queue")
                    except Exception as dlq_err:
                        logger.warning(f"Failed to add job {job.job_id} to DLQ: {dlq_err}")

                    if job.future:
                        try:
                            job.future.set_result(record)
                        except asyncio.InvalidStateError:
                            logger.debug(f"Future already done for failed job {job.job_id}")

                finally:
                    # Cleanup with lock protection
                    async with self._jobs_lock:
                        self._active_jobs.pop(job.job_id, None)
                        self._running_jobs.pop(job.job_id, None)

            except Exception as e:
                logger.error(f"Worker {worker_id} unexpected error: {e}", exc_info=True)
                # Ensure job future is resolved even on unexpected errors
                # to prevent clients from hanging indefinitely
                if job is not None and job.future:
                    try:
                        from tako_vm.security import sanitize_error

                        error_record = ExecutionRecord(
                            execution_id=job.job_id,
                            status="failed",
                            job_type=job.job_data.get("job_type") or "default",
                            job_ref=f"{job.job_data.get('job_type') or 'default'}@latest",
                            created_at=job.created_at,
                            queued_at=job.created_at,
                            code_hash=sha256_content(job.job_data.get("code", "")),
                            input_hash=sha256_json(job.job_data.get("input_data", {})),
                            client_ip=job.client_ip,
                            error=ExecutionError(
                                type="internal_error", message=sanitize_error(str(e))
                            ),
                        )
                        await self.storage.save_record(error_record)
                        job.future.set_result(error_record)
                    except Exception as recovery_err:
                        logger.error(
                            f"Failed to recover from error for job {job.job_id}: {recovery_err}"
                        )
                        try:
                            job.future.cancel()
                        except Exception:
                            pass
                await asyncio.sleep(1)  # Prevent tight loop on errors

        logger.info(f"Worker {worker_id} stopped")

    async def _execute_job(self, job: QueuedJob) -> ExecutionRecord:
        """
        Execute job in thread pool with timeout protection.

        Args:
            job: Job to execute

        Returns:
            ExecutionRecord with results

        Raises:
            asyncio.TimeoutError: If execution exceeds timeout + buffer
        """
        loop = asyncio.get_running_loop()

        # Get job timeout with buffer for Docker startup overhead
        job_timeout = job.job_data.get("timeout", 30)
        executor_timeout = job_timeout + 60  # 60s buffer for container setup

        # Run synchronous execution in thread pool with timeout protection
        record = await asyncio.wait_for(
            loop.run_in_executor(self._thread_pool, self._run_job_sync, job),
            timeout=executor_timeout,
        )

        return record

    def _run_job_sync(self, job: QueuedJob) -> ExecutionRecord:
        """
        Synchronously execute job (runs in thread pool).

        Args:
            job: Job to execute

        Returns:
            ExecutionRecord with results
        """
        return self.executor.execute_job_with_record(
            job_id=job.job_id,
            job=job.job_data,
            client_ip=job.client_ip,
        )
