# retoor <retoor@molodetz.nl>

from __future__ import annotations

import logging
from datetime import datetime, timezone
from typing import Any

import dataset
import sqlalchemy

from .guards import (
    AutomationDenied,
    REASON_CREATE_QUOTA,
    REASON_NESTED,
    REASON_NOT_A_USER,
    automation_allowed,
    creation_denial,
    nesting_allowed,
)
from .limits import create_quota, insert_task_within_quota
from .schedule import now_utc

logger = logging.getLogger("devii.tasks.store")

TABLE = "devii_tasks"
INDEXED_COLUMNS = (
    ["owner_kind", "owner_id"],
    ["uid"],
    ["enabled"],
    ["next_run_at"],
    ["status"],
)
TASK_COLUMNS = (
    ("uid", ""),
    ("owner_kind", ""),
    ("owner_id", ""),
    ("label", ""),
    ("prompt", ""),
    ("enabled", True),
    ("status", ""),
    ("created_at", ""),
    ("next_run_at", ""),
    ("last_run_at", ""),
    ("run_count", 0),
    ("last_result", ""),
    ("last_error", ""),
    ("kind", ""),
    ("run_at", ""),
    ("delay_seconds", 0),
    ("every_seconds", 0),
    ("start_at", ""),
    ("cron", ""),
    ("max_runs", 0),
    ("deleted_at", ""),
    ("deleted_by", ""),
    ("notify", 0),
    ("tz", ""),
    ("expires_at", ""),
    ("failure_count", 0),
)
DUE_SQL = (
    "SELECT * FROM devii_tasks WHERE enabled = 1 AND status = 'pending' "
    "AND deleted_at IS NULL AND next_run_at IS NOT NULL AND next_run_at <= :now "
    "ORDER BY next_run_at LIMIT :limit"
)
CLAIM_SQL = (
    "UPDATE devii_tasks SET status = 'running' WHERE uid = :uid "
    "AND status = 'pending' AND enabled = 1 AND deleted_at IS NULL"
)


def memory_db() -> Any:
    return dataset.connect("sqlite:///:memory:")


ACTIVE_STATUSES = ("pending", "running")


def due_rows(db: Any, now_iso: str, limit: int) -> list[dict[str, Any]]:
    if TABLE not in db.tables:
        return []
    return list(db.query(DUE_SQL, now=now_iso, limit=limit))


def claim(db: Any, uid: str) -> bool:
    statement = sqlalchemy.text(CLAIM_SQL)
    with db:
        result = db.executable.execute(statement, {"uid": uid})
    return result.rowcount == 1


class TaskStore:
    def __init__(
        self, db: Any, owner_kind: str, owner_id: str, operator: bool = False
    ) -> None:
        self._db = db
        self._owner_kind = owner_kind
        self._owner_id = owner_id
        self._operator = operator
        self._ensure_indexes()

    def _ensure_indexes(self) -> None:
        if TABLE not in self._db.tables:
            return
        self._ensure_schema()
        for columns in INDEXED_COLUMNS:
            try:
                self._table.create_index(columns)
            except sqlalchemy.exc.OperationalError as exc:
                if "already exists" not in str(exc).lower():
                    raise
                logger.debug("Index on %s was created by another process", columns)

    def _ensure_schema(self) -> None:
        table = self._db[TABLE]
        for column, example in TASK_COLUMNS:
            if table.has_column(column):
                continue
            try:
                table.create_column_by_example(column, example)
            except sqlalchemy.exc.OperationalError as exc:
                if "duplicate column" not in str(exc).lower():
                    raise
                logger.debug("Column %s was added by another process", column)
                table._reflect_table()

    @property
    def db(self) -> Any:
        return self._db

    def automation_allowed(self) -> bool:
        if self._operator:
            return True
        return automation_allowed(self._owner_kind, self._owner_id)

    def require_automation(self) -> None:
        if not self.automation_allowed():
            raise AutomationDenied(REASON_NOT_A_USER)

    def require_scheduling_allowed(self) -> None:
        self.require_automation()
        if self._operator:
            return
        if not nesting_allowed(self._owner_id):
            raise AutomationDenied(REASON_NESTED)

    def require_creation_allowed(self) -> None:
        if self._operator:
            return
        denial = creation_denial(self._db, self._owner_kind, self._owner_id, now_utc())
        if denial is not None:
            raise denial

    def count_active(self) -> int:
        rows = self._table.find(enabled=True, deleted_at=None, **self._scope)
        return sum(1 for row in rows if row.get("status") in ACTIVE_STATUSES)

    @property
    def _table(self) -> Any:
        return self._db[TABLE]

    @property
    def _scope(self) -> dict[str, str]:
        return {"owner_kind": self._owner_kind, "owner_id": self._owner_id}

    def create(self, record: dict[str, Any]) -> None:
        self.require_creation_allowed()
        self._ensure_schema()
        row = {
            "deleted_at": None,
            "deleted_by": None,
            "failure_count": 0,
            **record,
            **self._scope,
        }
        if self._operator:
            self._table.insert(row)
        elif not insert_task_within_quota(
            self._db, row, self._owner_kind, self._owner_id, now_utc()
        ):
            quota = create_quota(self._db, self._owner_kind, self._owner_id, now_utc())
            raise AutomationDenied(REASON_CREATE_QUOTA, quota.free_at)
        logger.info(
            "Task created uid=%s owner=%s/%s",
            record.get("uid"),
            self._owner_kind,
            self._owner_id,
        )

    def get(self, uid: str) -> dict[str, Any] | None:
        return self._table.find_one(uid=uid, deleted_at=None, **self._scope)

    def list(
        self, enabled_only: bool = False, status: str | None = None
    ) -> list[dict[str, Any]]:
        criteria: dict[str, Any] = dict(self._scope)
        criteria["deleted_at"] = None
        if enabled_only:
            criteria["enabled"] = True
        if status:
            criteria["status"] = status
        return list(self._table.find(**criteria))

    def update(self, uid: str, changes: dict[str, Any]) -> None:
        if changes.get("enabled"):
            self.require_scheduling_allowed()
        changes = {**changes, "uid": uid, **self._scope}
        self._table.update(changes, ["uid", "owner_kind", "owner_id"])
        logger.debug("Task updated uid=%s changes=%s", uid, list(changes))

    def delete(self, uid: str) -> bool:
        row = self._table.find_one(uid=uid, deleted_at=None, **self._scope)
        if not row:
            return False
        self._table.update(
            {
                "uid": uid,
                "deleted_at": datetime.now(timezone.utc).isoformat(),
                "deleted_by": f"{self._owner_kind}:{self._owner_id}",
                **self._scope,
            },
            ["uid", "owner_kind", "owner_id"],
        )
        logger.info("Task soft-deleted uid=%s", uid)
        return True

    def recover_running(self) -> int:
        if TABLE not in self._db.tables:
            return 0
        stuck = list(self._table.find(status="running", deleted_at=None, **self._scope))
        for row in stuck:
            self._table.update(
                {"uid": row["uid"], "status": "pending", **self._scope},
                ["uid", "owner_kind", "owner_id"],
            )
        if stuck:
            logger.info("Recovered %d task(s) stuck in running", len(stuck))
        return len(stuck)

    def has_pending(self) -> bool:
        rows = self._table.find(enabled=True, deleted_at=None, **self._scope)
        return any(row.get("status") in ACTIVE_STATUSES for row in rows)

    def due(self, now_iso: str) -> list[dict[str, Any]]:
        rows = self._table.find(
            enabled=True, status="pending", deleted_at=None, **self._scope
        )
        return [
            row
            for row in rows
            if row.get("next_run_at") and row["next_run_at"] <= now_iso
        ]