"""
Task backend for distributed task processing.

This provides an abstraction for task storage and execution:
- BrokerTaskBackend: Uses PostgreSQL as broker (production API servers)
- WorkerTaskBackend: No-op submit_task (production workers — child tasks are polled)
- SyncTaskBackend: Executes tasks immediately (testing/embedded)
"""

import json
import logging
from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable
from typing import Any

logger = logging.getLogger(__name__)


def fq_table(table: str, schema: str | None = None) -> str:
    """Get fully-qualified table name with optional schema prefix."""
    from .schema import fq_table_explicit

    return fq_table_explicit(table, schema)


class TaskBackend(ABC):
    """
    Abstract base class for task execution backends.

    Implementations must:
    1. Store/publish task events (as serializable dicts)
    2. Execute tasks through a provided executor callback (optional)

    The backend treats tasks as pure dictionaries that can be serialized
    and stored in the database. The executor (typically MemoryEngine.execute_task)
    receives the dict and routes it to the appropriate handler.
    """

    def __init__(self):
        """Initialize the task backend."""
        self._executor: Callable[[dict[str, Any]], Awaitable[None]] | None = None
        self._initialized = False

    def set_executor(self, executor: Callable[[dict[str, Any]], Awaitable[None]]):
        """
        Set the executor callback for processing tasks.

        Args:
            executor: Async function that takes a task dict and executes it
        """
        self._executor = executor

    @abstractmethod
    async def initialize(self):
        """
        Initialize the backend (e.g., connect to database).
        """
        pass

    @abstractmethod
    async def submit_task(self, task_dict: dict[str, Any]):
        """
        Submit a task for execution.

        Args:
            task_dict: Task as a dictionary (must be serializable)
        """
        pass

    @abstractmethod
    async def shutdown(self):
        """
        Shutdown the backend gracefully.
        """
        pass

    async def _execute_task(self, task_dict: dict[str, Any]):
        """
        Execute a task through the registered executor.

        Args:
            task_dict: Task dictionary to execute

        Raises:
            Exception: Re-raised from executor on failure.
        """
        if self._executor is None:
            task_type = task_dict.get("type", "unknown")
            logger.warning(f"No executor registered, skipping task {task_type}")
            return

        await self._executor(task_dict)


class SyncTaskBackend(TaskBackend):
    """
    Synchronous task backend that executes tasks immediately.

    This is useful for tests and embedded/CLI usage where we don't want
    background workers. Tasks are executed inline rather than being queued.
    """

    async def initialize(self):
        """No-op for sync backend."""
        self._initialized = True
        logger.debug("SyncTaskBackend initialized")

    async def submit_task(self, task_dict: dict[str, Any]):
        """
        Execute the task immediately (synchronously).

        Args:
            task_dict: Task dictionary to execute
        """
        if not self._initialized:
            await self.initialize()

        await self._execute_task(task_dict)

    async def shutdown(self):
        """No-op for sync backend."""
        self._initialized = False
        logger.debug("SyncTaskBackend shutdown")


class WorkerTaskBackend(TaskBackend):
    """
    Task backend for worker processes.

    Workers execute tasks directly via the poller (claim → execute), so they
    don't need submit_task to run anything.  When engine code running *inside*
    a worker-executed task calls submit_task (e.g. retain triggers consolidation),
    the async-operation row has already been persisted (with task_payload) by
    _submit_async_operation — so submit_task is a no-op.  The new task will be
    picked up by a worker on the next poll cycle instead of being executed inline,
    which avoids blocking the parent task.
    """

    async def initialize(self):
        self._initialized = True
        logger.debug("WorkerTaskBackend initialized")

    async def submit_task(self, task_dict: dict[str, Any]):
        """No-op: the row already exists in async_operations; a worker will claim it."""
        task_type = task_dict.get("type", "unknown")
        logger.debug(f"WorkerTaskBackend: submit_task no-op for {task_type} (will be picked up by poller)")

    async def shutdown(self):
        self._initialized = False
        logger.debug("WorkerTaskBackend shutdown")


class BrokerTaskBackend(TaskBackend):
    """
    Task backend using PostgreSQL as broker.

    submit_task() stores task_payload in async_operations table.
    Actual polling and execution is handled separately by WorkerPoller.

    This backend is used by the API to store tasks. Workers poll
    the database separately to claim and execute tasks.
    """

    def __init__(
        self,
        pool_getter: Callable[[], Any],
        schema: str | None = None,
        schema_getter: Callable[[], str | None] | None = None,
    ):
        """
        Initialize the broker task backend.

        Args:
            pool_getter: Callable that returns the asyncpg connection pool
            schema: Database schema for multi-tenant support (optional, static)
            schema_getter: Callable that returns current schema dynamically (optional).
                          If set, takes precedence over static schema for submit_task.
        """
        super().__init__()
        self._pool_getter = pool_getter
        self._schema = schema
        self._schema_getter = schema_getter

    async def initialize(self):
        """Initialize the backend."""
        self._initialized = True
        logger.info("BrokerTaskBackend initialized")

    async def submit_task(self, task_dict: dict[str, Any]):
        """
        Store task payload in async_operations table.

        The task_dict should contain an 'operation_id' if updating an existing
        operation record, otherwise a new operation will be created.

        Args:
            task_dict: Task dictionary to store (must be JSON serializable)
        """
        if not self._initialized:
            await self.initialize()

        pool = self._pool_getter()
        operation_id = task_dict.get("operation_id")
        task_type = task_dict.get("type", "unknown")
        bank_id = task_dict.get("bank_id")

        # Custom encoder to handle datetime objects
        from datetime import datetime

        def datetime_encoder(obj):
            if isinstance(obj, datetime):
                return obj.isoformat()
            raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")

        payload_json = json.dumps(task_dict, default=datetime_encoder)

        schema = self._schema_getter() if self._schema_getter else self._schema
        table = fq_table("async_operations", schema)

        from .db_utils import acquire_with_retry

        if operation_id:
            # Callers now include task_payload in the same INSERT that creates the
            # async_operations row (see MemoryEngine._submit_async_operation). The
            # WHERE clause guards against overwriting that payload — the UPDATE is a
            # no-op when the row is already claimable, and only fills in a NULL payload
            # for any legacy caller that still creates the row first.
            async with acquire_with_retry(pool) as conn:
                await conn.execute(
                    f"""
                    UPDATE {table}
                    SET task_payload = $1::jsonb, updated_at = now()
                    WHERE operation_id = $2 AND task_payload IS NULL
                    """,
                    payload_json,
                    operation_id,
                )
            logger.debug(f"submit_task UPDATE for operation {operation_id} (no-op if payload already set)")
        else:
            # Insert new operation (for tasks without pre-created records)
            # e.g., access_count_update tasks
            import uuid

            new_id = uuid.uuid4()
            async with acquire_with_retry(pool) as conn:
                await conn.execute(
                    f"""
                    INSERT INTO {table} (operation_id, bank_id, operation_type, status, task_payload)
                    VALUES ($1, $2, $3, 'pending', $4::jsonb)
                    """,
                    new_id,
                    bank_id,
                    task_type,
                    payload_json,
                )
            logger.debug(f"Created new operation {new_id} for task type {task_type}")

    async def shutdown(self):
        """Shutdown the backend."""
        self._initialized = False
        logger.info("BrokerTaskBackend shutdown")

    async def wait_for_pending_tasks(self, timeout: float = 120.0):
        """
        Wait for pending tasks to be processed.

        In the broker model, this polls the database to check if tasks
        for this process have been completed. This is useful in tests
        when worker_enabled=True (API processes its own tasks).

        Args:
            timeout: Maximum time to wait in seconds
        """
        import asyncio

        from .db_utils import acquire_with_retry

        pool = self._pool_getter()
        schema = self._schema_getter() if self._schema_getter else self._schema
        table = fq_table("async_operations", schema)

        start_time = asyncio.get_event_loop().time()
        while asyncio.get_event_loop().time() - start_time < timeout:
            # Check if there are any pending tasks with payloads
            async with acquire_with_retry(pool) as conn:
                count = await conn.fetchval(
                    f"""
                    SELECT COUNT(*) FROM {table}
                    WHERE status = 'pending' AND task_payload IS NOT NULL
                    """
                )

            if count == 0:
                return

            await asyncio.sleep(0.5)

        logger.warning(f"Timeout waiting for pending tasks after {timeout}s")