from pydantic import BaseModel
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger
from app.services.session_manager import session_manager
from app.services.storage import storage
from app.core.logging import logger

SCHEDULE_KEY = "schedules"


class ScheduleTask(BaseModel):
    id: str
    name: str
    cron: str
    prompt: str
    enabled: bool = True


class SchedulerService:
    def __init__(self):
        self.scheduler = AsyncIOScheduler()
        self.scheduled_tasks: dict[str, ScheduleTask] = {}

    def load(self):
        data = storage.load(SCHEDULE_KEY, {})
        if isinstance(data, dict):
            for k, v in data.items():
                self.scheduled_tasks[k] = ScheduleTask(**v)
        logger.info(f"Loaded {len(self.scheduled_tasks)} scheduled tasks")

    def save(self):
        data = {k: v.model_dump() for k, v in self.scheduled_tasks.items()}
        storage.save(SCHEDULE_KEY, data)

    def start(self):
        self.scheduler.start()
        for task in self.scheduled_tasks.values():
            if task.enabled:
                trigger = CronTrigger.from_crontab(task.cron)
                self.scheduler.add_job(
                    self._run_scheduled_task,
                    trigger,
                    args=[task.id],
                    id=task.id,
                )
                logger.info(f"Loaded schedule: {task.name} ({task.cron})")
        logger.info("Scheduler started")

    def shutdown(self):
        self.scheduler.shutdown()
        logger.info("Scheduler shutdown")

    async def add_schedule(self, task: ScheduleTask) -> ScheduleTask:
        self.scheduled_tasks[task.id] = task
        self.save()

        if task.enabled:
            trigger = CronTrigger.from_crontab(task.cron)
            self.scheduler.add_job(
                self._run_scheduled_task,
                trigger,
                args=[task.id],
                id=task.id,
            )
            logger.info(f"Added schedule: {task.name} ({task.cron})")

        return task

    async def remove_schedule(self, task_id: str) -> bool:
        task = self.scheduled_tasks.pop(task_id, None)
        if not task:
            return False

        try:
            self.scheduler.remove_job(task_id)
        except Exception:
            pass

        self.save()
        logger.info(f"Removed schedule: {task.name}")
        return True

    async def list_schedules(self) -> list[ScheduleTask]:
        return list(self.scheduled_tasks.values())

    async def _run_scheduled_task(self, task_id: str):
        task = self.scheduled_tasks.get(task_id)
        if not task:
            return

        logger.info(f"Running scheduled task: {task.name}")

        from app.models.session import CreateSessionRequest, SessionType

        request = CreateSessionRequest(
            type=SessionType.SCHEDULED,
            prompt=task.prompt,
        )
        task_obj = await session_manager.create_task(request)

        try:
            await session_manager.execute_task_async(task_obj.id)
            logger.info(f"Scheduled task {task.name} started")
        except Exception as e:
            logger.error(f"Scheduled task {task.name} failed to start: {e}")


scheduler_service = SchedulerService()