from __future__ import annotations

import threading
import time
import uuid
from typing import Any


class JobStore:
    def __init__(self) -> None:
        self._jobs: dict[str, dict[str, Any]] = {}
        self._lock = threading.RLock()

    def create(self, session_id: str, user_id: str, message: str) -> dict[str, Any]:
        job_id = uuid.uuid4().hex
        job = {
            "job_id": job_id,
            "session_id": session_id,
            "user_id": user_id,
            "message": message,
            "status": "queued",
            "created_at": time.time(),
            "updated_at": time.time(),
            "events": [],
            "answer": None,
            "usage": None,
            "error": None,
        }
        with self._lock:
            self._jobs[job_id] = job
        return job

    def get(self, job_id: str) -> dict[str, Any] | None:
        with self._lock:
            job = self._jobs.get(job_id)
            if not job:
                return None
            return {
                key: (value.copy() if isinstance(value, list) else value)
                for key, value in job.items()
            }

    def append_event(self, job_id: str, event: dict[str, Any]) -> None:
        with self._lock:
            job = self._jobs[job_id]
            seq = len(job["events"]) + 1
            job["events"].append({"seq": seq, **event})
            job["updated_at"] = time.time()

    def set_status(self, job_id: str, status: str) -> None:
        with self._lock:
            job = self._jobs[job_id]
            job["status"] = status
            job["updated_at"] = time.time()

    def finish(
        self,
        job_id: str,
        *,
        answer: str,
        usage: dict[str, Any] | None,
    ) -> None:
        with self._lock:
            job = self._jobs[job_id]
            job["status"] = "completed"
            job["answer"] = answer
            job["usage"] = usage
            job["updated_at"] = time.time()

    def fail(self, job_id: str, error_message: str) -> None:
        with self._lock:
            job = self._jobs[job_id]
            job["status"] = "failed"
            job["error"] = error_message
            job["updated_at"] = time.time()