"""Admission control — per-user GPU concurrency limits."""

from __future__ import annotations

import logging

from radix_edge.config import get_config
from radix_edge.db import get_db

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


def check_admission(user_id: str, num_gpus: int, db_conn=None) -> tuple[bool, str]:
    """Check if a job submission is allowed under the user's quota.

    Returns (allowed, reason).
    """
    config = get_config()

    def _check(db):
        # Get user quota (or defaults)
        quota = db.execute(
            "SELECT * FROM user_quotas WHERE user_id = ?", (user_id,)
        ).fetchone()

        max_concurrent = quota["max_concurrent_gpu_jobs"] if quota else config.default_max_concurrent_gpu_jobs
        max_gpus = quota["max_gpus_total"] if quota else config.default_max_gpus_total

        # Count current active jobs for this user
        active = db.execute(
            "SELECT COUNT(*) as c FROM jobs WHERE user_id = ? AND status IN ('queued', 'scheduled', 'running')",
            (user_id,),
        ).fetchone()["c"]

        if active >= max_concurrent:
            return False, f"Concurrent job limit reached ({active}/{max_concurrent})"

        # Count current GPU usage for this user
        gpu_usage = db.execute(
            "SELECT COALESCE(SUM(num_gpus), 0) as total FROM jobs WHERE user_id = ? AND status IN ('scheduled', 'running')",
            (user_id,),
        ).fetchone()["total"]

        if gpu_usage + num_gpus > max_gpus:
            return False, f"GPU limit would be exceeded ({gpu_usage}+{num_gpus} > {max_gpus})"

        return True, "admitted"

    if db_conn is not None:
        return _check(db_conn)

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


def get_user_quota(user_id: str, db_conn=None) -> dict:
    """Get the effective quota for a user."""
    config = get_config()

    def _get(db):
        row = db.execute("SELECT * FROM user_quotas WHERE user_id = ?", (user_id,)).fetchone()
        if row:
            return dict(row)
        return {
            "user_id": user_id,
            "max_concurrent_gpu_jobs": config.default_max_concurrent_gpu_jobs,
            "max_gpus_total": config.default_max_gpus_total,
            "fair_share_weight": config.default_fair_share_weight,
        }

    if db_conn is not None:
        return _get(db_conn)

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


def set_user_quota(user_id: str, max_concurrent: int = None, max_gpus: int = None,
                   fair_share_weight: float = None, db_conn=None):
    """Create or update a user's quota."""
    config = get_config()

    def _set(db):
        existing = db.execute("SELECT * FROM user_quotas WHERE user_id = ?", (user_id,)).fetchone()
        if existing:
            if max_concurrent is not None:
                db.execute("UPDATE user_quotas SET max_concurrent_gpu_jobs = ? WHERE user_id = ?", (max_concurrent, user_id))
            if max_gpus is not None:
                db.execute("UPDATE user_quotas SET max_gpus_total = ? WHERE user_id = ?", (max_gpus, user_id))
            if fair_share_weight is not None:
                db.execute("UPDATE user_quotas SET fair_share_weight = ? WHERE user_id = ?", (fair_share_weight, user_id))
        else:
            db.execute(
                "INSERT INTO user_quotas (user_id, max_concurrent_gpu_jobs, max_gpus_total, fair_share_weight) VALUES (?, ?, ?, ?)",
                (
                    user_id,
                    max_concurrent if max_concurrent is not None else config.default_max_concurrent_gpu_jobs,
                    max_gpus if max_gpus is not None else config.default_max_gpus_total,
                    fair_share_weight if fair_share_weight is not None else config.default_fair_share_weight,
                ),
            )

    if db_conn is not None:
        _set(db_conn)
    else:
        with get_db() as db:
            _set(db)
