from __future__ import annotations

import logging
import uuid
from typing import TYPE_CHECKING, Any
from datetime import datetime

if TYPE_CHECKING:
    from collections.abc import Callable, Awaitable

    from taskiq import AsyncBroker

from ui_router.events import EventData
from ui_router.exceptions import EventSchedulingError
from ui_router.schema import EventSourceType
from ui_router.integrations._utils import parse_delay

try:
    import taskiq

    HAS_TASKIQ = True
except ImportError:
    HAS_TASKIQ = False

logger = logging.getLogger(__name__)


class TaskiqEventScheduler:
    def __init__(
        self,
        event_callback: Callable[[EventData], Awaitable[None]],
        broker: Any,
        schedule_source: Any | None = None,
    ) -> None:
        if not HAS_TASKIQ:
            msg = "taskiq is not installed. Install it with: pip install aiogram-ui-router[taskiq]"
            raise ImportError(msg)

        self._event_callback = event_callback
        self._broker: AsyncBroker = broker
        self._schedule_source = schedule_source
        self._schedules: dict[str, Any] = {}

        self._task = self._broker.register_task(
            self._fire_event,
            task_name="ui_router_fire_event",
        )

    async def _fire_event(
        self,
        event_name: str,
        event_data: dict[str, Any],
        bot_id: str | int,
        user_id: int | None,
        chat_id: int | None,
    ) -> None:
        event = EventData(
            event_name=event_name,
            source_type=EventSourceType.SCHEDULED,
            data=event_data,
            bot_id=bot_id,
            user_id=user_id,
            chat_id=chat_id,
        )
        try:
            await self._event_callback(event)
        except Exception:
            logger.exception("Error in scheduled event callback for '%s'", event_name)

    def _require_schedule_source(self) -> Any:
        if self._schedule_source is None:
            msg = "schedule_source is required for scheduling. Provide a ScheduleSource instance."
            raise EventSchedulingError(msg)
        return self._schedule_source

    async def schedule_once(
        self,
        event_name: str,
        delay: str,
        event_data: dict[str, Any],
        bot_id: str | int,
        user_id: int | None = None,
        chat_id: int | None = None,
    ) -> str:
        source = self._require_schedule_source()
        delay_td = parse_delay(delay)
        run_at = datetime.now() + delay_td

        task_id = str(uuid.uuid4())
        schedule = await self._task.schedule_by_time(
            source,
            run_at,
            event_name=event_name,
            event_data=event_data,
            bot_id=bot_id,
            user_id=user_id,
            chat_id=chat_id,
            schedule_id=task_id,
        )
        self._schedules[task_id] = schedule
        logger.debug("Scheduled once event '%s' at %s, task_id: %s", event_name, run_at, task_id)
        return task_id

    async def schedule_cron(
        self,
        event_name: str,
        cron_expr: str,
        event_data: dict[str, Any],
        bot_id: str | int,
        user_id: int | None = None,
        chat_id: int | None = None,
    ) -> str:
        source = self._require_schedule_source()
        task_id = str(uuid.uuid4())

        schedule = await self._task.schedule_by_cron(
            source,
            cron_expr,
            event_name=event_name,
            event_data=event_data,
            bot_id=bot_id,
            user_id=user_id,
            chat_id=chat_id,
            schedule_id=task_id,
        )
        self._schedules[task_id] = schedule
        logger.debug("Scheduled cron event '%s' with '%s', task_id: %s", event_name, cron_expr, task_id)
        return task_id

    async def schedule_interval(
        self,
        event_name: str,
        interval_seconds: int,
        event_data: dict[str, Any],
        bot_id: str | int,
        user_id: int | None = None,
        chat_id: int | None = None,
    ) -> str:
        source = self._require_schedule_source()
        task_id = str(uuid.uuid4())

        schedule = await self._task.schedule_by_interval(
            source,
            interval_seconds,
            event_name=event_name,
            event_data=event_data,
            bot_id=bot_id,
            user_id=user_id,
            chat_id=chat_id,
            schedule_id=task_id,
        )
        self._schedules[task_id] = schedule
        logger.debug("Scheduled interval event '%s' every %ds, task_id: %s", event_name, interval_seconds, task_id)
        return task_id

    async def cancel_scheduled(self, task_id: str) -> bool:
        schedule = self._schedules.pop(task_id, None)
        if schedule is None:
            logger.warning("Schedule %s not found for cancellation", task_id)
            return False

        try:
            await schedule.unschedule()
        except Exception:
            logger.exception("Error unscheduling task %s", task_id)
            return False
        else:
            logger.debug("Cancelled scheduled task %s", task_id)
            return True
