from __future__ import annotations
import os
import time
import uuid
import threading
from concurrent.futures import ThreadPoolExecutor
from enum import Enum
from typing import Dict, Optional, Any
from pydantic import BaseModel

from codestar.core.schemas import JudgeResult, JudgeStatus
from codestar.core.engine import JudgeEngine
from codestar.problems.loader import repo
from codestar.server.db import create_submission, calculate_beats


class TaskStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


class JudgeTaskInfo(BaseModel):
    task_id: str
    problem_slug: str
    status: TaskStatus
    created_at: float
    started_at: Optional[float] = None
    completed_at: Optional[float] = None
    result: Optional[JudgeResult] = None
    error: Optional[str] = None


class JudgeTaskQueue:
    def __init__(self, max_workers: Optional[int] = None, max_tasks: int = 500):
        workers = max_workers or int(os.environ.get("JUDGE_WORKERS", 4))
        self._executor = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="judge-worker")
        self._tasks: Dict[str, JudgeTaskInfo] = {}
        self._max_tasks = max_tasks
        self._lock = threading.Lock()
        self._engine = JudgeEngine()

    def submit_task(
        self,
        problem_slug: str,
        code: str,
        is_test: bool = False,
        custom_testcase: Optional[str] = None,
        user_id: Optional[int] = None,
    ) -> str:
        task_id = f"task_{uuid.uuid4().hex[:12]}"
        now = time.time()
        task_info = JudgeTaskInfo(
            task_id=task_id,
            problem_slug=problem_slug,
            status=TaskStatus.PENDING,
            created_at=now,
        )

        with self._lock:
            self._cleanup_old_tasks()
            self._tasks[task_id] = task_info

        self._executor.submit(
            self._execute_task,
            task_id=task_id,
            problem_slug=problem_slug,
            code=code,
            is_test=is_test,
            custom_testcase=custom_testcase,
            user_id=user_id,
        )
        return task_id

    def _execute_task(
        self,
        task_id: str,
        problem_slug: str,
        code: str,
        is_test: bool,
        custom_testcase: Optional[str],
        user_id: Optional[int],
    ):
        with self._lock:
            if task_id not in self._tasks:
                return
            self._tasks[task_id].status = TaskStatus.RUNNING
            self._tasks[task_id].started_at = time.time()

        try:
            problem = repo.get(problem_slug)
            if not problem:
                with self._lock:
                    self._tasks[task_id].status = TaskStatus.FAILED
                    self._tasks[task_id].error = f"Problem '{problem_slug}' not found"
                    self._tasks[task_id].completed_at = time.time()
                return

            result = self._engine.judge(problem, code, custom_testcase=custom_testcase)

            # Calculate beats for AC submissions
            if result.status == JudgeStatus.AC:
                beats = calculate_beats(problem_slug, result.execution_time_ms, result.peak_memory_mb)
                result.beats_runtime_pct = beats["runtime_pct"]
                result.beats_memory_pct = beats["memory_pct"]

            # Persist submission if not test
            if not is_test:
                passed_cases = sum(1 for c in result.cases if c.passed)
                total_cases = len(result.cases)
                status_str = result.status.value if hasattr(result.status, "value") else str(result.status)
                err_msg = result.message if result.status != JudgeStatus.AC else None
                try:
                    create_submission(
                        user_id=user_id,
                        problem_slug=problem_slug,
                        status=status_str,
                        passed_cases=passed_cases,
                        total_cases=total_cases,
                        runtime_ms=result.execution_time_ms,
                        code=code,
                        error_message=err_msg,
                        peak_memory_mb=result.peak_memory_mb,
                    )
                except Exception:
                    pass

            with self._lock:
                self._tasks[task_id].status = TaskStatus.COMPLETED
                self._tasks[task_id].result = result
                self._tasks[task_id].completed_at = time.time()

        except Exception as e:
            with self._lock:
                self._tasks[task_id].status = TaskStatus.FAILED
                self._tasks[task_id].error = str(e)
                self._tasks[task_id].completed_at = time.time()

    def get_task(self, task_id: str) -> Optional[JudgeTaskInfo]:
        with self._lock:
            return self._tasks.get(task_id)

    def _cleanup_old_tasks(self, max_age_seconds: float = 3600.0):
        """Clean up tasks completed more than an hour ago, or evict oldest completed tasks if exceeding max capacity."""
        now = time.time()
        # 1. Clean up expired tasks
        to_delete = [
            tid for tid, info in self._tasks.items()
            if info.completed_at and (now - info.completed_at > max_age_seconds)
        ]
        for tid in to_delete:
            del self._tasks[tid]

        # 2. Strict capacity bounding: if still exceeding max_tasks, evict oldest completed
        if len(self._tasks) > self._max_tasks:
            completed_tasks = sorted(
                [info for info in self._tasks.values() if info.completed_at],
                key=lambda x: x.completed_at or 0.0,
            )
            excess = len(self._tasks) - self._max_tasks
            for info in completed_tasks[:excess]:
                self._tasks.pop(info.task_id, None)


# Global singleton instance
judge_queue = JudgeTaskQueue()