"""Main scheduler loop — creates DagRuns, dispatches tasks, collects results."""
from __future__ import annotations

import asyncio
import json
import logging
from datetime import datetime, timezone

from sqlalchemy import select, update

from crondag.config import Settings
from crondag.dag_bag import DagBag
from crondag.executor.thread import ThreadExecutor
from crondag.models.base import get_session_factory
from crondag.models.dag import DagModel
from crondag.models.dag_run import DagRunModel
from crondag.models.task_instance import TaskInstanceModel
from crondag.models.log import TaskLogModel
from crondag.scheduler.dependency import evaluate_task_readiness
from crondag.scheduler.evaluator import is_due, next_run_time
from crondag.sdk.dag import DAG
from crondag.sdk.schedule import Cron, Interval
from crondag.sdk.state import DagRunState, TaskState
from crondag.sdk.trigger_rule import TriggerRule

logger = logging.getLogger(__name__)


class Scheduler:
    """The main scheduler that orchestrates DAG runs and task execution."""

    def __init__(self, settings: Settings):
        self.settings = settings
        self.dag_bag = DagBag(settings.dags_dir)
        self.executor = ThreadExecutor(max_workers=settings.scheduler_max_active_tasks)
        self._session_factory = get_session_factory(settings.database_url)
        self._running_tasks: dict[int, asyncio.Task] = {}  # ti.id -> asyncio.Task
        self._stop_event = asyncio.Event()

    async def start(self):
        """Run the scheduler loop."""
        logger.info("Scheduler started")
        # Initial DAG scan
        await self._sync_dags()

        while not self._stop_event.is_set():
            try:
                await self._tick()
            except Exception:
                logger.exception("Error in scheduler tick")

            try:
                await asyncio.wait_for(self._stop_event.wait(), timeout=self.settings.scheduler_loop_interval)
                break  # Stop event was set
            except asyncio.TimeoutError:
                pass  # Normal timeout, continue loop

        await self.executor.shutdown()
        logger.info("Scheduler stopped")

    async def stop(self):
        """Signal the scheduler to stop."""
        self._stop_event.set()

    async def _tick(self):
        """One iteration of the scheduler loop."""
        # 1. Create DagRuns for DAGs that are due
        await self._create_dag_runs()
        # 2. Evaluate task dependencies and schedule ready tasks
        await self._evaluate_tasks()
        # 3. Dispatch scheduled tasks to executor
        await self._dispatch_tasks()
        # 4. Collect completed tasks
        await self._collect_results()
        # 5. Finalize completed DagRuns
        await self._finalize_dag_runs()

    async def _sync_dags(self):
        """Scan DAG files and sync to database."""
        dags = self.dag_bag.collect_dags()
        now = datetime.now(timezone.utc)

        async with self._session_factory() as session:
            for dag_id, dag_obj in dags.items():
                existing = await session.get(DagModel, dag_id)

                sched_type = None
                sched_value = None
                if isinstance(dag_obj.schedule, Cron):
                    sched_type = "cron"
                    sched_value = dag_obj.schedule.expression
                elif isinstance(dag_obj.schedule, Interval):
                    sched_type = "interval"
                    sched_value = str(int(dag_obj.schedule.delta.total_seconds()))

                if existing:
                    existing.description = dag_obj.description
                    existing.schedule_type = sched_type
                    existing.schedule_value = sched_value
                    existing.fileloc = dag_obj.fileloc
                    existing.last_parsed_at = now
                    existing.tags = json.dumps(dag_obj.tags)
                    # Recompute next_run_at if schedule changed
                    if sched_type and not existing.next_run_at:
                        existing.next_run_at = next_run_time(dag_obj.schedule, now)
                else:
                    nrt = next_run_time(dag_obj.schedule, now) if dag_obj.schedule else None
                    dag_model = DagModel(
                        dag_id=dag_id,
                        description=dag_obj.description,
                        schedule_type=sched_type,
                        schedule_value=sched_value,
                        is_paused=False,
                        fileloc=dag_obj.fileloc,
                        next_run_at=nrt,
                        last_parsed_at=now,
                        tags=json.dumps(dag_obj.tags),
                    )
                    session.add(dag_model)

            await session.commit()
        logger.info("Synced %d DAGs to database", len(dags))

    async def _create_dag_runs(self):
        """Create DagRuns for DAGs whose next_run_at has passed."""
        now = datetime.now(timezone.utc)

        async with self._session_factory() as session:
            result = await session.execute(
                select(DagModel).where(
                    DagModel.is_paused == False,
                    DagModel.next_run_at != None,
                    DagModel.next_run_at <= now,
                )
            )
            due_dags = result.scalars().all()

            for dag_model in due_dags:
                dag_obj = self.dag_bag.dags.get(dag_model.dag_id)
                if not dag_obj:
                    continue

                # Check max active runs
                active_runs = await session.execute(
                    select(DagRunModel).where(
                        DagRunModel.dag_id == dag_model.dag_id,
                        DagRunModel.state.in_([DagRunState.QUEUED, DagRunState.RUNNING]),
                    )
                )
                if len(active_runs.scalars().all()) >= self.settings.scheduler_max_active_runs_per_dag:
                    continue

                # Create the DagRun
                logical_date = dag_model.next_run_at
                run_id = f"scheduled__{logical_date.strftime('%Y-%m-%dT%H:%M:%S')}"

                dag_run = DagRunModel(
                    dag_id=dag_model.dag_id,
                    run_id=run_id,
                    state=DagRunState.QUEUED,
                    run_type="scheduled",
                    logical_date=logical_date,
                    created_at=now,
                )
                session.add(dag_run)

                # Create TaskInstances for all tasks in the DAG
                for task_id in dag_obj.get_task_ids():
                    task_obj = dag_obj.get_task(task_id)
                    ti = TaskInstanceModel(
                        task_id=task_id,
                        dag_id=dag_model.dag_id,
                        run_id=run_id,
                        state=TaskState.PENDING,
                        try_number=1,
                        max_retries=task_obj.retries,
                    )
                    session.add(ti)

                # Update next_run_at
                dag_model.next_run_at = next_run_time(dag_obj.schedule, logical_date)

                logger.info("Created DagRun %s for %s", run_id, dag_model.dag_id)

            await session.commit()

    async def _evaluate_tasks(self):
        """Evaluate task dependencies, move PENDING → SCHEDULED or UPSTREAM_FAILED."""
        async with self._session_factory() as session:
            # Get all running DagRuns
            result = await session.execute(
                select(DagRunModel).where(
                    DagRunModel.state.in_([DagRunState.QUEUED, DagRunState.RUNNING])
                )
            )
            active_runs = result.scalars().all()

            for dag_run in active_runs:
                dag_obj = self.dag_bag.dags.get(dag_run.dag_id)
                if not dag_obj:
                    continue

                # Mark QUEUED runs as RUNNING
                if dag_run.state == DagRunState.QUEUED:
                    dag_run.state = DagRunState.RUNNING
                    dag_run.start_date = datetime.now(timezone.utc)

                dep_map = dag_obj.get_dependency_map()

                # Get all task instances for this run
                ti_result = await session.execute(
                    select(TaskInstanceModel).where(
                        TaskInstanceModel.dag_id == dag_run.dag_id,
                        TaskInstanceModel.run_id == dag_run.run_id,
                    )
                )
                task_instances = {ti.task_id: ti for ti in ti_result.scalars().all()}

                for task_id, ti in task_instances.items():
                    if ti.state != TaskState.PENDING:
                        continue

                    upstream_ids = dep_map.get(task_id, [])
                    upstream_states = {}
                    for uid in upstream_ids:
                        upstream_ti = task_instances.get(uid)
                        if upstream_ti:
                            upstream_states[uid] = TaskState(upstream_ti.state)

                    task_obj = dag_obj.get_task(task_id)
                    trigger_rule = task_obj.trigger_rule if task_obj else TriggerRule.ALL_SUCCESS

                    new_state = evaluate_task_readiness(
                        task_id=task_id,
                        current_state=TaskState(ti.state),
                        trigger_rule=trigger_rule,
                        upstream_states=upstream_states,
                    )

                    if new_state != TaskState(ti.state):
                        ti.state = new_state
                        logger.debug("Task %s.%s: %s → %s", dag_run.dag_id, task_id, ti.state, new_state)

            await session.commit()

    async def _dispatch_tasks(self):
        """Dispatch SCHEDULED tasks to the executor."""
        # Limit concurrent tasks
        available_slots = self.settings.scheduler_max_active_tasks - len(self._running_tasks)
        if available_slots <= 0:
            return

        async with self._session_factory() as session:
            result = await session.execute(
                select(TaskInstanceModel)
                .where(TaskInstanceModel.state == TaskState.SCHEDULED)
                .limit(available_slots)
            )
            scheduled_tasks = result.scalars().all()

            for ti in scheduled_tasks:
                dag_obj = self.dag_bag.dags.get(ti.dag_id)
                if not dag_obj:
                    continue

                task_obj = dag_obj.get_task(ti.task_id)
                if not task_obj:
                    continue

                # Mark as RUNNING
                ti.state = TaskState.RUNNING
                ti.start_date = datetime.now(timezone.utc)

                # Resolve arguments from upstream return values
                dep_map = dag_obj.get_dependency_map()
                upstream_ids = dep_map.get(ti.task_id, [])
                args = []
                for uid in upstream_ids:
                    upstream_result = await session.execute(
                        select(TaskInstanceModel).where(
                            TaskInstanceModel.dag_id == ti.dag_id,
                            TaskInstanceModel.run_id == ti.run_id,
                            TaskInstanceModel.task_id == uid,
                        )
                    )
                    upstream_ti = upstream_result.scalar_one_or_none()
                    if upstream_ti and upstream_ti.return_value:
                        try:
                            args.append(json.loads(upstream_ti.return_value))
                        except json.JSONDecodeError:
                            args.append(upstream_ti.return_value)

                ti_id = ti.id
                # Launch async task for execution
                async_task = asyncio.create_task(
                    self._execute_task(ti_id, task_obj, args)
                )
                self._running_tasks[ti_id] = async_task

                logger.info("Dispatched task %s.%s (ti_id=%d)", ti.dag_id, ti.task_id, ti_id)

            await session.commit()

    async def _execute_task(self, ti_id: int, task_obj, args: list):
        """Execute a task in the thread pool and update the database."""
        try:
            result = await self.executor.submit(task_obj.execute, *args)

            # Update task instance
            async with self._session_factory() as session:
                ti = await session.get(TaskInstanceModel, ti_id)
                if ti:
                    now = datetime.now(timezone.utc)
                    ti.state = TaskState.SUCCESS
                    ti.end_date = now
                    if ti.start_date:
                        ti.duration = (now - ti.start_date).total_seconds()
                    if result is not None:
                        ti.return_value = json.dumps(result)

                    # Log success
                    log_entry = TaskLogModel(
                        task_instance_id=ti_id,
                        line_number=1,
                        timestamp=now,
                        level="INFO",
                        message=f"Task {ti.task_id} completed successfully",
                    )
                    session.add(log_entry)
                    await session.commit()

            logger.info("Task ti_id=%d succeeded", ti_id)

        except Exception as e:
            async with self._session_factory() as session:
                ti = await session.get(TaskInstanceModel, ti_id)
                if ti:
                    now = datetime.now(timezone.utc)
                    ti.end_date = now
                    if ti.start_date:
                        ti.duration = (now - ti.start_date).total_seconds()
                    ti.error = str(e)

                    # Check retries
                    if ti.try_number < ti.max_retries + 1:
                        ti.state = TaskState.SCHEDULED
                        ti.try_number += 1
                        ti.start_date = None
                        ti.end_date = None
                        ti.duration = None
                        logger.info("Task ti_id=%d failed, retrying (%d/%d)", ti_id, ti.try_number, ti.max_retries + 1)
                    else:
                        ti.state = TaskState.FAILED
                        logger.error("Task ti_id=%d failed: %s", ti_id, e)

                    # Log failure
                    log_entry = TaskLogModel(
                        task_instance_id=ti_id,
                        line_number=1,
                        timestamp=now,
                        level="ERROR",
                        message=f"Task failed: {e}",
                    )
                    session.add(log_entry)
                    await session.commit()

        finally:
            self._running_tasks.pop(ti_id, None)

    async def _collect_results(self):
        """Check for completed running tasks (handled by _execute_task callbacks)."""
        # Clean up finished tasks
        done = [tid for tid, t in self._running_tasks.items() if t.done()]
        for tid in done:
            task = self._running_tasks.pop(tid)
            if task.exception():
                logger.error("Task runner ti_id=%d had unhandled error: %s", tid, task.exception())

    async def _finalize_dag_runs(self):
        """Check if DagRuns are complete and update their state."""
        async with self._session_factory() as session:
            result = await session.execute(
                select(DagRunModel).where(DagRunModel.state == DagRunState.RUNNING)
            )
            running_runs = result.scalars().all()

            for dag_run in running_runs:
                ti_result = await session.execute(
                    select(TaskInstanceModel).where(
                        TaskInstanceModel.dag_id == dag_run.dag_id,
                        TaskInstanceModel.run_id == dag_run.run_id,
                    )
                )
                task_instances = ti_result.scalars().all()

                if not task_instances:
                    continue

                all_done = all(
                    TaskState(ti.state) in {TaskState.SUCCESS, TaskState.FAILED, TaskState.UPSTREAM_FAILED}
                    for ti in task_instances
                )

                if not all_done:
                    continue

                any_failed = any(
                    TaskState(ti.state) in {TaskState.FAILED, TaskState.UPSTREAM_FAILED}
                    for ti in task_instances
                )

                # Also check no tasks are still scheduled for retry
                any_retrying = any(
                    TaskState(ti.state) == TaskState.SCHEDULED
                    for ti in task_instances
                )
                if any_retrying:
                    continue

                now = datetime.now(timezone.utc)
                dag_run.end_date = now
                dag_run.state = DagRunState.FAILED if any_failed else DagRunState.SUCCESS
                logger.info("DagRun %s finalized: %s", dag_run.run_id, dag_run.state)

            await session.commit()

    async def trigger_dag(self, dag_id: str, conf: dict | None = None) -> DagRunModel | None:
        """Manually trigger a DAG run."""
        dag_obj = self.dag_bag.dags.get(dag_id)
        if not dag_obj:
            logger.error("DAG %s not found", dag_id)
            return None

        now = datetime.now(timezone.utc)
        run_id = f"manual__{now.strftime('%Y-%m-%dT%H:%M:%S.%f')}"

        async with self._session_factory() as session:
            dag_run = DagRunModel(
                dag_id=dag_id,
                run_id=run_id,
                state=DagRunState.QUEUED,
                run_type="manual",
                logical_date=now,
                conf=json.dumps(conf or {}),
                created_at=now,
            )
            session.add(dag_run)

            for task_id in dag_obj.get_task_ids():
                task_obj = dag_obj.get_task(task_id)
                ti = TaskInstanceModel(
                    task_id=task_id,
                    dag_id=dag_id,
                    run_id=run_id,
                    state=TaskState.PENDING,
                    try_number=1,
                    max_retries=task_obj.retries,
                )
                session.add(ti)

            await session.commit()
            await session.refresh(dag_run)

        logger.info("Manually triggered DAG %s: %s", dag_id, run_id)
        return dag_run

    async def rescan_dags(self):
        """Re-scan DAG files."""
        await self._sync_dags()
