"""Database-backed job source (validation_jobs table)."""

from __future__ import annotations

import asyncio
import logging
import uuid
from datetime import datetime, timedelta, timezone
from functools import partial
from typing import Any

from sqlalchemy import func, text, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

from pycharter.db.models.base import get_engine, get_schema_prefix
from pycharter.db.models.validation_job import ValidationJobModel
from pycharter.worker.config import WorkerConfig
from pycharter.worker.job_sources.base import JobSource
from pycharter.worker.models import ClaimedJob, JobStatus, ValidationJob

logger = logging.getLogger(__name__)

_DEFAULT_CLAIM_LEASE_TIMEOUT_S = 300
_DEFAULT_RETRY_BACKOFF_BASE_MS = 1000
_DEFAULT_RETRY_BACKOFF_MAX_MS = 300000


def _utc_now() -> datetime:
    return datetime.now(timezone.utc)


def _status_to_api(s: str) -> str:
    """Map DB status to API status."""
    if s == "pending":
        return "queued"
    if s == "claimed":
        return "running"
    if s in ("failed", "dead_letter"):
        return "failed"
    return s


class DatabaseJobSource(JobSource):
    """Job source backed by the validation_jobs table."""

    def __init__(self, db_url: str, config: WorkerConfig | None = None) -> None:
        self._db_url = db_url
        self._config = config
        self._engine = get_engine(db_url)
        self._is_postgres = not db_url.startswith("sqlite://")
        self._prefix = get_schema_prefix(db_url)
        self._table = f"{self._prefix}validation_jobs"
        self._SessionFactory = __import__(
            "sqlalchemy.orm", fromlist=["sessionmaker"]
        ).sessionmaker(bind=self._engine)

        self._claim_lease_timeout_s = (
            config.claim_lease_timeout_s if config else _DEFAULT_CLAIM_LEASE_TIMEOUT_S
        )
        self._retry_backoff_base_ms = (
            config.retry_backoff_base_ms if config else _DEFAULT_RETRY_BACKOFF_BASE_MS
        )
        self._retry_backoff_max_ms = (
            config.retry_backoff_max_ms if config else _DEFAULT_RETRY_BACKOFF_MAX_MS
        )
        # Table is created by Alembic (pycharter db upgrade). No create_all here to avoid
        # "table validation_jobs already exists" when schema stripping differs for SQLite.

    def _session(self) -> Session:
        return self._SessionFactory()

    async def _run_sync(self, fn: Any, *args: Any, **kwargs: Any) -> Any:
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, partial(fn, *args, **kwargs))

    # ------------------------------------------------------------------
    # JobSource interface
    # ------------------------------------------------------------------

    async def submit(self, job: ValidationJob) -> str:
        return await self._run_sync(self._submit_sync, job)

    async def claim_next(self, worker_id: str) -> ClaimedJob | None:
        return await self._run_sync(self._claim_next_sync, worker_id)

    async def complete(
        self, event_id: str, result: dict[str, Any] | None = None
    ) -> None:
        await self._run_sync(self._complete_sync, event_id, result)

    async def fail(self, event_id: str, error: str, *, retry: bool = True) -> None:
        await self._run_sync(self._fail_sync, event_id, error, retry)

    async def health_check(self) -> bool:
        try:
            return await self._run_sync(self._health_check_sync)
        except Exception:
            return False

    async def get_job(self, job_id: str) -> dict[str, Any] | None:
        return await self._run_sync(self._get_job_sync, job_id)

    async def list_jobs(
        self,
        status: str | None = None,
        limit: int | None = None,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        return await self._run_sync(self._list_jobs_sync, status, limit, offset)

    async def cancel_job(self, job_id: str) -> bool:
        return await self._run_sync(self._cancel_job_sync, job_id)

    async def delete_job(self, job_id: str) -> bool:
        return await self._run_sync(self._delete_job_sync, job_id)

    async def get_stats(self) -> dict[str, int]:
        return await self._run_sync(self._get_stats_sync)

    async def pending_count(self) -> int:
        try:
            return await self._run_sync(self._pending_count_sync)
        except Exception:
            return -1

    async def close(self) -> None:
        if self._engine is not None:
            self._engine.dispose()

    # ------------------------------------------------------------------
    # Sync implementations
    # ------------------------------------------------------------------

    def _submit_sync(self, job: ValidationJob) -> str:
        session = self._session()
        try:
            event_id = job.job_id or str(uuid.uuid4())
            now = _utc_now()
            fires_at = job.fires_at or now
            data_json = {
                "data_source": job.data_source,
                "options": job.options or {},
            }

            if job.idempotency_key and self._is_postgres:
                result = session.execute(
                    text(f"""
                        INSERT INTO {self._table}
                        (id, contract_name, contract_version, data_json, status,
                         fires_at, attempt, max_attempts, idempotency_key, created_at)
                        VALUES
                        (:id, :contract_name, :contract_version,
                         :data_json, 'pending', :fires_at,
                         0, :max_attempts, :idempotency_key, :created_at)
                        ON CONFLICT (idempotency_key) DO UPDATE
                            SET id = {self._table}.id
                        RETURNING id
                    """),
                    {
                        "id": event_id,
                        "contract_name": job.contract_name,
                        "contract_version": job.contract_version,
                        "data_json": data_json,
                        "fires_at": fires_at,
                        "max_attempts": job.max_attempts,
                        "idempotency_key": job.idempotency_key,
                        "created_at": now,
                    },
                )
                row = result.mappings().first()
                session.commit()
                return str(row["id"])

            if job.idempotency_key and not self._is_postgres:
                try:
                    row = ValidationJobModel(
                        id=event_id,
                        contract_name=job.contract_name,
                        contract_version=job.contract_version,
                        data_json=data_json,
                        status=JobStatus.PENDING.value,
                        fires_at=fires_at,
                        attempt=0,
                        max_attempts=job.max_attempts,
                        idempotency_key=job.idempotency_key,
                        created_at=now,
                    )
                    session.add(row)
                    session.commit()
                    return event_id
                except IntegrityError:
                    session.rollback()
                    existing = (
                        session.query(ValidationJobModel.id)
                        .filter(
                            ValidationJobModel.idempotency_key == job.idempotency_key
                        )
                        .first()
                    )
                    if existing:
                        return str(existing[0])
                    raise

            row = ValidationJobModel(
                id=event_id,
                contract_name=job.contract_name,
                contract_version=job.contract_version,
                data_json=data_json,
                status=JobStatus.PENDING.value,
                fires_at=fires_at,
                attempt=0,
                max_attempts=job.max_attempts,
                idempotency_key=job.idempotency_key,
                created_at=now,
            )
            session.add(row)
            session.commit()
            return event_id
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def _release_stale_claims_sync(self) -> None:
        session = self._session()
        try:
            session.execute(
                text(f"""
                    UPDATE {self._table}
                    SET status = 'pending', claimed_by = NULL, claimed_at = NULL
                    WHERE status = 'claimed'
                      AND claimed_at < :cutoff
                """),
                {"cutoff": _utc_now() - timedelta(seconds=self._claim_lease_timeout_s)},
            )
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def _claim_next_sync(self, worker_id: str) -> ClaimedJob | None:
        self._release_stale_claims_sync()
        if self._is_postgres:
            return self._claim_next_postgres(worker_id)
        return self._claim_next_sqlite(worker_id)

    def _claim_next_postgres(self, worker_id: str) -> ClaimedJob | None:
        session = self._session()
        now = _utc_now()
        table = self._table
        try:
            result = session.execute(
                text(f"""
                    WITH claimable AS (
                        SELECT id FROM {table}
                        WHERE status = 'pending' AND fires_at <= :now
                        ORDER BY fires_at ASC
                        LIMIT 1
                        FOR UPDATE SKIP LOCKED
                    )
                    UPDATE {table}
                    SET status = 'claimed', claimed_by = :worker_id, claimed_at = :now,
                        attempt = attempt + 1
                    FROM claimable
                    WHERE {table}.id = claimable.id
                    RETURNING *
                """),
                {"now": now, "worker_id": worker_id},
            )
            row = result.mappings().first()
            session.commit()
            if row is None:
                return None
            return self._row_to_claimed(row, now)
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def _claim_next_sqlite(self, worker_id: str) -> ClaimedJob | None:
        session = self._session()
        now = _utc_now()
        table = self._table
        try:
            result = session.execute(
                text(f"""
                    SELECT id, contract_name, contract_version, data_json,
                           attempt, max_attempts, fires_at
                    FROM {table}
                    WHERE status = 'pending' AND fires_at <= :now
                    ORDER BY fires_at ASC
                    LIMIT 1
                """),
                {"now": now},
            )
            candidate = result.mappings().first()
            if candidate is None:
                return None
            event_id = str(candidate["id"])
            new_attempt = (candidate["attempt"] or 0) + 1
            session.execute(
                text(f"""
                    UPDATE {table}
                    SET status = 'claimed', claimed_by = :worker_id, claimed_at = :now,
                        attempt = :attempt
                    WHERE id = :event_id
                """),
                {
                    "worker_id": worker_id,
                    "now": now,
                    "attempt": new_attempt,
                    "event_id": event_id,
                },
            )
            session.commit()
            return self._row_to_claimed(
                {
                    "id": event_id,
                    "contract_name": candidate["contract_name"],
                    "contract_version": candidate["contract_version"],
                    "data_json": candidate["data_json"],
                    "attempt": new_attempt,
                    "max_attempts": candidate["max_attempts"],
                    "fires_at": candidate["fires_at"],
                },
                now,
            )
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def _row_to_claimed(self, row: Any, claimed_at: datetime) -> ClaimedJob:
        import json as _json

        raw = (
            row["data_json"]
            if hasattr(row, "__getitem__")
            else getattr(row, "data_json", None)
        )
        if isinstance(raw, str):
            data = _json.loads(raw)
        elif isinstance(raw, dict):
            data = raw
        else:
            data = {}
        return ClaimedJob(
            event_id=str(row["id"]),
            contract_name=row["contract_name"],
            contract_version=row["contract_version"],
            data_source=data.get("data_source", ""),
            options=data.get("options") or {},
            attempt=row["attempt"],
            max_attempts=row["max_attempts"],
            fires_at=row["fires_at"],
            claimed_at=claimed_at,
        )

    def _complete_sync(
        self, event_id: str, result: dict[str, Any] | None = None
    ) -> None:
        session = self._session()
        try:
            session.execute(
                update(ValidationJobModel)
                .where(ValidationJobModel.id == event_id)
                .values(
                    status=JobStatus.COMPLETED.value,
                    completed_at=_utc_now(),
                    result_json=result,
                )
            )
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def _fail_sync(self, event_id: str, error: str, retry: bool) -> None:
        session = self._session()
        try:
            row = (
                session.query(ValidationJobModel)
                .filter(ValidationJobModel.id == event_id)
                .first()
            )
            if row is None:
                return
            can_retry = retry and row.attempt < row.max_attempts
            if can_retry:
                backoff_ms = min(
                    self._retry_backoff_max_ms,
                    self._retry_backoff_base_ms * (2**row.attempt),
                )
                row.status = JobStatus.PENDING.value
                row.claimed_by = None
                row.claimed_at = None
                row.error_message = error
                row.fires_at = _utc_now() + timedelta(milliseconds=backoff_ms)
                logger.info(
                    "Job %s failed (attempt %d/%d), re-queued in %dms: %s",
                    event_id,
                    row.attempt,
                    row.max_attempts,
                    backoff_ms,
                    error,
                )
            else:
                row.status = JobStatus.DEAD_LETTER.value
                row.error_message = error
                row.completed_at = _utc_now()
                logger.warning(
                    "Job %s moved to dead letter after %d attempts: %s",
                    event_id,
                    row.attempt,
                    error,
                )
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def _health_check_sync(self) -> bool:
        session = self._session()
        try:
            session.execute(text("SELECT 1"))
            return True
        finally:
            session.close()

    def _get_job_sync(self, job_id: str) -> dict[str, Any] | None:
        session = self._session()
        try:
            row = (
                session.query(ValidationJobModel)
                .filter(ValidationJobModel.id == job_id)
                .first()
            )
            if row is None:
                return None
            return self._model_to_job_response(row)
        finally:
            session.close()

    def _model_to_job_response(self, row: ValidationJobModel) -> dict[str, Any]:
        duration_ms = None
        try:
            if row.completed_at and row.claimed_at:
                duration_ms = (row.completed_at - row.claimed_at).total_seconds() * 1000
            elif row.completed_at and row.created_at:
                duration_ms = (row.completed_at - row.created_at).total_seconds() * 1000
        except TypeError:
            pass  # naive/aware datetime mismatch (SQLite)
        return {
            "job_id": str(row.id),
            "status": _status_to_api(row.status),
            "contract_name": row.contract_name,
            "contract_version": row.contract_version,
            "result": row.result_json,
            "error": row.error_message,
            "created_at": row.created_at,
            "started_at": row.claimed_at,
            "completed_at": row.completed_at,
            "duration_ms": duration_ms,
        }

    def _list_jobs_sync(
        self,
        status: str | None = None,
        limit: int | None = None,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        session = self._session()
        try:
            q = session.query(ValidationJobModel).order_by(
                ValidationJobModel.created_at.desc()
            )
            if status is not None:
                api_to_db = {"queued": "pending", "running": "claimed"}
                db_status = api_to_db.get(status, status)
                q = q.filter(ValidationJobModel.status == db_status)
            if offset > 0:
                q = q.offset(offset)
            if limit is not None:
                q = q.limit(limit)
            rows = q.all()
            return [self._model_to_job_response(r) for r in rows]
        finally:
            session.close()

    def _cancel_job_sync(self, job_id: str) -> bool:
        session = self._session()
        try:
            row = (
                session.query(ValidationJobModel)
                .filter(ValidationJobModel.id == job_id)
                .first()
            )
            if row is None or row.status not in (
                JobStatus.PENDING.value,
                JobStatus.CLAIMED.value,
            ):
                return False
            row.status = JobStatus.CANCELLED.value
            row.completed_at = _utc_now()
            session.commit()
            return True
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def _delete_job_sync(self, job_id: str) -> bool:
        session = self._session()
        try:
            deleted = (
                session.query(ValidationJobModel)
                .filter(ValidationJobModel.id == job_id)
                .delete()
            )
            session.commit()
            return deleted > 0
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def _get_stats_sync(self) -> dict[str, int]:
        session = self._session()
        try:
            total = session.query(ValidationJobModel).count()
            by_status: dict[str, int] | None = {}
            for s in JobStatus:
                by_status[s.value] = (
                    session.query(ValidationJobModel)
                    .filter(ValidationJobModel.status == s.value)
                    .count()
                )
            out = {
                "queued": by_status.get("pending", 0),
                "running": by_status.get("claimed", 0),
                "completed": by_status.get("completed", 0),
                "failed": by_status.get("failed", 0) + by_status.get("dead_letter", 0),
                "cancelled": by_status.get("cancelled", 0),
                "total": total,
            }
            return out
        finally:
            session.close()

    def _pending_count_sync(self) -> int:
        session = self._session()
        try:
            return (
                session.query(func.count(ValidationJobModel.id))
                .filter(ValidationJobModel.status == JobStatus.PENDING.value)
                .scalar()
                or 0
            )
        finally:
            session.close()
