from __future__ import annotations

import json
import re
from collections.abc import Iterable
from datetime import timedelta

from ai_orchestrator.schemas.queue_models import QueueJob, QueueJobStatus
from ai_orchestrator.schemas.task_state import TaskState, utc_now
from ai_orchestrator.schemas.task_status import TaskLifecycleStatus


_ENCODING_TO_CODEC = {
    "WIN1252": "cp1252",
}

_ESCAPED_CHAR_PATTERN = re.compile(
    r"\\u[0-9a-fA-F]{4}|\\U[0-9a-fA-F]{8}|\\x[0-9a-fA-F]{2}"
)


def _sanitize_payload_for_encoding(value, *, codec: str | None):
    if codec is None:
        return value
    if isinstance(value, str):
        try:
            value.encode(codec)
            return value
        except UnicodeEncodeError:
            return value.encode(codec, "backslashreplace").decode(codec)
    if isinstance(value, list):
        return [_sanitize_payload_for_encoding(item, codec=codec) for item in value]
    if isinstance(value, dict):
        return {
            _sanitize_payload_for_encoding(key, codec=codec): _sanitize_payload_for_encoding(item, codec=codec)
            for key, item in value.items()
        }
    return value


def _restore_payload_from_encoding(value):
    if isinstance(value, str):
        return _ESCAPED_CHAR_PATTERN.sub(_decode_escaped_character, value)
    if isinstance(value, list):
        return [_restore_payload_from_encoding(item) for item in value]
    if isinstance(value, dict):
        return {
            _restore_payload_from_encoding(key): _restore_payload_from_encoding(item)
            for key, item in value.items()
        }
    return value


def _decode_escaped_character(match: re.Match[str]) -> str:
    token = match.group(0)
    if token.startswith("\\u"):
        return chr(int(token[2:], 16))
    if token.startswith("\\U"):
        return chr(int(token[2:], 16))
    if token.startswith("\\x"):
        return chr(int(token[2:], 16))
    return token


def _safe_json_payload(payload: dict, *, codec: str | None) -> str:
    sanitized = _sanitize_payload_for_encoding(payload, codec=codec)
    return json.dumps(sanitized, ensure_ascii=True)


def _safe_text(value: str | None, *, codec: str | None) -> str | None:
    if value is None:
        return None
    sanitized = _sanitize_payload_for_encoding(value, codec=codec)
    return str(sanitized)


def _load_psycopg():
    try:
        import psycopg
    except ImportError as exc:
        raise RuntimeError(
            "Postgres support requires `psycopg`. Install the postgres extra from pyproject.toml."
        ) from exc
    return psycopg


class PostgresTaskStore:
    def __init__(self, dsn: str) -> None:
        self.dsn = dsn
        self._psycopg = _load_psycopg()
        self._text_codec = self._resolve_text_codec()
        self._initialize()

    def _connect(self):
        return self._psycopg.connect(self.dsn)

    def _resolve_text_codec(self) -> str | None:
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute("SHOW SERVER_ENCODING")
            row = cursor.fetchone()
        if row is None:
            return None
        return _ENCODING_TO_CODEC.get(str(row[0]).upper())

    def _initialize(self) -> None:
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute(
                """
                CREATE TABLE IF NOT EXISTS tasks (
                    task_id TEXT PRIMARY KEY,
                    status TEXT NOT NULL,
                    payload TEXT NOT NULL,
                    created_at TIMESTAMPTZ NOT NULL,
                    updated_at TIMESTAMPTZ NOT NULL
                )
                """
            )
            cursor.execute(
                """
                SELECT data_type
                FROM information_schema.columns
                WHERE table_name = 'tasks' AND column_name = 'payload'
                """
            )
            row = cursor.fetchone()
            if row is not None and str(row[0]).lower() == "jsonb":
                cursor.execute(
                    """
                    ALTER TABLE tasks
                    ALTER COLUMN payload TYPE TEXT
                    USING payload::text
                    """
                )
            connection.commit()

    def create(self, task: TaskState) -> None:
        payload = task.model_dump(mode="json")
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO tasks (task_id, status, payload, created_at, updated_at)
                VALUES (%s, %s, %s, %s, %s)
                """,
                (
                    task.task_id,
                    task.status.value,
                    _safe_json_payload(payload, codec=self._text_codec),
                    task.timestamps.created_at,
                    task.timestamps.updated_at,
                ),
            )
            connection.commit()

    def save(self, task: TaskState) -> None:
        task.touch()
        payload = task.model_dump(mode="json")
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE tasks
                SET status = %s, payload = %s, updated_at = %s
                WHERE task_id = %s
                """,
                (
                    task.status.value,
                    _safe_json_payload(payload, codec=self._text_codec),
                    task.timestamps.updated_at,
                    task.task_id,
                ),
            )
            connection.commit()

    def get(self, task_id: str) -> TaskState | None:
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute("SELECT payload FROM tasks WHERE task_id = %s", (task_id,))
            row = cursor.fetchone()
        if row is None:
            return None
        payload = _restore_payload_from_encoding(json.loads(row[0]))
        return TaskState.model_validate(payload)

    def list_by_statuses(self, statuses: Iterable[TaskLifecycleStatus]) -> list[TaskState]:
        values = [status.value for status in statuses]
        if not values:
            return []
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute("SELECT payload FROM tasks WHERE status = ANY(%s)", (values,))
            rows = cursor.fetchall()
        return [TaskState.model_validate(_restore_payload_from_encoding(json.loads(row[0]))) for row in rows]


class PostgresQueueStore:
    def __init__(self, dsn: str) -> None:
        self.dsn = dsn
        self._psycopg = _load_psycopg()
        self._text_codec = self._resolve_text_codec()
        self._initialize()

    def _connect(self):
        return self._psycopg.connect(self.dsn)

    def _resolve_text_codec(self) -> str | None:
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute("SHOW SERVER_ENCODING")
            row = cursor.fetchone()
        if row is None:
            return None
        return _ENCODING_TO_CODEC.get(str(row[0]).upper())

    def _initialize(self) -> None:
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute(
                """
                CREATE TABLE IF NOT EXISTS jobs (
                    job_id TEXT PRIMARY KEY,
                    task_id TEXT NOT NULL,
                    engine_name TEXT NOT NULL,
                    status TEXT NOT NULL,
                    attempts INTEGER NOT NULL,
                    max_attempts INTEGER NOT NULL,
                    worker_id TEXT,
                    lease_until TIMESTAMPTZ,
                    available_at TIMESTAMPTZ NOT NULL,
                    last_heartbeat TIMESTAMPTZ,
                    error_message TEXT,
                    created_at TIMESTAMPTZ NOT NULL,
                    updated_at TIMESTAMPTZ NOT NULL
                )
                """
            )
            cursor.execute(
                "CREATE INDEX IF NOT EXISTS idx_jobs_status_available ON jobs (status, available_at)"
            )
            connection.commit()

    def enqueue(self, *, task_id: str, engine_name: str, max_attempts: int) -> QueueJob:
        from uuid import uuid4

        now = utc_now()
        job = QueueJob(
            job_id=f"job_{uuid4().hex[:12]}",
            task_id=task_id,
            engine_name=engine_name,
            status=QueueJobStatus.queued,
            attempts=0,
            max_attempts=max_attempts,
            available_at=now,
            created_at=now,
            updated_at=now,
        )
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO jobs (
                    job_id, task_id, engine_name, status, attempts, max_attempts,
                    worker_id, lease_until, available_at, last_heartbeat, error_message,
                    created_at, updated_at
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """,
                (
                    job.job_id,
                    job.task_id,
                    job.engine_name,
                    job.status.value,
                    job.attempts,
                    job.max_attempts,
                    job.worker_id,
                    job.lease_until,
                    job.available_at,
                    job.last_heartbeat,
                    _safe_text(job.error_message, codec=self._text_codec),
                    job.created_at,
                    job.updated_at,
                ),
            )
            connection.commit()
        return job

    def lease_next(self, *, worker_id: str, lease_seconds: int) -> QueueJob | None:
        now = utc_now()
        lease_until = now + timedelta(seconds=lease_seconds)
        with self._connect() as connection, connection.cursor(row_factory=self._psycopg.rows.dict_row) as cursor:
            cursor.execute(
                """
                WITH candidate AS (
                    SELECT job_id
                    FROM jobs
                    WHERE (
                        status = %s AND available_at <= %s
                    ) OR (
                        status = %s AND lease_until IS NOT NULL AND lease_until <= %s
                    )
                    ORDER BY created_at ASC
                    LIMIT 1
                    FOR UPDATE SKIP LOCKED
                )
                UPDATE jobs
                SET status = %s,
                    worker_id = %s,
                    attempts = jobs.attempts + 1,
                    lease_until = %s,
                    last_heartbeat = %s,
                    updated_at = %s
                FROM candidate
                WHERE jobs.job_id = candidate.job_id
                RETURNING jobs.*
                """,
                (
                    QueueJobStatus.queued.value,
                    now,
                    QueueJobStatus.leased.value,
                    now,
                    QueueJobStatus.leased.value,
                    worker_id,
                    lease_until,
                    now,
                    now,
                ),
            )
            row = cursor.fetchone()
            connection.commit()
        if row is None:
            return None
        return QueueJob.model_validate(row)

    def heartbeat(self, *, job_id: str, worker_id: str, lease_seconds: int) -> None:
        now = utc_now()
        lease_until = now + timedelta(seconds=lease_seconds)
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE jobs
                SET lease_until = %s, last_heartbeat = %s, updated_at = %s
                WHERE job_id = %s AND worker_id = %s AND status = %s
                """,
                (
                    lease_until,
                    now,
                    now,
                    job_id,
                    worker_id,
                    QueueJobStatus.leased.value,
                ),
            )
            connection.commit()

    def mark_completed(self, job_id: str) -> None:
        self._mark_terminal(job_id, QueueJobStatus.completed, None)

    def mark_failed(self, job_id: str, error_message: str) -> None:
        self._mark_terminal(job_id, QueueJobStatus.failed, error_message)

    def mark_cancelled(self, job_id: str) -> None:
        self._mark_terminal(job_id, QueueJobStatus.cancelled, None)

    def cancel_for_task(self, task_id: str) -> None:
        now = utc_now()
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE jobs
                SET status = %s, updated_at = %s
                WHERE task_id = %s AND status = ANY(%s)
                """,
                (
                    QueueJobStatus.cancelled.value,
                    now,
                    task_id,
                    [QueueJobStatus.queued.value, QueueJobStatus.leased.value],
                ),
            )
            connection.commit()

    def requeue_expired(self) -> int:
        now = utc_now()
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE jobs
                SET status = %s, worker_id = NULL, lease_until = NULL, updated_at = %s
                WHERE status = %s AND lease_until IS NOT NULL AND lease_until <= %s
                """,
                (
                    QueueJobStatus.queued.value,
                    now,
                    QueueJobStatus.leased.value,
                    now,
                ),
            )
            rowcount = cursor.rowcount
            connection.commit()
        return int(rowcount or 0)

    def get_by_task_id(self, task_id: str) -> list[QueueJob]:
        with self._connect() as connection, connection.cursor(row_factory=self._psycopg.rows.dict_row) as cursor:
            cursor.execute("SELECT * FROM jobs WHERE task_id = %s ORDER BY created_at ASC", (task_id,))
            rows = cursor.fetchall()
        return [QueueJob.model_validate(row) for row in rows]

    def list_by_statuses(self, statuses: Iterable[QueueJobStatus]) -> list[QueueJob]:
        values = [status.value for status in statuses]
        if not values:
            return []
        with self._connect() as connection, connection.cursor(row_factory=self._psycopg.rows.dict_row) as cursor:
            cursor.execute("SELECT * FROM jobs WHERE status = ANY(%s) ORDER BY created_at ASC", (values,))
            rows = cursor.fetchall()
        return [QueueJob.model_validate(row) for row in rows]

    def _mark_terminal(self, job_id: str, status: QueueJobStatus, error_message: str | None) -> None:
        now = utc_now()
        with self._connect() as connection, connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE jobs
                SET status = %s, error_message = %s, updated_at = %s, lease_until = NULL
                WHERE job_id = %s
                """,
                (
                    status.value,
                    _safe_text(error_message, codec=self._text_codec),
                    now,
                    job_id,
                ),
            )
            connection.commit()
