# -*- coding: utf-8 -*-
from __future__ import annotations

import asyncio
import logging
import random
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Literal, Optional, Union
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

from apscheduler.events import (
    EVENT_JOB_MAX_INSTANCES,
    EVENT_JOB_MISSED,
    JobExecutionEvent,
    JobSubmissionEvent,
)
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.date import DateTrigger
from apscheduler.triggers.interval import IntervalTrigger
from qwenpaw.exceptions import ConfigurationException

from ...config import get_heartbeat_config, get_dream_cron
from ..inbox_store import append_event as append_inbox_event

from ..console_push_store import append as push_store_append
from .executor import CronExecutor
from .heartbeat import (
    is_cron_expression,
    parse_heartbeat_cron,
    parse_heartbeat_every,
    run_heartbeat_once,
)
from .models import CronExecutionRecord, CronJobSpec, CronJobState
from .repo.base import BaseJobRepository
from ...api_action import ManagerBase, api_action

HEARTBEAT_JOB_ID = "_heartbeat"
DREAM_JOB_ID = "_dream"
HEARTBEAT_MISFIRE_GRACE_SECONDS = 60
DREAM_MISFIRE_GRACE_SECONDS = 600
# Spread scheduled Auto-Dream calls across the minute after the configured
# cron time so installations using the default schedule do not start at once.
DREAM_JITTER_MAX_SECONDS = 60
INTERNAL_JOB_IDS = frozenset({HEARTBEAT_JOB_ID, DREAM_JOB_ID})
CRON_HISTORY_LIMIT = 50

logger = logging.getLogger(__name__)


@dataclass
class _Runtime:
    sem: asyncio.Semaphore


class CronManager(ManagerBase):
    endpoint_prefix = "crons"

    def __init__(
        self,
        *,
        repo: BaseJobRepository,
        workspace: Any,
        channel_manager: Any,
        timezone: str = "UTC",  # pylint: disable=redefined-outer-name
        agent_id: Optional[str] = None,
    ):
        self._repo = repo
        self._workspace = workspace
        self._channel_manager = channel_manager
        self._agent_id = agent_id
        self._scheduler = AsyncIOScheduler(timezone=timezone)
        self._executor = CronExecutor(
            workspace=workspace,
            channel_manager=channel_manager,
        )

        self._lock = asyncio.Lock()
        self._states: Dict[str, CronJobState] = {}
        self._history: Dict[str, list[CronExecutionRecord]] = {}
        self._rt: Dict[str, _Runtime] = {}
        self._started = False

    async def start(self) -> None:
        async with self._lock:
            if self._started:
                return
            jobs_file = await self._repo.load()
            valid_job_ids = {
                job.id for job in jobs_file.jobs if job.id is not None
            }
            await self._repo.prune_orphan_history(valid_job_ids)

            self._register_scheduler_listeners()
            self._scheduler.start()
            for job in jobs_file.jobs:
                try:
                    await self._register_or_update(job)
                except Exception as e:  # pylint: disable=broad-except
                    logger.warning(
                        "Skipping invalid cron job during startup: "
                        "job_id=%s name=%s schedule_type=%s cron=%s "
                        "run_at=%s error=%s",
                        job.id,
                        job.name,
                        job.schedule.type,
                        job.schedule.cron,
                        job.schedule.run_at,
                        repr(e),
                    )
                    if job.enabled:
                        disabled_job = job.model_copy(
                            update={"enabled": False},
                        )
                        await self._repo.upsert_job(disabled_job)
                        logger.warning(
                            "Auto-disabled invalid cron job: "
                            "job_id=%s name=%s",
                            job.id,
                            job.name,
                        )

            # Heartbeat: scheduled job when enabled in config
            hb = get_heartbeat_config(self._agent_id)
            if getattr(hb, "enabled", False):
                trigger = self._build_heartbeat_trigger(hb.every)
                self._scheduler.add_job(
                    self._heartbeat_callback,
                    trigger=trigger,
                    id=HEARTBEAT_JOB_ID,
                    misfire_grace_time=HEARTBEAT_MISFIRE_GRACE_SECONDS,
                    replace_existing=True,
                )
                logger.info(
                    "Heartbeat job scheduled for agent %s: every=%s",
                    self._agent_id,
                    hb.every,
                )

            # Dream-based memory optimization: cron job from config
            dream_cron = get_dream_cron(self._agent_id)
            if dream_cron:
                try:
                    trigger = CronTrigger.from_crontab(
                        dream_cron,
                        timezone=self._scheduler.timezone,
                    )
                    self._scheduler.add_job(
                        self._dream_callback,
                        trigger=trigger,
                        id=DREAM_JOB_ID,
                        misfire_grace_time=DREAM_MISFIRE_GRACE_SECONDS,
                        replace_existing=True,
                    )
                    logger.info(
                        f"Dream-based memory optimization job scheduled for "
                        f"agent {self._agent_id}: cron={dream_cron}",
                    )
                except Exception as e:  # pylint: disable=broad-except
                    logger.error(
                        f"Failed to schedule dream-based memory optimization"
                        f"for  agent {self._agent_id}: error={repr(e)}",
                    )

            self._started = True

    async def stop(self) -> None:
        async with self._lock:
            if not self._started:
                return
            self._scheduler.shutdown(wait=False)
            self._started = False

    # ----- read/state -----

    @api_action(
        methods={"http", "cli", "slash"},
        http_method="GET",
        http_path="/crons/jobs",
        slash_command="cron-list",
    )
    async def list_jobs(self) -> list[CronJobSpec]:
        return await self._repo.list_jobs()

    async def get_job(self, job_id: str) -> Optional[CronJobSpec]:
        return await self._repo.get_job(job_id)

    def get_state(self, job_id: str) -> CronJobState:
        return self._states.get(job_id, CronJobState())

    async def get_history(self, job_id: str) -> list[CronExecutionRecord]:
        if job_id not in self._history:
            self._history[job_id] = await self._repo.get_history(job_id)
        return self._history[job_id]

    # ----- write/control -----

    @api_action(
        methods={"http", "cli", "slash"},
        http_method="POST",
        http_path="/crons/jobs",
        request_model=CronJobSpec,
        slash_command="cron-create",
    )
    async def create_or_replace_job(self, spec: CronJobSpec) -> None:
        async with self._lock:
            await self._repo.upsert_job(spec)
            if self._started:
                await self._register_or_update(spec)

    @api_action(
        methods={"http", "cli", "slash"},
        http_method="DELETE",
        http_path="/crons/jobs/{job_id}",
        slash_command="cron-delete",
    )
    async def delete_job(self, job_id: str) -> bool:
        async with self._lock:
            if self._started and self._scheduler.get_job(job_id):
                self._scheduler.remove_job(job_id)
            self._states.pop(job_id, None)
            self._history.pop(job_id, None)
            await self._repo.delete_history(job_id)
            self._rt.pop(job_id, None)
            return await self._repo.delete_job(job_id)

    async def pause_job(self, job_id: str) -> None:
        async with self._lock:
            self._scheduler.pause_job(job_id)

    async def resume_job(self, job_id: str) -> None:
        async with self._lock:
            self._scheduler.resume_job(job_id)

    async def reschedule_heartbeat(self) -> None:
        """Reload heartbeat config and update or remove the heartbeat job.

        Note: CronManager should always be started during workspace
        initialization, so this method assumes self._started is True.
        """
        async with self._lock:
            if not self._started:
                logger.warning(
                    f"CronManager not started for agent {self._agent_id}, "
                    f"cannot reschedule heartbeat. This should not happen.",
                )
                return

            hb = get_heartbeat_config(self._agent_id)

            # Remove existing heartbeat job if present
            if self._scheduler.get_job(HEARTBEAT_JOB_ID):
                self._scheduler.remove_job(HEARTBEAT_JOB_ID)

            # Add heartbeat job if enabled
            if getattr(hb, "enabled", False):
                trigger = self._build_heartbeat_trigger(hb.every)
                self._scheduler.add_job(
                    self._heartbeat_callback,
                    trigger=trigger,
                    id=HEARTBEAT_JOB_ID,
                    misfire_grace_time=HEARTBEAT_MISFIRE_GRACE_SECONDS,
                    replace_existing=True,
                )
                logger.info(
                    "heartbeat rescheduled: every=%s",
                    hb.every,
                )
            else:
                logger.info("heartbeat disabled, job removed")

    async def reschedule_dream(self) -> None:
        """Reschedule the dream-based memory optimization job based on
        configuration.

        Note: CronManager should always be started during workspace
        initialization, so this method assumes self._started is True.
        """
        async with self._lock:
            if not self._started:
                logger.warning(
                    f"CronManager not started for agent {self._agent_id}, "
                    "cannot reschedule dream-based memory optimization."
                    "This should not happen.",
                )
                return

            # Check if dream-based memory optimization is enabled in config
            dream_cron = get_dream_cron(self._agent_id)

            # Remove existing job if any
            if self._scheduler.get_job(DREAM_JOB_ID):
                self._scheduler.remove_job(DREAM_JOB_ID)
                logger.info(
                    "Dream-based memory optimization job removed for "
                    f"agent {self._agent_id}",
                )

            # Add new job if cron expression is valid
            if dream_cron:
                try:
                    trigger = CronTrigger.from_crontab(
                        dream_cron,
                        timezone=self._scheduler.timezone,
                    )
                    self._scheduler.add_job(
                        self._dream_callback,
                        trigger=trigger,
                        id=DREAM_JOB_ID,
                        misfire_grace_time=DREAM_MISFIRE_GRACE_SECONDS,
                        replace_existing=True,
                    )
                    logger.info(
                        "Dream-based memory optimization job rescheduled"
                        f"for agent {self._agent_id}: cron={dream_cron}",
                    )
                except Exception as e:  # pylint: disable=broad-except
                    logger.error(
                        "Failed to reschedule dream-based memory  "
                        f"optimization for agent {self._agent_id}: "
                        f"error={repr(e)}",
                    )
            else:
                logger.info(
                    "dream-based memory optimization disabled, job removed",
                )

    async def run_job(self, job_id: str) -> None:
        """Trigger a job to run in the background (fire-and-forget).

        Raises KeyError if the job does not exist.
        The actual execution happens asynchronously; errors are logged
        and reflected in the job state but NOT propagated to the caller.
        """
        job = await self._repo.get_job(job_id)
        if not job:
            raise KeyError(f"Job not found: {job_id}")
        logger.info(
            "cron run_job (async): job_id=%s channel=%s task_type=%s "
            "target_user_id=%s target_session_id=%s",
            job_id,
            job.dispatch.channel,
            job.task_type,
            (job.dispatch.target.user_id or "")[:40],
            (job.dispatch.target.session_id or "")[:40],
        )
        task = asyncio.create_task(
            self._execute_once(
                job,
                trigger="manual",
            ),
            name=f"cron-run-{job_id}",
        )
        task.add_done_callback(lambda t: self._task_done_cb(t, job))

    # ----- callbacks -----

    def _task_done_cb(self, task: asyncio.Task, job: CronJobSpec) -> None:
        """Suppress and log exceptions from fire-and-forget tasks.

        On failure, push an error message to the console push store so
        the frontend can display it.
        """
        if task.cancelled():
            return
        exc = task.exception()
        if exc is not None:
            logger.error(
                "cron background task %s failed: %s",
                task.get_name(),
                repr(exc),
            )
            # Push error to the console for the frontend to display
            session_id = job.dispatch.target.session_id
            if session_id:
                error_text = f"❌ Cron job [{job.name}] failed: {exc}"
                asyncio.ensure_future(
                    push_store_append(session_id, error_text),
                )

    # ----- internal -----

    def _register_scheduler_listeners(self) -> None:
        mask = EVENT_JOB_MISSED | EVENT_JOB_MAX_INSTANCES
        self._scheduler.add_listener(self._on_scheduler_event, mask=mask)

    def _on_scheduler_event(
        self,
        event: JobExecutionEvent | JobSubmissionEvent,
    ) -> None:
        if event.code == EVENT_JOB_MISSED:
            asyncio.create_task(self._handle_job_missed(event))
        elif event.code == EVENT_JOB_MAX_INSTANCES:
            asyncio.create_task(self._handle_job_max_instances(event))

    async def _handle_job_missed(self, event: JobExecutionEvent) -> None:
        job_id = event.job_id
        if job_id in INTERNAL_JOB_IDS:
            return

        job = await self._repo.get_job(job_id)
        if not job:
            return

        scheduled = event.scheduled_run_time
        if scheduled.tzinfo is None:
            scheduled = scheduled.replace(tzinfo=timezone.utc)
        late_seconds = max(
            0,
            int((datetime.now(timezone.utc) - scheduled).total_seconds()),
        )
        grace = job.runtime.misfire_grace_seconds
        error_msg = (
            f"missed scheduled run at {scheduled.isoformat()}: "
            f"late by {late_seconds}s, grace={grace}s"
        )
        await self._record_skipped(job, error_msg)

    async def _handle_job_max_instances(
        self,
        event: JobSubmissionEvent,
    ) -> None:
        job_id = event.job_id
        if job_id in INTERNAL_JOB_IDS:
            return

        job = await self._repo.get_job(job_id)
        if not job:
            return

        scheduled_times = event.scheduled_run_times or []
        if scheduled_times:
            # coalesce may queue multiple due times;
            # [-1] is the latest skipped slot.
            scheduled = scheduled_times[-1]
            if scheduled.tzinfo is None:
                scheduled = scheduled.replace(tzinfo=timezone.utc)
            scheduled_text = scheduled.isoformat()
        else:
            scheduled_text = "unknown"
        error_msg = (
            f"skipped scheduled run at {scheduled_text}: "
            f"maximum running instances reached "
            f"({job.runtime.max_concurrency})"
        )
        await self._record_skipped(job, error_msg)

    async def _record_skipped(self, job: CronJobSpec, error_msg: str) -> None:
        if job.id is None:
            logger.error(
                "cron _record_skipped: job.id is None, skipping record",
            )
            return
        logger.warning(
            "cron job skipped: job_id=%s name=%s %s",
            job.id,
            job.name,
            error_msg,
        )

        st = self._states.get(job.id, CronJobState())
        st.last_status = "skipped"
        st.last_error = error_msg
        aps_job = self._scheduler.get_job(job.id)
        st.next_run_at = aps_job.next_run_time if aps_job else st.next_run_at
        self._states[job.id] = st

        record = CronExecutionRecord(
            run_at=self._now_in_job_timezone(job),
            status="skipped",
            error=error_msg,
            trigger="scheduled",
        )
        records = await self._repo.append_history(
            job.id,
            record,
            limit=CRON_HISTORY_LIMIT,
        )
        self._history[job.id] = records

    async def _register_or_update(self, spec: CronJobSpec) -> None:
        # Validate and build trigger first. If schedule is invalid, fail fast
        # without mutating scheduler/runtime state.
        assert spec.id is not None, "Job must have an id"
        trigger = self._build_trigger(spec)

        # per-job concurrency semaphore
        self._rt[spec.id] = _Runtime(
            sem=asyncio.Semaphore(spec.runtime.max_concurrency),
        )

        # replace existing
        if self._scheduler.get_job(spec.id):
            self._scheduler.remove_job(spec.id)

        self._scheduler.add_job(
            self._scheduled_callback,
            trigger=trigger,
            id=spec.id,
            args=[spec.id],
            misfire_grace_time=spec.runtime.misfire_grace_seconds,
            replace_existing=True,
        )

        if not spec.enabled:
            self._scheduler.pause_job(spec.id)

        # update next_run
        aps_job = self._scheduler.get_job(spec.id)
        st = self._states.get(spec.id, CronJobState())
        st.next_run_at = aps_job.next_run_time if aps_job else None
        self._states[spec.id] = st

    def _build_trigger(
        self,
        spec: CronJobSpec,
    ) -> Union[CronTrigger, DateTrigger, IntervalTrigger]:
        if spec.schedule.type == "once":
            assert spec.schedule.run_at is not None
            if spec.schedule.repeat_every_days:
                end_date: datetime | None = None
                if (
                    spec.schedule.repeat_end_type == "until"
                    and spec.schedule.repeat_until is not None
                ):
                    end_date = spec.schedule.repeat_until
                elif (
                    spec.schedule.repeat_end_type == "count"
                    and spec.schedule.repeat_count is not None
                ):
                    end_date = spec.schedule.run_at + timedelta(
                        days=spec.schedule.repeat_every_days
                        * (spec.schedule.repeat_count - 1),
                    )
                return IntervalTrigger(
                    days=spec.schedule.repeat_every_days,
                    start_date=spec.schedule.run_at,
                    end_date=end_date,
                    timezone=spec.schedule.timezone,
                )
            return DateTrigger(
                run_date=spec.schedule.run_at,
                timezone=spec.schedule.timezone,
            )

        # enforce 5 fields (no seconds)
        assert spec.schedule.cron is not None
        parts = [p for p in spec.schedule.cron.split() if p]
        if len(parts) != 5:
            raise ConfigurationException(
                config_key="cron.schedule.cron",
                message=(
                    f"cron must have 5 fields, "
                    f"got {len(parts)}: {spec.schedule.cron}"
                ),
            )

        minute, hour, day, month, day_of_week = parts
        return CronTrigger(
            minute=minute,
            hour=hour,
            day=day,
            month=month,
            day_of_week=day_of_week,
            timezone=spec.schedule.timezone,
        )

    def _build_heartbeat_trigger(
        self,
        every: str,
    ) -> Union[CronTrigger, IntervalTrigger]:
        """Build a trigger from the heartbeat *every* value.

        Returns CronTrigger for cron expressions,
        IntervalTrigger for interval strings.
        """
        if is_cron_expression(every):
            minute, hour, day, month, day_of_week = parse_heartbeat_cron(every)
            return CronTrigger(
                minute=minute,
                hour=hour,
                day=day,
                month=month,
                day_of_week=day_of_week,
            )
        interval_seconds = parse_heartbeat_every(every)
        return IntervalTrigger(seconds=interval_seconds)

    async def _scheduled_callback(self, job_id: str) -> None:
        job = await self._repo.get_job(job_id)
        if not job:
            return

        await self._execute_once(
            job,
            trigger="scheduled",
        )

        # refresh next_run
        aps_job = self._scheduler.get_job(job_id)
        st = self._states.get(job_id, CronJobState())
        st.next_run_at = aps_job.next_run_time if aps_job else None
        self._states[job_id] = st

    @staticmethod
    def _now_in_job_timezone(job: CronJobSpec) -> datetime:
        tz_name = job.schedule.timezone or "UTC"
        try:
            tz = ZoneInfo(tz_name)
        except (ZoneInfoNotFoundError, ValueError):
            logger.warning(
                "Invalid cron job timezone, using UTC: job_id=%s "
                "timezone=%s",
                job.id,
                tz_name,
            )
            tz = timezone.utc
        return datetime.now(tz)

    async def _heartbeat_callback(self) -> None:
        """Run one heartbeat (HEARTBEAT.md as query, optional dispatch)."""
        try:
            workspace_dir = getattr(
                self._workspace,
                "workspace_dir",
                None,
            )

            await run_heartbeat_once(
                workspace=self._workspace,
                channel_manager=self._channel_manager,
                agent_id=self._agent_id,
                workspace_dir=workspace_dir,
            )
        except asyncio.CancelledError:
            logger.info("heartbeat cancelled")
            raise
        except Exception:  # pylint: disable=broad-except
            logger.exception("heartbeat run failed")

    async def _dream_callback(self) -> None:
        """Run one scheduled dream task after a random startup delay."""
        try:
            delay_seconds = random.randint(0, DREAM_JITTER_MAX_SECONDS)
            logger.info(
                "Dream task for agent %s will start in %s seconds",
                self._agent_id,
                delay_seconds,
            )
            await asyncio.sleep(delay_seconds)
            await self._workspace.memory_manager.dream()
            logger.debug("Dream task executed successfully")
        except asyncio.CancelledError:
            logger.info("Dream task was cancelled")
            raise
        except Exception as e:  # pylint: disable=broad-except
            logger.error(f"Failed to execute dream task: {e}", exc_info=True)

    # pylint: disable-next=too-many-branches,too-many-statements
    async def _execute_once(
        self,
        job: CronJobSpec,
        *,
        trigger: Literal["scheduled", "manual"] = "scheduled",
    ) -> None:
        assert job.id is not None, "Job must have an id"
        rt = self._rt.get(job.id)
        if not rt:
            rt = _Runtime(sem=asyncio.Semaphore(job.runtime.max_concurrency))
            self._rt[job.id] = rt

        async with rt.sem:
            st = self._states.get(job.id, CronJobState())
            st.last_status = "running"
            self._states[job.id] = st
            execution_result: dict[str, Any] = {}
            execution_succeeded = False
            delivery_failed = False

            try:
                execution_result = await self._executor.execute(job)
                execution_succeeded = True
                delivery_failed = (
                    execution_result.get("delivery_status") == "failed"
                )
                if delivery_failed:
                    st.last_status = "error"
                    delivery_error = (
                        execution_result.get("delivery_error")
                        or "delivery failed"
                    )
                    st.last_error = f"delivery failed: {delivery_error}"
                else:
                    st.last_status = "success"
                    st.last_error = None
                logger.info(
                    "cron _execute_once: job_id=%s status=success",
                    job.id,
                )
            except asyncio.CancelledError:
                st.last_status = "cancelled"
                st.last_error = "Job was cancelled"
                logger.info(
                    "cron _execute_once: job_id=%s status=cancelled",
                    job.id,
                )
                raise
            except Exception as e:  # pylint: disable=broad-except
                st.last_status = "error"
                st.last_error = repr(e)
                logger.warning(
                    "cron _execute_once: job_id=%s status=error error=%s",
                    job.id,
                    repr(e),
                )
                raise
            finally:
                st.last_run_at = self._now_in_job_timezone(job)
                self._states[job.id] = st
                record = CronExecutionRecord(
                    run_at=st.last_run_at,
                    status=st.last_status or "error",
                    error=st.last_error,
                    trigger=trigger,
                )
                records = await self._repo.append_history(
                    job.id,
                    record,
                    limit=CRON_HISTORY_LIMIT,
                )
                self._history[job.id] = records
                if execution_succeeded:
                    if delivery_failed:
                        try:
                            await append_inbox_event(
                                agent_id=self._agent_id,
                                source_type="cron",
                                source_id=job.id,
                                event_type="cron_delivery_failed_fallback",
                                status="error",
                                severity="error",
                                title=f"Cron result not delivered: {job.name}",
                                body=(
                                    "Task executed successfully, "
                                    "but channel delivery failed."
                                ),
                                payload={
                                    "job_id": job.id,
                                    "job_name": job.name,
                                    "task_type": job.task_type,
                                    "trigger": trigger,
                                    "run_id": execution_result.get("run_id"),
                                    "delivery_error": execution_result.get(
                                        "delivery_error",
                                    ),
                                },
                            )
                        except Exception:  # pylint: disable=broad-except
                            logger.exception(
                                "failed to append cron fallback event",
                            )
                    elif job.save_result_to_inbox:
                        if job.task_type == "text":
                            body = (job.text or "").strip()
                        else:
                            body = "Agent cron task finished successfully."
                        try:
                            await append_inbox_event(
                                agent_id=self._agent_id,
                                source_type="cron",
                                source_id=job.id,
                                event_type="cron_result",
                                status="success",
                                severity="info",
                                title=f"Cron result: {job.name}",
                                body=body,
                                payload={
                                    "job_id": job.id,
                                    "job_name": job.name,
                                    "task_type": job.task_type,
                                    "trigger": trigger,
                                    "run_id": execution_result.get("run_id"),
                                    "save_result_to_inbox": (
                                        job.save_result_to_inbox
                                    ),
                                },
                            )
                        except Exception:  # pylint: disable=broad-except
                            logger.exception(
                                "failed to append cron result inbox event",
                            )
