"""Fair-share accounting — tracks GPU-seconds per user."""

from __future__ import annotations

import logging

from radix_edge.db import get_db

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


def get_user_usage(user_id: str, hours: int = 24, db_conn=None) -> dict:
    """Get a user's resource usage in the recent window.

    Returns dict with gpu_seconds, job_count, avg_runtime.
    """
    def _get(db):
        row = db.execute(
            """SELECT
                COALESCE(SUM(runtime_seconds * num_gpus), 0.0) as gpu_seconds,
                COUNT(*) as job_count,
                COALESCE(AVG(runtime_seconds), 0.0) as avg_runtime
            FROM jobs
            WHERE user_id = ? AND status IN ('succeeded', 'failed')
            AND completed_at > datetime('now', ? || ' hours')""",
            (user_id, f"-{hours}"),
        ).fetchone()
        return {
            "user_id": user_id,
            "gpu_seconds": row["gpu_seconds"],
            "job_count": row["job_count"],
            "avg_runtime": row["avg_runtime"],
            "window_hours": hours,
        }

    if db_conn is not None:
        return _get(db_conn)

    with get_db() as db:
        return _get(db)


def get_all_user_usage(hours: int = 24, db_conn=None) -> list[dict]:
    """Get resource usage for all users in the recent window."""
    def _get(db):
        rows = db.execute(
            """SELECT user_id,
                COALESCE(SUM(runtime_seconds * num_gpus), 0.0) as gpu_seconds,
                COUNT(*) as job_count
            FROM jobs
            WHERE status IN ('succeeded', 'failed')
            AND completed_at > datetime('now', ? || ' hours')
            GROUP BY user_id
            ORDER BY gpu_seconds DESC""",
            (f"-{hours}",),
        ).fetchall()
        return [{"user_id": r["user_id"], "gpu_seconds": r["gpu_seconds"], "job_count": r["job_count"]} for r in rows]

    if db_conn is not None:
        return _get(db_conn)

    with get_db() as db:
        return _get(db)


def get_user_priority(user_id: str, db_conn=None) -> float:
    """Get scheduling priority multiplier for a user (inverse of usage ratio).

    Higher = should be prioritized more. Users with less usage get higher priority.
    """
    def _calc(db):
        # Get this user's usage
        usage = get_user_usage(user_id, hours=24, db_conn=db)

        # Get fair share weight
        quota = db.execute(
            "SELECT fair_share_weight FROM user_quotas WHERE user_id = ?", (user_id,)
        ).fetchone()
        weight = quota["fair_share_weight"] if quota else 1.0

        if usage["gpu_seconds"] <= 0:
            return 2.0  # High priority for unused users

        # Compare to average usage
        all_usage = get_all_user_usage(hours=24, db_conn=db)
        if not all_usage:
            return 1.0

        avg_gpu_seconds = sum(u["gpu_seconds"] for u in all_usage) / len(all_usage)
        if avg_gpu_seconds <= 0:
            return 1.0

        usage_ratio = usage["gpu_seconds"] / (avg_gpu_seconds * weight)
        return 1.0 / max(usage_ratio, 0.1)

    if db_conn is not None:
        return _calc(db_conn)

    with get_db() as db:
        return _calc(db)
