import threading
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Callable, Dict, List, Optional


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


@dataclass
class DatasetJob:
    job_id: str
    name: str
    project: str
    status: JobStatus = JobStatus.PENDING
    created_at: str = ""
    completed_at: Optional[str] = None
    error: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "job_id": self.job_id,
            "dataset_name": self.name,
            "project": self.project,
            "status": self.status.value,
            "created_at": self.created_at,
            "completed_at": self.completed_at,
            "error": self.error,
        }


class DatasetJobManager:
    """Thread-safe in-memory job manager for dataset creation jobs."""

    def __init__(self, max_workers: int = 2):
        self._jobs: Dict[str, DatasetJob] = {}
        self._lock = threading.Lock()
        self._executor = ThreadPoolExecutor(
            max_workers=max_workers, thread_name_prefix="dataset-job"
        )

    def submit_job(
        self,
        name: str,
        project: str,
        task_fn: Callable[[DatasetJob], None],
        metadata: Optional[Dict[str, Any]] = None,
    ) -> DatasetJob:
        job_id = str(uuid.uuid4())
        job = DatasetJob(
            job_id=job_id,
            name=name,
            project=project,
            status=JobStatus.PENDING,
            created_at=datetime.now(timezone.utc).isoformat(),
            metadata=metadata or {},
        )

        with self._lock:
            self._jobs[job_id] = job

        self._executor.submit(self._run_job, job, task_fn)
        return job

    def _run_job(self, job: DatasetJob, task_fn: Callable[[DatasetJob], None]):
        with self._lock:
            job.status = JobStatus.RUNNING

        try:
            task_fn(job)
            with self._lock:
                job.status = JobStatus.COMPLETED
                job.completed_at = datetime.now(timezone.utc).isoformat()
        except Exception as e:
            with self._lock:
                job.status = JobStatus.FAILED
                job.completed_at = datetime.now(timezone.utc).isoformat()
                job.error = str(e)

    def get_job(self, job_id: str) -> Optional[DatasetJob]:
        with self._lock:
            return self._jobs.get(job_id)

    def list_jobs(
        self,
        project: Optional[str] = None,
        status: Optional[JobStatus] = None,
    ) -> List[DatasetJob]:
        with self._lock:
            jobs = list(self._jobs.values())

        if project:
            jobs = [j for j in jobs if j.project == project]
        if status:
            jobs = [j for j in jobs if j.status == status]

        return sorted(jobs, key=lambda j: j.created_at, reverse=True)

    def cleanup_completed(self, max_age_seconds: int = 3600):
        """Remove completed/failed jobs older than max_age_seconds."""
        now = datetime.now(timezone.utc)
        with self._lock:
            to_remove = []
            for job_id, job in self._jobs.items():
                if job.status in (JobStatus.COMPLETED, JobStatus.FAILED):
                    if job.completed_at:
                        completed = datetime.fromisoformat(job.completed_at)
                        if (now - completed).total_seconds() > max_age_seconds:
                            to_remove.append(job_id)
            for job_id in to_remove:
                del self._jobs[job_id]

    def shutdown(self):
        self._executor.shutdown(wait=False)


# Singleton instance for use across the REST API
_dataset_job_manager: Optional[DatasetJobManager] = None
_manager_lock = threading.Lock()


def get_dataset_job_manager() -> DatasetJobManager:
    global _dataset_job_manager
    if _dataset_job_manager is None:
        with _manager_lock:
            if _dataset_job_manager is None:
                _dataset_job_manager = DatasetJobManager()
    return _dataset_job_manager