from __future__ import annotations

import queue
import threading
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from .metadata_store import MetadataStore
from .schemas import JobDetail, JobEvent, JobStatus


@dataclass
class JobState:
    status: JobStatus
    events: list[JobEvent] = field(default_factory=list)
    subscribers: list[queue.Queue] = field(default_factory=list)


class JobManager:
    def __init__(self, store: Optional[MetadataStore] = None) -> None:
        self._lock = threading.Lock()
        self.store = store
        self._jobs: dict[str, JobState] = {}
        self._restore_from_store()

    def _restore_from_store(self) -> None:
        if not self.store:
            return
        for status in self.store.load_jobs():
            self._jobs[status.job_id] = JobState(
                status=status,
                events=self.store.load_events(status.job_id, limit=100),
            )

    def create_job(self) -> str:
        job_id = uuid.uuid4().hex
        now = datetime.now(timezone.utc)
        status = JobStatus(
            job_id=job_id,
            status="queued",
            stage="queued",
            progress_percent=0,
            message="任务已创建",
            created_at=now,
            updated_at=now,
            artifacts={},
            error="",
        )
        with self._lock:
            self._jobs[job_id] = JobState(status=status)
        if self.store:
            self.store.upsert_job(status)
        return job_id

    def exists(self, job_id: str) -> bool:
        with self._lock:
            return job_id in self._jobs

    def get_status(self, job_id: str) -> JobStatus:
        with self._lock:
            return self._jobs[job_id].status

    def get_detail(self, job_id: str) -> JobDetail:
        with self._lock:
            state = self._jobs[job_id]
            return JobDetail(status=state.status, recent_events=state.events[-100:])

    def delete_job(self, job_id: str) -> None:
        with self._lock:
            self._jobs.pop(job_id, None)
        if self.store:
            self.store.delete_job(job_id)

    def set_artifacts(self, job_id: str, artifacts: dict[str, str]) -> None:
        with self._lock:
            status = self._jobs[job_id].status
            status.artifacts = artifacts
            status.updated_at = datetime.now(timezone.utc)
        if self.store:
            self.store.upsert_job(status)

    def set_status(
        self,
        job_id: str,
        *,
        status: str,
        stage: str,
        progress_percent: int,
        message: str,
        error: str | None = None,
        elapsed_seconds: float | None = None,
    ) -> None:
        with self._lock:
            s = self._jobs[job_id].status
            s.status = status
            s.stage = stage
            s.progress_percent = progress_percent
            s.message = message
            s.updated_at = datetime.now(timezone.utc)
            if error is not None:
                s.error = error
            if elapsed_seconds is not None:
                s.elapsed_seconds = elapsed_seconds
        if self.store:
            self.store.upsert_job(s)

    def add_event(
        self,
        job_id: str,
        *,
        kind: str,
        stage: str,
        progress_percent: int,
        message: str,
        elapsed_seconds: float | None = None,
    ) -> None:
        event = JobEvent(
            type=kind,
            stage=stage,
            progress_percent=progress_percent,
            message=message,
            timestamp=datetime.now(timezone.utc),
            elapsed_seconds=elapsed_seconds,
        )
        with self._lock:
            state = self._jobs[job_id]
            state.events.append(event)
            state.status.stage = stage
            state.status.progress_percent = progress_percent
            state.status.message = message
            state.status.updated_at = event.timestamp
            if elapsed_seconds is not None:
                state.status.elapsed_seconds = elapsed_seconds
            for sub in state.subscribers:
                sub.put(event)
        if self.store:
            self.store.add_event(job_id, event)
            self.store.upsert_job(state.status)

    def subscribe(self, job_id: str) -> queue.Queue:
        q: queue.Queue = queue.Queue()
        with self._lock:
            state = self._jobs[job_id]
            # A fast job can emit preview_ready/completed before the browser
            # finishes opening the SSE connection. Replay the existing event
            # history into this subscriber so timing and preview updates are
            # never lost; the lock also makes the snapshot/registration
            # atomic with respect to new events.
            for event in state.events:
                q.put(event)
            state.subscribers.append(q)
        return q

    def unsubscribe(self, job_id: str, q: queue.Queue) -> None:
        with self._lock:
            subs = self._jobs[job_id].subscribers
            if q in subs:
                subs.remove(q)

    @staticmethod
    def resolve_artifact_path(status: JobStatus, kind: str) -> Path:
        if kind not in status.artifacts:
            raise KeyError(kind)
        p = status.artifacts[kind]
        if not p:
            raise FileNotFoundError(kind)
        return Path(p)