from __future__ import annotations

from datetime import datetime
from typing import Any
from uuid import uuid4

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.date import DateTrigger
from tzlocal import get_localzone

from traderai.memory import MemoryStore, iso_now, time_since


class WakeScheduler:
    def __init__(self, memory: MemoryStore) -> None:
        self.memory = memory
        self.scheduler = AsyncIOScheduler(timezone=get_localzone())
        self.agent = None

    def bind_agent(self, agent: Any) -> None:
        self.agent = agent

    def start(self) -> None:
        if not self.scheduler.running:
            self.scheduler.start()
        for job in self.memory.list_jobs():
            self._schedule_existing(job)

    def shutdown(self) -> None:
        if self.scheduler.running:
            self.scheduler.shutdown(wait=False)

    def schedule_date(self, run_at: str, prompt: str, job_id: str | None = None) -> dict[str, Any]:
        parsed = datetime.fromisoformat(run_at)
        job_id = job_id or f"wake-{uuid4()}"
        trigger = DateTrigger(run_date=parsed)
        self.scheduler.add_job(self._run_job, trigger=trigger, id=job_id, args=[job_id, prompt], replace_existing=True)
        return self.memory.add_job(job_id, prompt, "date", run_at, parsed.isoformat())

    def schedule_cron(self, cron: str, prompt: str, job_id: str | None = None) -> dict[str, Any]:
        job_id = job_id or f"wake-{uuid4()}"
        trigger = CronTrigger.from_crontab(cron)
        self.scheduler.add_job(self._run_job, trigger=trigger, id=job_id, args=[job_id, prompt], replace_existing=True)
        next_run = self.scheduler.get_job(job_id).next_run_time
        return self.memory.add_job(job_id, prompt, "cron", cron, next_run.isoformat() if next_run else None)

    def list_jobs(self) -> list[dict[str, Any]]:
        return self.memory.list_jobs()

    def _schedule_existing(self, job: dict[str, Any]) -> None:
        if job["trigger_type"] == "cron":
            trigger = CronTrigger.from_crontab(job["trigger_value"])
        elif job["trigger_type"] == "date":
            trigger = DateTrigger(run_date=datetime.fromisoformat(job["trigger_value"]))
        else:
            return
        self.scheduler.add_job(
            self._run_job,
            trigger=trigger,
            id=job["id"],
            args=[job["id"], job["prompt"]],
            replace_existing=True,
        )

    async def _run_job(self, job_id: str, prompt: str) -> None:
        last = self.memory.last_interaction()
        last_text = f"{last['created_at']} ({time_since(last['created_at'])})" if last else "never"
        wake_message = (
            f"Scheduled wake job fired. Current time is {iso_now()}. "
            f"The last chat interaction was {last_text}. Job instruction: {prompt}"
        )
        if self.agent is None:
            self.memory.add_outbox(wake_message)
            return

        text = await self.agent.generate_wake_response(wake_message)
        self.memory.add_outbox(text)
        self.memory.mark_job_run(job_id)