IPSExtractor/services/bot_descarga/infrastructure/store.py

"""Repositorio sobre PostgreSQL (tablas download_* + sync de tasks)."""
from __future__ import annotations

from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional

from sqlalchemy import select

from domain.models import (
    Job,
    JobEventRecord,
    JobStatus,
    Worker,
    WorkerStatus,
    utc_now,
)
from infrastructure.database import (
    DownloadJobEventRow,
    DownloadJobRow,
    DownloadWorkerRow,
    Task,
    _iso,
    init_db,
    session_scope,
    utcnow,
)
from infrastructure.storage_minio import downloads_prefix

HEARTBEAT_OFFLINE_SECONDS = 45


def _parse_iso(value: str | None) -> datetime | None:
    if not value:
        return None
    dt = datetime.fromisoformat(value)
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return dt


class JobStore:
    """API de persistencia compatible con el store SQLite anterior."""

    def __init__(self, db_path: str | None = None):
        # db_path ignorado (compat); usa DATABASE_URL.
        init_db()

    def _row_worker(self, row: DownloadWorkerRow) -> Worker:
        worker = Worker(
            id=row.id,
            hostname=row.hostname or "",
            max_concurrent=int(row.max_concurrent or 1),
            status=row.status,
            last_heartbeat=_iso(row.last_heartbeat) or utc_now(),
            current_job_id=row.current_job_id,
        )
        return self._refresh_worker_status(worker)

    def _row_job(self, row: DownloadJobRow) -> Job:
        return Job(
            id=row.id,
            worker_id=row.worker_id,
            status=row.status,
            title=row.title or "",
            excel_key=row.excel_key or "",
            downloads_prefix=row.downloads_prefix or "",
            original_filename=row.original_filename or "",
            cancel_requested=bool(row.cancel_requested),
            result=dict(row.result or {}),
            error=row.error,
            created_at=_iso(row.created_at) or utc_now(),
            started_at=_iso(row.started_at),
            finished_at=_iso(row.finished_at),
            job_id=row.job_id,
            task_id=row.task_id,
        )

    def _refresh_worker_status(self, worker: Worker) -> Worker:
        try:
            hb = datetime.fromisoformat(worker.last_heartbeat)
            if hb.tzinfo is None:
                hb = hb.replace(tzinfo=timezone.utc)
            age = datetime.now(timezone.utc) - hb.astimezone(timezone.utc)
            if age > timedelta(seconds=HEARTBEAT_OFFLINE_SECONDS):
                worker.status = WorkerStatus.OFFLINE.value
            elif worker.current_job_id:
                worker.status = WorkerStatus.BUSY.value
            else:
                worker.status = WorkerStatus.ONLINE.value
        except Exception:
            worker.status = WorkerStatus.OFFLINE.value
        return worker

    def register_worker(
        self,
        worker_id: str,
        hostname: str = "",
        max_concurrent: int = 1,
    ) -> Worker:
        with session_scope() as session:
            row = session.get(DownloadWorkerRow, worker_id)
            now = utcnow()
            if row:
                row.hostname = hostname
                row.max_concurrent = max_concurrent
                row.status = WorkerStatus.ONLINE.value
                row.last_heartbeat = now
                row.current_job_id = None
            else:
                row = DownloadWorkerRow(
                    id=worker_id,
                    hostname=hostname,
                    max_concurrent=max_concurrent,
                    status=WorkerStatus.ONLINE.value,
                    last_heartbeat=now,
                    current_job_id=None,
                )
                session.add(row)
            session.flush()
            return self._row_worker(row)

    def heartbeat(self, worker_id: str) -> Optional[Worker]:
        with session_scope() as session:
            row = session.get(DownloadWorkerRow, worker_id)
            if not row:
                return None
            row.last_heartbeat = utcnow()
            row.status = (
                WorkerStatus.BUSY.value if row.current_job_id else WorkerStatus.ONLINE.value
            )
            session.flush()
            return self._row_worker(row)

    def get_worker(self, worker_id: str) -> Optional[Worker]:
        with session_scope() as session:
            row = session.get(DownloadWorkerRow, worker_id)
            if not row:
                return None
            return self._row_worker(row)

    def list_workers(self) -> List[Worker]:
        with session_scope() as session:
            rows = session.scalars(
                select(DownloadWorkerRow).order_by(DownloadWorkerRow.id)
            ).all()
            return [self._row_worker(r) for r in rows]

    def is_worker_online(self, worker_id: str) -> bool:
        worker = self.get_worker(worker_id)
        if not worker:
            return False
        return worker.status in (WorkerStatus.ONLINE.value, WorkerStatus.BUSY.value)

    def update_excel_key(self, job_id: str, excel_key: str) -> Optional[Job]:
        with session_scope() as session:
            row = session.get(DownloadJobRow, job_id)
            if not row:
                return None
            row.excel_key = excel_key
            session.flush()
            return self._row_job(row)

    def create_job(
        self,
        worker_id: str,
        excel_key: str,
        original_filename: str = "",
        title: str = "",
        *,
        job_id: Optional[str] = None,
        task_id: Optional[str] = None,
    ) -> Job:
        import uuid as uuid_mod

        download_id = str(uuid_mod.uuid4())
        prefix = downloads_prefix(download_id)
        with session_scope() as session:
            row = DownloadJobRow(
                id=download_id,
                job_id=job_id,
                task_id=task_id,
                worker_id=worker_id,
                status=JobStatus.PENDING.value,
                title=title or f"Descarga {original_filename}",
                excel_key=excel_key,
                downloads_prefix=prefix,
                original_filename=original_filename,
                cancel_requested=False,
                result={},
                created_at=utcnow(),
            )
            session.add(row)
            session.flush()
            return self._row_job(row)

    def get_job(self, job_id: str) -> Optional[Job]:
        with session_scope() as session:
            row = session.get(DownloadJobRow, job_id)
            if not row:
                return None
            return self._row_job(row)

    def list_jobs(
        self,
        limit: int = 100,
        *,
        job_id: Optional[str] = None,
    ) -> List[Job]:
        with session_scope() as session:
            stmt = select(DownloadJobRow).order_by(DownloadJobRow.created_at.desc()).limit(limit)
            if job_id:
                stmt = (
                    select(DownloadJobRow)
                    .where(DownloadJobRow.job_id == job_id)
                    .order_by(DownloadJobRow.created_at.desc())
                    .limit(limit)
                )
            rows = session.scalars(stmt).all()
            return [self._row_job(r) for r in rows]

    def claim_pending(self, worker_id: str) -> Optional[Job]:
        with session_scope() as session:
            worker = session.get(DownloadWorkerRow, worker_id)
            if not worker or worker.current_job_id:
                return None

            row = session.scalars(
                select(DownloadJobRow)
                .where(
                    DownloadJobRow.worker_id == worker_id,
                    DownloadJobRow.status == JobStatus.PENDING.value,
                )
                .order_by(DownloadJobRow.created_at.asc())
                .limit(1)
            ).first()
            if not row:
                return None

            now = utcnow()
            row.status = JobStatus.RUNNING.value
            row.started_at = now
            worker.current_job_id = row.id
            worker.status = WorkerStatus.BUSY.value
            worker.last_heartbeat = now

            if row.task_id:
                task = session.get(Task, row.task_id)
                if task:
                    task.status = JobStatus.RUNNING.value
                    task.started_at = now
                    task.updated_at = now

            session.flush()
            return self._row_job(row)

    def request_cancel(self, job_id: str) -> Optional[Job]:
        with session_scope() as session:
            row = session.get(DownloadJobRow, job_id)
            if not row:
                return None
            if row.status in (
                JobStatus.COMPLETED.value,
                JobStatus.FAILED.value,
                JobStatus.CANCELLED.value,
            ):
                return self._row_job(row)

            now = utcnow()
            row.cancel_requested = True
            if row.status == JobStatus.PENDING.value:
                row.status = JobStatus.CANCELLED.value
                row.finished_at = now
                self._sync_task(session, row, status=JobStatus.CANCELLED.value)
            session.flush()
            return self._row_job(row)

    def is_cancel_requested(self, job_id: str) -> bool:
        job = self.get_job(job_id)
        return bool(job and job.cancel_requested)

    def _sync_task(
        self,
        session,
        row: DownloadJobRow,
        *,
        status: Optional[str] = None,
        result: Optional[Dict[str, Any]] = None,
        error: Optional[str] = None,
    ) -> None:
        if not row.task_id:
            return
        task = session.get(Task, row.task_id)
        if not task:
            return
        now = utcnow()
        if status:
            task.status = status
        if result is not None:
            task.result = {
                **dict(task.result or {}),
                "download_job_id": row.id,
                **result,
            }
        if error is not None:
            task.error = error
        task.updated_at = now
        if status in (
            JobStatus.COMPLETED.value,
            JobStatus.FAILED.value,
            JobStatus.CANCELLED.value,
        ):
            task.finished_at = now

    def patch_job(
        self,
        job_id: str,
        *,
        status: Optional[str] = None,
        result: Optional[Dict[str, Any]] = None,
        error: Optional[str] = None,
    ) -> Optional[Job]:
        with session_scope() as session:
            row = session.get(DownloadJobRow, job_id)
            if not row:
                return None

            new_status = status or row.status
            if result is not None:
                row.result = result
            if error is not None:
                row.error = error
            row.status = new_status

            finished = False
            if new_status in (
                JobStatus.COMPLETED.value,
                JobStatus.FAILED.value,
                JobStatus.CANCELLED.value,
            ):
                row.finished_at = row.finished_at or utcnow()
                finished = True

            self._sync_task(
                session,
                row,
                status=new_status if status else None,
                result=result,
                error=error,
            )

            if finished:
                worker = session.get(DownloadWorkerRow, row.worker_id)
                if worker and worker.current_job_id == job_id:
                    worker.current_job_id = None
                    worker.status = WorkerStatus.ONLINE.value

            session.flush()
            return self._row_job(row)

    def add_event(
        self,
        job_id: str,
        event_type: str,
        message: str,
        *,
        radicado: Optional[str] = None,
        nit: Optional[str] = None,
        index: Optional[int] = None,
        total: Optional[int] = None,
        metadata: Optional[Dict[str, Any]] = None,
        timestamp: Optional[str] = None,
    ) -> JobEventRecord:
        with session_scope() as session:
            ts = _parse_iso(timestamp) or utcnow()
            ev = DownloadJobEventRow(
                download_job_id=job_id,
                event_type=event_type,
                message=message,
                timestamp=ts,
                radicado=radicado,
                nit=nit,
                idx=index,
                total=total,
                metadata_json=metadata or {},
            )
            session.add(ev)
            session.flush()
            return JobEventRecord(
                id=int(ev.id),
                job_id=job_id,
                event_type=event_type,
                message=message,
                timestamp=_iso(ts) or utc_now(),
                radicado=radicado,
                nit=nit,
                index=index,
                total=total,
                metadata=metadata or {},
            )

    def list_events(self, job_id: str, since_id: int = 0) -> List[JobEventRecord]:
        with session_scope() as session:
            rows = session.scalars(
                select(DownloadJobEventRow)
                .where(
                    DownloadJobEventRow.download_job_id == job_id,
                    DownloadJobEventRow.id > since_id,
                )
                .order_by(DownloadJobEventRow.id.asc())
            ).all()
            return [
                JobEventRecord(
                    id=int(r.id),
                    job_id=r.download_job_id,
                    event_type=r.event_type,
                    message=r.message,
                    timestamp=_iso(r.timestamp) or utc_now(),
                    radicado=r.radicado,
                    nit=r.nit,
                    index=r.idx,
                    total=r.total,
                    metadata=dict(r.metadata_json or {}),
                )
                for r in rows
            ]

    def create_downloads_task(
        self,
        *,
        parent_job_id: str,
        download_job_id: str,
        worker_id: str,
        title: str,
        original_filename: str,
    ) -> str:
        """Crea Task(type=downloads) y la asocia al download_job. Devuelve task_id."""
        with session_scope() as session:
            now = utcnow()
            task = Task(
                job_id=parent_job_id,
                type="downloads",
                status=JobStatus.PENDING.value,
                title=title or f"Bot descarga ({original_filename})",
                request={
                    "worker_id": worker_id,
                    "original_filename": original_filename,
                    "download_job_id": download_job_id,
                },
                result={"download_job_id": download_job_id},
                created_at=now,
                updated_at=now,
            )
            session.add(task)
            session.flush()
            row = session.get(DownloadJobRow, download_job_id)
            if row:
                row.task_id = task.id
                row.job_id = parent_job_id
            session.flush()
            return task.id

    def parent_job_exists(self, parent_job_id: str) -> bool:
        from infrastructure.database import Job as IpsJob

        with session_scope() as session:
            return session.get(IpsJob, parent_job_id) is not None