"""Main scheduling loop — ties scoring model, policies, and executor together."""

from __future__ import annotations

import asyncio
import json
import logging

from radix_edge.config import get_config
from radix_edge.db import get_db
from radix_edge.executor.gpu_allocator import allocate, get_available_gpus
from radix_edge.executor.executor import run_job
from radix_edge.models import Job
from radix_edge.scheduler.scoring_model import SchedulerModel, JobFeatures
from radix_edge.scheduler.policies import order_jobs

logger = logging.getLogger("radix_edge.scheduler")

# Singleton model instance
_model: SchedulerModel | None = None


def get_model() -> SchedulerModel:
    """Get or create the singleton scoring model."""
    global _model
    if _model is None:
        _model = SchedulerModel()
    return _model


async def scheduling_loop():
    """Background task: pick queued jobs, score, allocate GPUs, dispatch."""
    config = get_config()
    model = get_model()
    checkpoint_counter = 0

    logger.info("Scheduler started (policy=%s)", config.scheduler_policy)

    while True:
        try:
            dispatched = 0

            with get_db() as db:
                # Get queued jobs
                rows = db.execute(
                    "SELECT * FROM jobs WHERE status = 'queued' ORDER BY created_at"
                ).fetchall()
                if not rows:
                    await asyncio.sleep(1.0)
                    continue

                queued = [Job.from_row(r) for r in rows]

                # Get available GPUs
                available = get_available_gpus(db)

                # Check if there are any jobs that can run (GPU jobs need
                # available GPUs; CPU-only jobs can always run)
                has_cpu_jobs = any(j.num_gpus == 0 for j in queued)
                if not available and not has_cpu_jobs:
                    await asyncio.sleep(1.0)
                    continue

                # Get distinct GPU types for scoring
                gpu_rows = db.execute(
                    "SELECT DISTINCT name FROM gpus WHERE status = 'idle'"
                ).fetchall()
                gpu_types = [r["name"] for r in gpu_rows] or ["generic"]

                # Order by policy
                ordered = order_jobs(queued, config.scheduler_policy, db)

                # Score and dispatch
                for job in ordered:
                    # CPU-only jobs: skip GPU allocation entirely
                    if job.num_gpus == 0:
                        db.execute(
                            "UPDATE jobs SET status = 'scheduled', scheduled_at = strftime('%Y-%m-%dT%H:%M:%fZ','now'), scheduled_gpus = '[]' WHERE job_id = ?",
                            (job.job_id,),
                        )
                        db.commit()
                        asyncio.create_task(run_job(job.job_id, []))
                        dispatched += 1
                        logger.info("Scheduled CPU-only job %s (no GPUs)", job.job_id)
                        continue

                    if not available:
                        break

                    if job.num_gpus > len(available):
                        continue  # Not enough GPUs for this job

                    # Score using IT model
                    features = JobFeatures(
                        job_type=job.job_type,
                        gpu_mem_gb=job.gpu_mem_required_mb / 1024.0,
                        num_gpus=job.num_gpus,
                        tenant=job.user_id,
                    )
                    result = model.score_job(features, gpu_types)

                    # Allocate GPUs
                    indices = allocate(job.job_id, job.num_gpus)
                    if indices is None:
                        continue

                    # Mark as scheduled
                    db.execute(
                        "UPDATE jobs SET status = 'scheduled', scheduled_at = strftime('%Y-%m-%dT%H:%M:%fZ','now'), scheduled_gpus = ? WHERE job_id = ?",
                        (json.dumps(indices), job.job_id),
                    )
                    db.commit()

                    # Dispatch execution (non-blocking)
                    asyncio.create_task(run_job(job.job_id, indices))
                    dispatched += 1

                    # Update available list
                    available = [i for i in available if i not in indices]

                    logger.info(
                        "Scheduled job %s → GPUs %s (score=%.1f, type=%s)",
                        job.job_id, indices, result.priority_score, result.chosen_gpu,
                    )

            # Periodic checkpoint
            checkpoint_counter += 1
            if checkpoint_counter % 60 == 0:
                model.checkpoint()

        except Exception:
            logger.exception("Scheduler error")

        await asyncio.sleep(0.5)


def feed_observation(job_id: str):
    """Called when a job completes — feeds runtime back to the scoring model."""
    try:
        with get_db() as db:
            row = db.execute(
                "SELECT job_type, runtime_seconds, scheduled_gpus FROM jobs WHERE job_id = ?",
                (job_id,),
            ).fetchone()
            if not row or not row["runtime_seconds"]:
                return

            # Get GPU type from the first scheduled GPU
            gpu_indices = json.loads(row["scheduled_gpus"]) if row["scheduled_gpus"] else []
            if not gpu_indices:
                return

            gpu_row = db.execute(
                "SELECT name FROM gpus WHERE gpu_index = ?", (gpu_indices[0],)
            ).fetchone()
            gpu_type = gpu_row["name"] if gpu_row else "generic"

            model = get_model()
            model.observe(row["job_type"], gpu_type, row["runtime_seconds"])
    except Exception as e:
        logger.warning("Failed to feed observation for job %s: %s", job_id, e)
