"""Scheduler engine — manages cron-based agent invocations via APScheduler 3.x."""

from __future__ import annotations

import asyncio
import contextlib
import logging
import time
import uuid
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

from apscheduler.jobstores.memory import MemoryJobStore
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger

from agent_gateway.persistence.domain import ScheduleRecord
from agent_gateway.scheduler.handler import run_scheduled_job, set_scheduler_engine

if TYPE_CHECKING:
    from collections.abc import Callable, Coroutine

    from agent_gateway.config import SchedulerConfig
    from agent_gateway.persistence.protocols import ExecutionRepository, ScheduleRepository
    from agent_gateway.queue.protocol import ExecutionQueue
    from agent_gateway.scheduler.lock import DistributedLock
    from agent_gateway.workspace.agent import AgentDefinition, ScheduleConfig

logger = logging.getLogger(__name__)


class SchedulerEngine:
    """Manages cron-based agent scheduling via APScheduler 3.x AsyncIOScheduler.

    Wraps APScheduler behind a thin interface so it can be swapped
    to v4 when it stabilizes.
    """

    def __init__(
        self,
        config: SchedulerConfig,
        schedule_repo: ScheduleRepository,
        execution_repo: ExecutionRepository,
        queue: ExecutionQueue,
        invoke_fn: Callable[..., Coroutine[Any, Any, Any]],
        track_task: Callable[[asyncio.Task[None]], None],
        timezone: str = "UTC",
        distributed_lock: DistributedLock | None = None,
    ) -> None:
        self._config = config
        self._schedule_repo = schedule_repo
        self._execution_repo = execution_repo
        self._queue = queue
        self._invoke_fn = invoke_fn
        self._track_task = track_task
        self._timezone = timezone

        # Distributed lock for multi-instance duplicate prevention
        from agent_gateway.scheduler.lock import NullDistributedLock

        self._distributed_lock: DistributedLock = distributed_lock or NullDistributedLock()

        # APScheduler instance, created in start()
        self._scheduler: AsyncIOScheduler | None = None

        # Overlap prevention: tracks schedule_ids with monotonic fire time
        self._active_scheduled: dict[str, float] = {}
        self._active_lock = asyncio.Lock()

        # Schedule definitions keyed by schedule_id for lookup at fire time
        self._schedule_configs: dict[str, ScheduleConfig] = {}
        self._agent_map: dict[str, str] = {}  # schedule_id -> agent_id
        self._known_agent_ids: set[str] = set()  # all loaded agent IDs

    async def start(
        self,
        agents: dict[str, AgentDefinition],
    ) -> None:
        """Register schedules and start the APScheduler background loop."""
        await self._distributed_lock.initialize()

        # Track all known agent IDs (for validation in create_admin_schedule)
        self._known_agent_ids = set(agents.keys())

        # Build schedule_id -> config mapping
        for agent in agents.values():
            for sched in agent.schedules:
                schedule_id = f"{agent.id}:{sched.name}"
                self._schedule_configs[schedule_id] = sched
                self._agent_map[schedule_id] = agent.id

        # Configure APScheduler with in-memory job store.
        # SQLAlchemyJobStore is intentionally avoided because APScheduler 3.x
        # uses pickle serialization which is an RCE risk.
        self._scheduler = AsyncIOScheduler(
            jobstores={"default": MemoryJobStore()},
            timezone=self._timezone,
        )

        # Set the module-level reference so the handler can reach us
        set_scheduler_engine(self)

        # Register each schedule as a cron job
        for schedule_id, sched_config in self._schedule_configs.items():
            agent_id = self._agent_map[schedule_id]
            await self._register_job(schedule_id, agent_id, sched_config)

        self._scheduler.start()

        # Sync to persistence (after start so next_run_time is computed)
        await self._sync_schedule_records(agents)

        # Load admin-created schedules from DB
        await self._load_admin_schedules()

        total = len(self._schedule_configs)
        enabled_count = sum(1 for c in self._schedule_configs.values() if c.enabled)
        logger.info(
            "Scheduler started: %d schedules (%d enabled)",
            total,
            enabled_count,
        )

    async def stop(self) -> None:
        """Stop the scheduler gracefully."""
        if self._scheduler is not None:
            self._scheduler.shutdown(wait=False)
            self._scheduler = None

        await self._distributed_lock.dispose()
        set_scheduler_engine(None)
        logger.info("Scheduler stopped")

    async def _register_job(
        self,
        schedule_id: str,
        agent_id: str,
        sched_config: ScheduleConfig,
    ) -> None:
        """Register a single cron job with APScheduler."""
        if self._scheduler is None:
            return

        # Per-schedule timezone takes priority; fall back to gateway default
        tz = sched_config.timezone or self._timezone

        trigger = CronTrigger.from_crontab(sched_config.cron, timezone=tz)

        input_data = dict(sched_config.input) if sched_config.input else {}
        input_data["source"] = "scheduled"
        input_data["schedule_id"] = schedule_id
        input_data["schedule_name"] = sched_config.name
        if sched_config.instructions:
            input_data["_schedule_instructions"] = sched_config.instructions

        self._scheduler.add_job(
            run_scheduled_job,
            trigger=trigger,
            id=schedule_id,
            name=f"schedule:{schedule_id}",
            kwargs={
                "schedule_id": schedule_id,
                "agent_id": agent_id,
                "message": sched_config.message,
                "input": input_data,
            },
            replace_existing=True,
            coalesce=self._config.coalesce,
            misfire_grace_time=self._config.misfire_grace_seconds,
            max_instances=self._config.max_instances,
        )

        # Pause if disabled
        if not sched_config.enabled:
            self._scheduler.pause_job(schedule_id)

    async def _sync_schedule_records(
        self,
        agents: dict[str, AgentDefinition],
    ) -> None:
        """Sync workspace schedule configs to persistence ScheduleRecords.

        Batches all upserts into a single transaction to avoid O(N)
        sequential round-trips at startup.
        """
        records: list[ScheduleRecord] = []
        now = datetime.now(UTC)
        for agent in agents.values():
            for sched in agent.schedules:
                schedule_id = f"{agent.id}:{sched.name}"
                next_run = self._get_next_run_time(schedule_id)
                records.append(
                    ScheduleRecord(
                        id=schedule_id,
                        agent_id=agent.id,
                        name=sched.name,
                        cron_expr=sched.cron,
                        message=sched.message,
                        instructions=sched.instructions,
                        input=dict(sched.input) if sched.input else None,
                        enabled=sched.enabled,
                        timezone=sched.timezone or self._timezone,
                        source="workspace",
                        next_run_at=next_run,
                        created_at=now,
                    )
                )
        await self._schedule_repo.upsert_batch(records)

    async def _load_admin_schedules(self) -> None:
        """Load admin-created schedules from persistence and register as APScheduler jobs."""
        from agent_gateway.workspace.agent import ScheduleConfig as _ScheduleConfig

        records = await self._schedule_repo.list_all()
        count = 0
        for rec in records:
            if rec.source != "admin" or rec.deleted_at is not None:
                continue
            config = _ScheduleConfig(
                name=rec.name,
                cron=rec.cron_expr,
                message=rec.message,
                instructions=rec.instructions,
                input=rec.input or {},
                enabled=rec.enabled,
                timezone=rec.timezone,
            )
            self._schedule_configs[rec.id] = config
            self._agent_map[rec.id] = rec.agent_id
            await self._register_job(rec.id, rec.agent_id, config)
            count += 1
        if count:
            logger.info("Loaded %d admin schedules from persistence", count)

    def _get_next_run_time(self, schedule_id: str) -> datetime | None:
        """Get the next fire time for a schedule from APScheduler."""
        if self._scheduler is None:
            return None
        job = self._scheduler.get_job(schedule_id)
        if job is None:
            return None
        next_time: datetime | None = job.next_run_time
        return next_time

    async def dispatch_scheduled_execution(
        self,
        schedule_id: str,
        agent_id: str,
        message: str,
        input: dict[str, Any],
    ) -> None:
        """Dispatch a scheduled agent execution. Called by the handler module."""
        # Overlap check with safety valve: auto-clear after 2x misfire grace
        timeout = self._config.misfire_grace_seconds * 2
        async with self._active_lock:
            fire_time = self._active_scheduled.get(schedule_id)
            if fire_time is not None:
                elapsed = time.monotonic() - fire_time
                if elapsed < timeout:
                    logger.warning(
                        "Schedule '%s' still running (%.0fs), skipping overlapping fire",
                        schedule_id,
                        elapsed,
                    )
                    return
                logger.warning(
                    "Schedule '%s' stuck for %.0fs (timeout=%ds), force-clearing",
                    schedule_id,
                    elapsed,
                    timeout,
                )
            self._active_scheduled[schedule_id] = time.monotonic()

        # Distributed lock: prevent duplicate fires across instances.
        # Round to 60s boundary so all instances derive the same key
        # even with seconds of wall-clock jitter between them.
        fire_time_key = str(int(time.time()) // 60 * 60)
        lock_name = f"schedule:{schedule_id}:{fire_time_key}"
        lock_ttl = self._config.distributed_lock.lock_ttl_seconds

        acquired = await self._distributed_lock.acquire(lock_name, lock_ttl)
        if not acquired:
            logger.info(
                "Schedule '%s' lock held by another instance, skipping",
                schedule_id,
            )
            async with self._active_lock:
                self._active_scheduled.pop(schedule_id, None)
            return

        try:
            await self._do_dispatch(schedule_id, agent_id, message, input)
        except Exception:
            logger.exception("Scheduled execution for '%s' failed", schedule_id)
            async with self._active_lock:
                self._active_scheduled.pop(schedule_id, None)
            # Update last_run_at even on failure
            await self._update_after_run(schedule_id)
        finally:
            await self._distributed_lock.release(lock_name)

    async def _do_dispatch(
        self,
        schedule_id: str,
        agent_id: str,
        message: str,
        input: dict[str, Any],
    ) -> None:
        """Internal dispatch — either enqueue to worker pool or invoke directly."""
        from agent_gateway.persistence.domain import ExecutionRecord
        from agent_gateway.queue.null import NullQueue

        execution_id = str(uuid.uuid4())

        # Strip internal keys (prefixed with '_') from persisted input
        # to avoid leaking schedule_instructions / notify_config via the API.
        # The original dict with internal keys is still passed to invoke_fn.
        persisted_input = {k: v for k, v in input.items() if not k.startswith("_")}

        # user_id=None for scheduler-initiated executions
        record = ExecutionRecord(
            id=execution_id,
            agent_id=agent_id,
            status="queued",
            message=message,
            input=persisted_input,
            schedule_id=schedule_id,
            schedule_name=input.get("schedule_name", ""),
            created_at=datetime.now(UTC),
        )
        await self._execution_repo.create(record)

        if not isinstance(self._queue, NullQueue):
            from agent_gateway.queue.models import ExecutionJob

            job = ExecutionJob(
                execution_id=execution_id,
                agent_id=agent_id,
                message=message,
                input=input,
                enqueued_at=datetime.now(UTC).isoformat(),
                schedule_id=schedule_id,
            )
            await self._queue.enqueue(job)
            logger.info(
                "Scheduled execution '%s' enqueued: agent=%s, execution=%s",
                schedule_id,
                agent_id,
                execution_id,
            )
        else:
            try:
                await self._invoke_fn(agent_id, message, input=input)
            finally:
                async with self._active_lock:
                    self._active_scheduled.pop(schedule_id, None)
                await self._update_after_run(schedule_id)

    async def on_execution_complete(self, schedule_id: str) -> None:
        """Called by the worker pool when a scheduled execution finishes.

        Removes the schedule from the active set and updates last_run_at.
        """
        async with self._active_lock:
            self._active_scheduled.pop(schedule_id, None)
        await self._update_after_run(schedule_id)

    async def _update_after_run(self, schedule_id: str) -> None:
        """Update the schedule's last_run_at and next_run_at in persistence."""
        next_run = self._get_next_run_time(schedule_id)
        await self._schedule_repo.update_last_run(
            schedule_id=schedule_id,
            last_run_at=datetime.now(UTC),
            next_run_at=next_run,
        )

    # --- Public API for schedule management (used by API routes) ---

    async def pause(self, schedule_id: str) -> bool:
        """Pause a schedule. Returns True if the schedule was found."""
        if self._scheduler is None:
            return False
        job = self._scheduler.get_job(schedule_id)
        if job is None:
            return False
        self._scheduler.pause_job(schedule_id)
        await self._schedule_repo.update_enabled(schedule_id, False)
        return True

    async def resume(self, schedule_id: str) -> bool:
        """Resume a paused schedule. Returns True if the schedule was found."""
        if self._scheduler is None:
            return False
        job = self._scheduler.get_job(schedule_id)
        if job is None:
            return False
        self._scheduler.resume_job(schedule_id)
        await self._schedule_repo.update_enabled(schedule_id, True)
        # Update next_run_at after resume (don't touch last_run_at)
        next_run = self._get_next_run_time(schedule_id)
        if next_run:
            await self._schedule_repo.update_next_run(schedule_id, next_run)
        return True

    async def trigger(self, schedule_id: str) -> str | None:
        """Manually trigger a scheduled job. Returns execution_id or None."""
        config = self._schedule_configs.get(schedule_id)
        agent_id = self._agent_map.get(schedule_id)
        if config is None or agent_id is None:
            return None

        execution_id = str(uuid.uuid4())
        input_data: dict[str, Any] = dict(config.input) if config.input else {}
        input_data["source"] = "manual_trigger"
        input_data["schedule_id"] = schedule_id
        input_data["schedule_name"] = config.name
        if config.instructions:
            input_data["_schedule_instructions"] = config.instructions

        from agent_gateway.persistence.domain import ExecutionRecord

        persisted_input = {k: v for k, v in input_data.items() if not k.startswith("_")}
        # user_id=None for scheduler-initiated executions
        record = ExecutionRecord(
            id=execution_id,
            agent_id=agent_id,
            status="queued",
            message=config.message,
            input=persisted_input,
            schedule_id=schedule_id,
            schedule_name=config.name,
            created_at=datetime.now(UTC),
        )
        await self._execution_repo.create(record)

        # Dispatch via gateway invoke (direct, not queued)
        task = asyncio.create_task(
            self._run_manual_trigger(execution_id, agent_id, config.message, input_data),
            name=f"manual-trigger-{execution_id}",
        )
        self._track_task(task)
        return execution_id

    async def _run_manual_trigger(
        self,
        execution_id: str,
        agent_id: str,
        message: str,
        input: dict[str, Any],
    ) -> None:
        """Run a manually triggered schedule execution in the background."""
        try:
            result = await self._invoke_fn(agent_id, message, input=input)
            await self._execution_repo.update_status(
                execution_id,
                "completed",
                completed_at=datetime.now(UTC),
            )
            await self._execution_repo.update_result(
                execution_id,
                result=result.to_dict(),
                usage=result.usage.to_dict(),
            )
        except Exception as exc:
            logger.exception("Manual trigger execution %s failed", execution_id)
            await self._execution_repo.update_status(
                execution_id,
                "failed",
                error=str(exc),
            )

    async def get_schedules(self) -> list[dict[str, Any]]:
        """List all schedules with live next_run_at from APScheduler."""
        records = await self._schedule_repo.list_all()
        result = []
        for rec in records:
            # Get live next_run_at from APScheduler if available
            next_run = self._get_next_run_time(rec.id) or rec.next_run_at
            result.append(
                {
                    "id": rec.id,
                    "agent_id": rec.agent_id,
                    "name": rec.name,
                    "cron_expr": rec.cron_expr,
                    "enabled": rec.enabled,
                    "timezone": rec.timezone,
                    "source": rec.source,
                    "next_run_at": next_run,
                    "last_run_at": rec.last_run_at,
                    "created_at": rec.created_at,
                }
            )
        return result

    async def register_user_schedule(
        self,
        schedule_id: str,
        agent_id: str,
        cron_expr: str,
        message: str,
        input_data: dict[str, Any] | None = None,
        timezone: str = "UTC",
        enabled: bool = True,
        notify: dict[str, Any] | None = None,
        instructions: str | None = None,
    ) -> None:
        """Register a per-user schedule as a cron job."""
        if self._scheduler is None:
            return

        trigger = CronTrigger.from_crontab(cron_expr, timezone=timezone)

        job_input = dict(input_data) if input_data else {}
        job_input["source"] = "user_scheduled"
        job_input["schedule_id"] = schedule_id
        if notify:
            job_input["_notify_config"] = notify
        if instructions:
            job_input["_schedule_instructions"] = instructions

        self._scheduler.add_job(
            run_scheduled_job,
            trigger=trigger,
            id=schedule_id,
            name=f"user-schedule:{schedule_id}",
            kwargs={
                "schedule_id": schedule_id,
                "agent_id": agent_id,
                "message": message,
                "input": job_input,
            },
            replace_existing=True,
            coalesce=self._config.coalesce,
            misfire_grace_time=self._config.misfire_grace_seconds,
            max_instances=1,
        )

        if not enabled:
            self._scheduler.pause_job(schedule_id)

        self._agent_map[schedule_id] = agent_id

    async def remove_user_schedule(self, schedule_id: str) -> None:
        """Remove a per-user schedule from APScheduler."""
        if self._scheduler is None:
            return
        with contextlib.suppress(Exception):
            self._scheduler.remove_job(schedule_id)
        self._agent_map.pop(schedule_id, None)

    async def update_schedule(
        self,
        schedule_id: str,
        cron_expr: str | None = None,
        message: str | None = None,
        timezone: str | None = None,
        enabled: bool | None = None,
        instructions: str | None = None,
    ) -> bool:
        """Update a system schedule's configuration at runtime.

        Returns True if the schedule was found and updated.
        """
        if self._scheduler is None:
            return False

        config = self._schedule_configs.get(schedule_id)
        agent_id = self._agent_map.get(schedule_id)
        if config is None or agent_id is None:
            return False

        # Update in-memory ScheduleConfig dataclass
        if cron_expr is not None:
            config.cron = cron_expr
        if message is not None:
            config.message = message
        if timezone is not None:
            config.timezone = timezone
        if enabled is not None:
            config.enabled = enabled
        if instructions is not None:
            config.instructions = instructions or None

        # Remove and re-register the APScheduler job with new settings
        with contextlib.suppress(Exception):
            self._scheduler.remove_job(schedule_id)
        await self._register_job(schedule_id, agent_id, config)

        # Update persistence record
        next_run = self._get_next_run_time(schedule_id)
        await self._schedule_repo.update_schedule(
            schedule_id,
            cron_expr=config.cron,
            message=config.message,
            timezone=config.timezone or self._timezone,
            next_run_at=next_run,
            instructions=config.instructions,
        )

        if enabled is not None:
            await self._schedule_repo.update_enabled(schedule_id, enabled)

        return True

    async def get_schedule(self, schedule_id: str) -> dict[str, Any] | None:
        """Get a single schedule with full details."""
        rec = await self._schedule_repo.get(schedule_id)
        if rec is None or rec.deleted_at is not None:
            return None

        next_run = self._get_next_run_time(schedule_id) or rec.next_run_at
        return {
            "id": rec.id,
            "agent_id": rec.agent_id,
            "name": rec.name,
            "cron_expr": rec.cron_expr,
            "message": rec.message,
            "instructions": rec.instructions,
            "input": rec.input or {},
            "enabled": rec.enabled,
            "timezone": rec.timezone,
            "source": rec.source,
            "next_run_at": next_run,
            "last_run_at": rec.last_run_at,
            "created_at": rec.created_at,
        }

    async def create_admin_schedule(
        self,
        agent_id: str,
        name: str,
        cron_expr: str,
        message: str,
        instructions: str | None = None,
        input_data: dict[str, Any] | None = None,
        timezone: str = "UTC",
        enabled: bool = True,
    ) -> str:
        """Create a new admin schedule. Returns the schedule_id."""
        from agent_gateway.exceptions import ScheduleConflictError, ScheduleValidationError
        from agent_gateway.workspace.agent import ScheduleConfig as _ScheduleConfig

        # Validate agent exists
        if agent_id not in self._known_agent_ids:
            raise ScheduleValidationError(f"Unknown agent: {agent_id}")

        schedule_id = f"admin:{agent_id}:{name}"

        # Check for duplicates
        existing = await self._schedule_repo.get(schedule_id)
        if existing is not None and existing.deleted_at is None:
            raise ScheduleConflictError(f"Schedule '{name}' already exists for agent '{agent_id}'")

        # Validate cron expression
        try:
            CronTrigger.from_crontab(cron_expr, timezone=timezone)
        except (ValueError, KeyError) as e:
            raise ScheduleValidationError(str(e)) from e

        config = _ScheduleConfig(
            name=name,
            cron=cron_expr,
            message=message,
            instructions=instructions,
            input=input_data or {},
            enabled=enabled,
            timezone=timezone,
        )

        # Persist FIRST
        now = datetime.now(UTC)
        record = ScheduleRecord(
            id=schedule_id,
            agent_id=agent_id,
            name=name,
            cron_expr=cron_expr,
            message=message,
            instructions=instructions,
            input=input_data,
            enabled=enabled,
            timezone=timezone,
            source="admin",
            created_at=now,
        )
        await self._schedule_repo.upsert(record)

        # Then register in memory + APScheduler
        try:
            self._schedule_configs[schedule_id] = config
            self._agent_map[schedule_id] = agent_id
            await self._register_job(schedule_id, agent_id, config)
        except Exception:
            self._schedule_configs.pop(schedule_id, None)
            self._agent_map.pop(schedule_id, None)
            raise

        return schedule_id

    async def delete_admin_schedule(self, schedule_id: str) -> bool:
        """Delete an admin-created schedule. Returns False if not found or not admin-sourced."""
        rec = await self._schedule_repo.get(schedule_id)
        if rec is None or rec.source != "admin":
            return False

        # Remove from APScheduler
        if self._scheduler is not None:
            with contextlib.suppress(Exception):
                self._scheduler.remove_job(schedule_id)

        # Remove from in-memory maps
        self._schedule_configs.pop(schedule_id, None)
        self._agent_map.pop(schedule_id, None)

        # Soft-delete via dedicated method (NOT upsert, which resets deleted_at)
        await self._schedule_repo.soft_delete(schedule_id)
        return True
