# #region SchedulerModule [C:3] [TYPE Module] [SEMANTICS scheduler, schedule, scheduler-service]
# @BRIEF Manages scheduled tasks using APScheduler.
# @LAYER Core
# @RELATION DEPENDS_ON -> TaskManager
import asyncio
from datetime import date, datetime, time, timedelta

from apscheduler.schedulers.background import BackgroundScheduler
from apscheduler.triggers.cron import CronTrigger

from .config_manager import ConfigManager
from .cot_logger import seed_trace_id
from .database import SessionLocal
from .logger import belief_scope, logger


# #region SchedulerService [C:3] [TYPE Class] [SEMANTICS scheduler, service, apscheduler]
# @BRIEF Provides a service to manage scheduled backup tasks.
# @RELATION DEPENDS_ON -> [ThrottledSchedulerConfigurator]
class SchedulerService:
    # #region __init__ [TYPE Function]
    # @PURPOSE: Initializes the scheduler service with task and config managers.
    # @PRE  task_manager and config_manager must be provided.
    # @POST  Scheduler instance is created but not started.
    def __init__(self, task_manager, config_manager: ConfigManager):
        with belief_scope("SchedulerService.__init__"):
            self.task_manager = task_manager
            self.config_manager = config_manager
            self.scheduler = BackgroundScheduler()
            self.loop = asyncio.get_event_loop()
    # #endregion __init__
    # #region start [TYPE Function]
    # @PURPOSE: Starts the background scheduler and loads initial schedules.
    # @PRE  Scheduler should be initialized.
    # @POST  Scheduler is running and schedules are loaded.
    def start(self):
        with belief_scope("SchedulerService.start"):
            if not self.scheduler.running:
                self.scheduler.start()
                logger.info("Scheduler started.")
                self.load_schedules()
    # #endregion start
    # #region stop [TYPE Function]
    # @PURPOSE: Stops the background scheduler.
    # @PRE  Scheduler should be running.
    # @POST  Scheduler is shut down.
    def stop(self):
        with belief_scope("SchedulerService.stop"):
            if self.scheduler.running:
                self.scheduler.shutdown()
                logger.info("Scheduler stopped.")
    # #endregion stop
    # #region load_schedules [C:4] [TYPE Function] [SEMANTICS scheduler,backup,translation,apscheduler]
    # @BRIEF Load backup and active translation schedules from config and DB, re-registering all jobs.
    # @PRE config_manager must have valid configuration; database is accessible.
    # @POST All enabled backup jobs and active translation schedules are re-registered in APScheduler.
    # @SIDE_EFFECT Removes all existing APScheduler jobs; queries translation_schedules table; registers APScheduler jobs for translation.
    def load_schedules(self):
        with belief_scope("SchedulerService.load_schedules"):
            # Clear existing jobs
            self.scheduler.remove_all_jobs()
            config = self.config_manager.get_config()
            for env in config.environments:
                if env.backup_schedule and env.backup_schedule.enabled:
                    self.add_backup_job(env.id, env.backup_schedule.cron_expression)
            # Re-register active translation schedules from the database
            try:
                from ..models.translate import TranslationSchedule
                db = SessionLocal()
                try:
                    active_schedules = (
                        db.query(TranslationSchedule)
                        .filter(TranslationSchedule.is_active == True)
                        .all()
                    )
                    for sched in active_schedules:
                        self.add_translation_job(
                            schedule_id=sched.id,
                            job_id=sched.job_id,
                            cron_expression=sched.cron_expression,
                            timezone=sched.timezone,
                            execution_mode=getattr(sched, 'execution_mode', 'full'),
                        )
                    if active_schedules:
                        logger.reason(f"Loaded {len(active_schedules)} active translation schedule(s)")
                finally:
                    db.close()
            except Exception as e:
                logger.explore("Failed to load translation schedules on startup",
                               extra={"error": str(e)})
            # Re-register active validation policies with schedules
            try:
                from ..models.llm import ValidationPolicy
                db = SessionLocal()
                try:
                    active_policies = (
                        db.query(ValidationPolicy)
                        .filter(
                            ValidationPolicy.is_active.is_(True),
                            ValidationPolicy.schedule_days.isnot(None),
                        )
                        .all()
                    )
                    for policy in active_policies:
                        schedule_days = policy.schedule_days or []
                        if not schedule_days:
                            continue
                        ws = policy.window_start
                        if ws is None:
                            continue
                        # Build cron: minute hour * * day-of-week
                        dow = ",".join(str(d) for d in sorted(schedule_days))
                        cron = f"{ws.minute} {ws.hour} * * {dow}"
                        self.add_validation_job(
                            policy_id=policy.id,
                            cron_expression=cron,
                            app_timezone=config.settings.app_timezone,
                        )
                    if active_policies:
                        logger.reason(f"Loaded {len(active_policies)} active validation schedule(s)")
                finally:
                    db.close()
            except Exception as e:
                logger.explore("Failed to load validation schedules on startup",
                               extra={"error": str(e)})
    # #endregion load_schedules
    # #region add_backup_job [TYPE Function]
    # @PURPOSE: Adds a scheduled backup job for an environment.
    # @PRE  env_id and cron_expression must be valid strings.
    # @POST  A new job is added to the scheduler or replaced if it already exists.
    # @PARAM  env_id (str) - The ID of the environment.
    # @PARAM  cron_expression (str) - The cron expression for the schedule.
    def add_backup_job(self, env_id: str, cron_expression: str):
        with belief_scope(
            "SchedulerService.add_backup_job",
            f"env_id={env_id}, cron={cron_expression}",
        ):
            job_id = f"backup_{env_id}"
            try:
                self.scheduler.add_job(
                    self._trigger_backup,
                    CronTrigger.from_crontab(cron_expression),
                    id=job_id,
                    args=[env_id],
                    replace_existing=True,
                )
                logger.info(
                    f"Scheduled backup job added for environment {env_id}: {cron_expression}"
                )
            except Exception as e:
                logger.error(f"Failed to add backup job for environment {env_id}: {e}")
    # #endregion add_backup_job
    # #region add_translation_job [C:4] [TYPE Function] [SEMANTICS scheduler,translation,cron,apscheduler]
    # @BRIEF Register a translation schedule with APScheduler.
    # @PRE schedule_id, job_id, and cron_expression are valid strings.
    # @POST A new APScheduler job is registered or replaced if it already exists.
    # @SIDE_EFFECT Mutates APScheduler state; calls execute_scheduled_translation on trigger.
    # @RELATION DEPENDS_ON -> [execute_scheduled_translation]
    # @RELATION DEPENDS_ON -> [SessionLocal]
    def add_translation_job(self, schedule_id: str, job_id: str, cron_expression: str, timezone: str = "UTC", execution_mode: str = "full"):
        with belief_scope(
            "SchedulerService.add_translation_job",
            f"schedule_id={schedule_id}, job_id={job_id}, cron={cron_expression}",
        ):
            from zoneinfo import ZoneInfo

            from ..plugins.translate.scheduler import execute_scheduled_translation

            job_id_aps = f"translate_{schedule_id}"
            try:
                tz = ZoneInfo(timezone)
                self.scheduler.add_job(
                    execute_scheduled_translation,
                    CronTrigger.from_crontab(cron_expression, timezone=tz),
                    id=job_id_aps,
                    args=[schedule_id, job_id, SessionLocal, self.config_manager, execution_mode],
                    replace_existing=True,
                )
                logger.reason(
                    f"Translation schedule registered: {job_id_aps} ({cron_expression})"
                )
            except Exception as e:
                logger.explore(
                    f"Failed to register translation schedule: {job_id_aps}",
                    extra={"error": str(e)},
                )
    # #endregion add_translation_job
    # #region remove_translation_job [C:4] [TYPE Function] [SEMANTICS scheduler,translation,remove,apscheduler]
    # @BRIEF Remove a translation schedule from APScheduler.
    # @PRE schedule_id is a valid string.
    # @POST The APScheduler job is removed if it exists; silently ignored otherwise.
    # @SIDE_EFFECT Mutates APScheduler state.
    def remove_translation_job(self, schedule_id: str):
        with belief_scope(
            "SchedulerService.remove_translation_job",
            f"schedule_id={schedule_id}",
        ):
            job_id_aps = f"translate_{schedule_id}"
            try:
                self.scheduler.remove_job(job_id_aps)
                logger.info(f"Translation schedule removed: {job_id_aps}")
            except Exception:
                logger.reason(f"Translation schedule not found (already removed): {job_id_aps}")
    # #endregion remove_translation_job
    # #region _trigger_backup [TYPE Function]
    # @PURPOSE: Triggered by the scheduler to start a backup task.
    # @PRE  env_id must be a valid environment ID.
    # @POST  A new backup task is created in the task manager if not already running.
    # @PARAM  env_id (str) - The ID of the environment.
    def _trigger_backup(self, env_id: str):
        seed_trace_id()
        with belief_scope("SchedulerService._trigger_backup", f"env_id={env_id}"):
            logger.info(f"Triggering scheduled backup for environment {env_id}")
            # Check if a backup is already running for this environment
            active_tasks = self.task_manager.get_tasks(limit=100)
            for task in active_tasks:
                if (
                    task.plugin_id == "superset-backup"
                    and task.status in ["PENDING", "RUNNING"]
                    and task.params.get("environment_id") == env_id
                ):
                    logger.warning(
                        f"Backup already running for environment {env_id}. Skipping scheduled run."
                    )
                    return
            # Run the backup task
            # We need to run this in the event loop since create_task is async
            asyncio.run_coroutine_threadsafe(
                self.task_manager.create_task(
                    "superset-backup", {"environment_id": env_id}
                ),
                self.loop,
            )
    # #endregion _trigger_backup
    # #region add_validation_job [C:3] [TYPE Function] [SEMANTICS scheduler,validation,cron,apscheduler]
    # @BRIEF Register a validation policy schedule with APScheduler.
    # @PRE policy_id and cron_expression are valid strings.
    # @POST A new APScheduler job is registered or replaced if it already exists.
    # @SIDE_EFFECT Mutates APScheduler state; calls _trigger_validation on trigger.
    # @RELATION DEPENDS_ON -> [_trigger_validation]
    def add_validation_job(self, policy_id: str, cron_expression: str, app_timezone: str = "Europe/Moscow") -> None:
        with belief_scope(
            "SchedulerService.add_validation_job",
            f"policy_id={policy_id}, cron={cron_expression}",
        ):
            from zoneinfo import ZoneInfo

            job_id_aps = f"validation_{policy_id}"
            try:
                tz = ZoneInfo(app_timezone)
                self.scheduler.add_job(
                    self._trigger_validation,
                    CronTrigger.from_crontab(cron_expression, timezone=tz),
                    id=job_id_aps,
                    args=[policy_id],
                    replace_existing=True,
                )
                logger.reason(
                    f"Validation schedule registered: {job_id_aps} ({cron_expression}) [{app_timezone}]"
                )
            except Exception as e:
                logger.explore(
                    f"Failed to register validation schedule: {job_id_aps}",
                    extra={"error": str(e)},
                )
    # #endregion add_validation_job
    # #region remove_validation_job [C:2] [TYPE Function] [SEMANTICS scheduler,validation,remove,apscheduler]
    # @BRIEF Remove a validation policy schedule from APScheduler.
    # @PRE policy_id is a valid string.
    # @POST The APScheduler job is removed if it exists; silently ignored otherwise.
    # @SIDE_EFFECT Mutates APScheduler state.
    def remove_validation_job(self, policy_id: str) -> None:
        with belief_scope(
            "SchedulerService.remove_validation_job",
            f"policy_id={policy_id}",
        ):
            job_id_aps = f"validation_{policy_id}"
            try:
                self.scheduler.remove_job(job_id_aps)
                logger.info(f"Validation schedule removed: {job_id_aps}")
            except Exception:
                logger.reason(f"Validation schedule not found (already removed): {job_id_aps}")
    # #endregion remove_validation_job
    # #region reload_validation_policy [C:2] [TYPE Function] [SEMANTICS scheduler,validation,reload,apscheduler]
    # @BRIEF Reload a single validation policy schedule — calls remove then add.
    # @PRE policy_id is a valid ValidationPolicy id with schedule data in DB.
    # @POST Old job is removed; new job is registered if policy is active and has schedule_days.
    # @SIDE_EFFECT Mutates APScheduler state; queries ValidationPolicy from DB.
    def reload_validation_policy(self, policy_id: str, app_timezone: str = "Europe/Moscow") -> None:
        with belief_scope(
            "SchedulerService.reload_validation_policy",
            f"policy_id={policy_id}",
        ):
            self.remove_validation_job(policy_id)
            try:
                from ..models.llm import ValidationPolicy
                db = SessionLocal()
                try:
                    policy = db.query(ValidationPolicy).filter(ValidationPolicy.id == policy_id).first()
                    if not policy:
                        logger.reason(f"Policy {policy_id} not found; schedule removed")
                        return
                    if not policy.is_active:
                        logger.reason(f"Policy {policy_id} is inactive; schedule removed")
                        return
                    schedule_days = policy.schedule_days or []
                    if not schedule_days or policy.window_start is None:
                        logger.reason(f"Policy {policy_id} has no schedule; nothing to register")
                        return
                    dow = ",".join(str(d) for d in sorted(schedule_days))
                    ws = policy.window_start
                    cron = f"{ws.minute} {ws.hour} * * {dow}"
                    self.add_validation_job(policy_id, cron, app_timezone=app_timezone)
                finally:
                    db.close()
            except Exception as e:
                logger.explore(
                    f"Failed to reload validation policy {policy_id}",
                    extra={"error": str(e)},
                )
    # #endregion reload_validation_policy
    # #region _trigger_validation [C:3] [TYPE Function] [SEMANTICS scheduler,validation,trigger]
    # @BRIEF APScheduler job handler — triggers validation runs for a policy.
    # @PRE policy_id is a valid ValidationPolicy id with dashboard_ids.
    # @POST A validation task is spawned via TaskManager for each dashboard in the policy.
    # @SIDE_EFFECT Creates TaskRecord entries for validation runs; logs progress.
    def _trigger_validation(self, policy_id: str) -> None:
        seed_trace_id()
        with belief_scope("SchedulerService._trigger_validation", f"policy_id={policy_id}"):
            logger.reason(f"Triggering scheduled validation for policy {policy_id}")
            db = SessionLocal()
            try:
                from ..models.llm import ValidationPolicy

                policy = db.query(ValidationPolicy).filter(ValidationPolicy.id == policy_id).first()
                if not policy:
                    logger.explore(f"Validation policy not found: {policy_id}")
                    return
                if not policy.is_active:
                    logger.reason(f"Validation policy {policy_id} is no longer active; skipping")
                    return
                dashboard_ids = list(policy.dashboard_ids or [])
                if not dashboard_ids:
                    logger.explore(f"Validation policy {policy_id} has no dashboards; skipping")
                    return

                ws = policy.window_start
                we = policy.window_end
                today = date.today()
                scheduled_times = ThrottledSchedulerConfigurator.calculate_schedule(
                    window_start=ws,
                    window_end=we or time(ws.hour, ws.minute + 30),
                    dashboard_ids=dashboard_ids,
                    current_date=today,
                )

                for idx, dash_id in enumerate(dashboard_ids):
                    sched_time = scheduled_times[idx] if idx < len(scheduled_times) else scheduled_times[-1]
                    params: dict = {
                        "dashboard_id": dash_id,
                        "environment_id": policy.environment_id,
                        "provider_id": policy.provider_id or "",
                        "policy_id": policy_id,
                    }
                    asyncio.run_coroutine_threadsafe(
                        self.task_manager.create_task(
                            plugin_id="llm_dashboard_validation", params=params
                        ),
                        self.loop,
                    )
                    logger.reason(
                        f"Scheduled validation for dashboard {dash_id} at {sched_time.isoformat()}"
                    )
            except Exception as e:
                logger.explore(
                    f"Error triggering validation for policy {policy_id}",
                    extra={"error": str(e)},
                )
            finally:
                db.close()
    # #endregion _trigger_validation
# #endregion SchedulerService
# #region ThrottledSchedulerConfigurator [C:5] [TYPE Class] [SEMANTICS scheduler, throttling, distribution]
# @BRIEF Distributes validation tasks evenly within an execution window.
# @PRE  Validation policies provide a finite dashboard list and a valid execution window.
# @POST  Produces deterministic per-dashboard run timestamps within the configured window.
# @RELATION DEPENDS_ON -> SchedulerModule
# @INVARIANT  Returned schedule size always matches number of dashboard IDs.
# @SIDE_EFFECT  Emits warning logs for degenerate or near-zero scheduling windows.
# @DATA_CONTRACT  Input[window_start, window_end, dashboard_ids, current_date] -> Output[List[EXT:Python:datetime]]
class ThrottledSchedulerConfigurator:
    # #region calculate_schedule [TYPE Function]
    # @PURPOSE: Calculates execution times for N tasks within a window.
    # @PRE  window_start, window_end (time), dashboard_ids (List), current_date (date).
    # @POST  Returns List[EXT:Python:datetime] of scheduled times.
    # @INVARIANT  Tasks are distributed with near-even spacing.
    @staticmethod
    def calculate_schedule(
        window_start: time, window_end: time, dashboard_ids: list, current_date: date
    ) -> list:
        with belief_scope("ThrottledSchedulerConfigurator.calculate_schedule"):
            n = len(dashboard_ids)
            if n == 0:
                return []
            start_dt = datetime.combine(current_date, window_start)
            end_dt = datetime.combine(current_date, window_end)
            # Handle window crossing midnight
            if end_dt < start_dt:
                end_dt += timedelta(days=1)
            total_seconds = (end_dt - start_dt).total_seconds()
            # Minimum interval of 1 second to avoid division by zero or negative
            if total_seconds <= 0:
                logger.warning(
                    f"[calculate_schedule] Window size is zero or negative. Falling back to start time for all {n} tasks."
                )
                return [start_dt] * n
            # If window is too small for even distribution (e.g. 10 tasks in 5 seconds),
            # we still distribute them but they might be very close.
            # The requirement says "near-even spacing".
            if n == 1:
                return [start_dt]
            interval = total_seconds / (n - 1) if n > 1 else 0
            # If interval is too small (e.g. < 1s), we might want a fallback,
            # but the spec says "handle too-small windows with explicit fallback/warning".
            if interval < 1:
                logger.warning(
                    f"[calculate_schedule] Window too small for {n} tasks (interval {interval:.2f}s). Tasks will be highly concentrated."
                )
            scheduled_times = []
            for i in range(n):
                scheduled_times.append(start_dt + timedelta(seconds=i * interval))
            return scheduled_times
    # #endregion calculate_schedule
# #endregion ThrottledSchedulerConfigurator
# #endregion SchedulerModule
