"""Shared helpers for API service handlers."""

from __future__ import annotations

import logging

from fastapi import HTTPException

from aistudio_api.api.response_models import (
    HealthResponse,
    ModelStatsResponse,
    StatsResponse,
    StatsTotalsResponse,
)
from aistudio_api.api.state import runtime_state

logger = logging.getLogger("aistudio.server")
MAX_RETRIES = 3


async def try_switch_account() -> bool:
    """尝试切换到下一个可用账号。返回是否成功切换。"""
    rotator = runtime_state.rotator
    if rotator is None:
        return False

    next_account = await rotator.get_next_account()
    if next_account is None:
        return False

    account_service = runtime_state.account_service
    client = runtime_state.client
    if account_service is None or client is None or client._session is None:
        return False

    result = await account_service.activate_account(
        next_account.id,
        client._session,
        runtime_state.snapshot_cache,
        None,  # skip lock — caller already holds it
        keep_snapshot_cache=False,
    )
    return result is not None


def require_busy_lock():
    busy_lock = runtime_state.busy_lock
    if busy_lock is None:
        raise HTTPException(
            503, detail={"message": "Server not ready", "type": "service_unavailable"}
        )
    if busy_lock.locked():
        raise HTTPException(
            503,
            detail={
                "message": "Server is busy with maximum concurrent requests",
                "type": "service_unavailable",
            },
            headers={"Retry-After": "2"},
        )
    return busy_lock


async def ensure_active_account(attempt: int) -> None:
    if attempt != 0:
        return
    account_svc = runtime_state.account_service
    if account_svc and not account_svc.get_active_account():
        await try_switch_account()


def record_rotator_event(event: str) -> None:
    rotator = runtime_state.rotator
    account_service = runtime_state.account_service
    account = account_service.get_active_account() if account_service else None
    if not rotator or account is None:
        return
    if event == "success":
        rotator.record_success(account.id)
    elif event == "rate_limited":
        rotator.record_rate_limited(account.id)
    elif event == "error":
        rotator.record_error(account.id)


def health_response() -> HealthResponse:
    busy_lock = runtime_state.busy_lock
    return HealthResponse(status="ok", busy=busy_lock.locked() if busy_lock else False)


def stats_response() -> StatsResponse:
    stats = dict(runtime_state.model_stats)
    totals = StatsTotalsResponse(
        requests=sum(s.requests for s in stats.values()),
        success=sum(s.success for s in stats.values()),
        rate_limited=sum(s.rate_limited for s in stats.values()),
        errors=sum(s.errors for s in stats.values()),
        prompt_tokens=sum(s.prompt_tokens for s in stats.values()),
        completion_tokens=sum(s.completion_tokens for s in stats.values()),
        total_tokens=sum(s.total_tokens for s in stats.values()),
    )
    models = {
        name: ModelStatsResponse(
            requests=s.requests,
            success=s.success,
            rate_limited=s.rate_limited,
            errors=s.errors,
            prompt_tokens=s.prompt_tokens,
            completion_tokens=s.completion_tokens,
            total_tokens=s.total_tokens,
            last_used=s.last_used,
        )
        for name, s in stats.items()
    }
    return StatsResponse(models=models, totals=totals)