# repositories/job_repo.py
from typing import Optional, Dict, Any
from uuid import UUID
from datetime import datetime, timezone
from sqlmodel import select
from sqlalchemy.ext.asyncio import AsyncSession

from ..models import Job, JobStatus, JobType


class JobRepository:
    """Replaces job_store.py with async SQLModel operations"""

    def __init__(self, session: AsyncSession):
        self.session = session

    async def create_job(
        self,
        job_id: str,
        project_id: str,
        account_id: Optional[str] = None,
        status: str = "queued",
        mode: Optional[str] = None,
        config: Optional[Dict[str, Any]] = None,
        **extra_meta,
    ) -> Job:
        """Create a new job (replaces write_job_row)"""
        now = datetime.now(timezone.utc)
        job = Job(
            id=job_id,  # Use string directly, no UUID conversion
            project_id=project_id,  # Use string directly
            account_id=account_id,
            status=JobStatus(status),
            job_type=JobType(mode or "match"),
            created_at=now,
            updated_at=now,
            config={**(config or {}), **extra_meta},
        )

        self.session.add(job)
        await self.session.commit()
        await self.session.refresh(job)
        return job

    async def update_job_status(
        self,
        job_id: str,
        status: str,
        results: Optional[Dict[str, Any]] = None,
        error: Optional[str] = None,
        results_path: Optional[str] = None,
        matches_found: Optional[int] = None,
    ) -> Optional[Job]:
        """Update job status (replaces update_status)"""
        stmt = select(Job).where(Job.id == job_id)  # Use string directly
        result = await self.session.execute(stmt)
        job = result.scalar_one_or_none()

        if not job:
            return None

        job.status = JobStatus(status)
        job.updated_at = datetime.utcnow()

        if results is not None:
            job.results = results
        if error is not None:
            job.error = error
        if results_path is not None:
            job.results_path = results_path
        if matches_found is not None:
            job.matches_found = matches_found

        self.session.add(job)
        await self.session.commit()
        await self.session.refresh(job)
        return job

    async def get_job(self, job_id: str) -> Optional[Job]:
        """Get job by ID (replaces fetch_job)"""
        stmt = select(Job).where(Job.id == job_id)  # Use string directly
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def get_jobs_for_project(
        self, project_id: UUID, status: Optional[JobStatus] = None, limit: int = 100
    ) -> list[Job]:
        """Get all jobs for a project with optional filtering"""
        stmt = select(Job).where(Job.project_id == project_id)

        if status:
            stmt = stmt.where(Job.status == status)

        stmt = stmt.order_by(Job.created_at.desc()).limit(limit)

        result = await self.session.execute(stmt)
        return result.scalars().all()
