"""Scheduling policies: FIFO, Priority, FairShare, SJF.

Simplified from engine/scheduler/policies.py — uses local models.py types.
"""

from __future__ import annotations

import logging

from radix_edge.models import Job

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


def order_jobs(jobs: list[Job], policy: str, db_conn=None) -> list[Job]:
    """Order jobs according to the specified policy.

    Args:
        jobs: List of queued jobs to order.
        policy: One of 'fifo', 'priority', 'fairshare', 'sjf'.
        db_conn: Optional DB connection (for fairshare usage queries).

    Returns:
        Jobs ordered by scheduling priority (first = schedule first).
    """
    if not jobs:
        return []

    if policy == "fifo":
        return _order_fifo(jobs)
    elif policy == "priority":
        return _order_priority(jobs)
    elif policy == "fairshare":
        return _order_fairshare(jobs, db_conn)
    elif policy == "sjf":
        return _order_sjf(jobs)
    else:
        logger.warning("Unknown policy '%s', falling back to priority", policy)
        return _order_priority(jobs)


def _order_fifo(jobs: list[Job]) -> list[Job]:
    """First-In-First-Out: order by creation time."""
    return sorted(jobs, key=lambda j: j.created_at or "")


def _order_priority(jobs: list[Job]) -> list[Job]:
    """Priority: higher priority first, ties broken by creation time (FIFO)."""
    return sorted(jobs, key=lambda j: (-j.priority, j.created_at or ""))


def _order_fairshare(jobs: list[Job], db_conn=None) -> list[Job]:
    """Fair share: prioritize users who have consumed fewer GPU-seconds recently.

    Falls back to priority ordering when no DB connection is available.
    """
    if db_conn is None:
        return _order_priority(jobs)

    # Get recent GPU-seconds per user (last 24 hours)
    usage_rows = db_conn.execute(
        """SELECT user_id, SUM(runtime_seconds * num_gpus) as gpu_seconds
        FROM jobs
        WHERE status = 'succeeded' AND completed_at > datetime('now', '-24 hours')
        GROUP BY user_id"""
    ).fetchall()

    usage_map: dict[str, float] = {r["user_id"]: r["gpu_seconds"] or 0.0 for r in usage_rows}

    # Get fair share weights
    quota_rows = db_conn.execute("SELECT user_id, fair_share_weight FROM user_quotas").fetchall()
    weight_map: dict[str, float] = {r["user_id"]: r["fair_share_weight"] for r in quota_rows}

    def fairshare_key(job: Job):
        usage = usage_map.get(job.user_id, 0.0)
        weight = weight_map.get(job.user_id, 1.0)
        # Lower usage_ratio = higher priority
        usage_ratio = usage / max(weight, 0.01)
        return (usage_ratio, -job.priority, job.created_at or "")

    return sorted(jobs, key=fairshare_key)


def _order_sjf(jobs: list[Job]) -> list[Job]:
    """Shortest Job First: prefer jobs with lower timeout (proxy for duration).

    Since we don't have exact duration estimates, we use timeout_seconds as a proxy.
    Ties broken by priority then creation time.
    """
    return sorted(jobs, key=lambda j: (j.timeout_seconds, -j.priority, j.created_at or ""))
