fair_use_admin.py

"""Admin endpoints for fair-use management."""

import hashlib
import hmac
import logging
import os
from datetime import datetime
from typing import Any, Dict, Optional

from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request
from pydantic import BaseModel, ConfigDict, Field
from models.shared import StatusResponse

import database.fair_use as fair_use_db
from database._client import db
from utils.other.endpoints import get_current_user_uid, rate_limit_dependency
from utils.fair_use import (
    get_rolling_speech_ms,
    get_dg_budget_status,
    invalidate_enforcement_cache,
    normalize_expired_restriction_state,
    FAIR_USE_ENABLED,
    FAIR_USE_DAILY_SPEECH_MS,
    FAIR_USE_3DAY_SPEECH_MS,
    FAIR_USE_WEEKLY_SPEECH_MS,
)

logger = logging.getLogger(__name__)

router = APIRouter()

ADMIN_KEY = os.getenv('ADMIN_KEY', '')


class FairUseLimitsResponse(BaseModel):
    daily_hours: float
    three_day_hours: float
    weekly_hours: float


class FairUseUsagePctResponse(BaseModel):
    daily: float
    three_day: float
    weekly: float


class FairUseDailyGenerationsBudgetResponse(BaseModel):
    daily_limit_ms: int
    used_ms: int
    remaining_ms: int
    exhausted: bool
    resets_at: Optional[str] = None


class FairUseStatusResponse(BaseModel):
    stage: str
    case_ref: str
    speech_hours_today: float
    speech_hours_3day: float
    speech_hours_weekly: float
    limits: FairUseLimitsResponse
    usage_pct: FairUseUsagePctResponse
    dg_budget: FairUseDailyGenerationsBudgetResponse
    message: str


class PublicFairUseCaseStatusResponse(BaseModel):
    case_ref: str
    stage: str
    message: str
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None
    support_email: str


class FlaggedUsersResponse(BaseModel):
    """Admin dashboard: users with active fair-use enforcement."""

    users: list[Dict[str, Any]] = Field(description='Users with active enforcement, each a fair-use state dict.')
    fair_use_enabled: bool = Field(description='Whether fair-use enforcement is globally enabled.')


class FairUseUserDetailResponse(BaseModel):
    """Admin per-user fair-use detail."""

    uid: str = Field(description='User UID.')
    state: Dict[str, Any] = Field(description='Current fair-use state document.')
    events: list[Dict[str, Any]] = Field(description='Recent fair-use events.')
    current_speech_ms: Dict[str, int] = Field(
        description='Rolling speech usage in milliseconds: daily_ms, three_day_ms, weekly_ms.'
    )


class FairUseSetStageResponse(BaseModel):
    """Ack for manual stage update; preserves the stage value."""

    status: str = Field(description='Ack status, e.g. "updated".')
    stage: str = Field(description='The enforcement stage that was set (none|warning|throttle|restrict).')


class FairUseCaseLookupResponse(BaseModel):
    """A fair-use event located by case reference (support team lookup).

    The Firestore event document carries dynamic fields beyond uid/event_id;
    extra='allow' passes them through without modelling every key.
    """

    model_config = ConfigDict(extra='allow')
    uid: str = Field(description='User UID who owns the event.')
    event_id: str = Field(description='Event identifier.')


def _verify_admin_key(x_admin_key: str = Header(..., alias='X-Admin-Key')) -> str:
    """Validate admin key from request header using constant-time comparison.

    Returns a short hash of the key for audit logging (not the key itself).
    """
    if not ADMIN_KEY or not hmac.compare_digest(x_admin_key, ADMIN_KEY):
        raise HTTPException(status_code=403, detail='Invalid admin key')
    return f'admin:{hashlib.sha256(x_admin_key.encode()).hexdigest()[:8]}'


# ---------------------------------------------------------------------------
# Dashboard
# ---------------------------------------------------------------------------


@router.get('/v1/admin/fair-use/flagged', tags=['admin'], response_model=FlaggedUsersResponse)
def get_flagged_users(
    admin_id: str = Depends(_verify_admin_key),
    stage: Optional[str] = None,
    limit: int = Query(default=50, le=200),
):
    """Get users with active fair-use enforcement."""
    # Clamp in-function (not only via Query) so direct/non-HTTP callers can't pass a
    # negative or huge limit straight through to the Firestore query.
    limit = max(1, min(limit, 200))
    users = fair_use_db.get_flagged_users(stage_filter=stage, limit=limit)
    return {'users': users, 'fair_use_enabled': FAIR_USE_ENABLED}


@router.get('/v1/admin/fair-use/user/{uid}', tags=['admin'], response_model=FairUseUserDetailResponse)
def get_user_fair_use_detail(uid: str, admin_id: str = Depends(_verify_admin_key)):
    """Get detailed fair-use state and events for a specific user."""
    state = fair_use_db.get_fair_use_state(uid)
    events = fair_use_db.get_fair_use_events(uid, limit=50)
    speech = get_rolling_speech_ms(uid)

    return {
        'uid': uid,
        'state': state,
        'events': events,
        'current_speech_ms': speech,
    }


# ---------------------------------------------------------------------------
# Admin actions
# ---------------------------------------------------------------------------


@router.post('/v1/admin/fair-use/user/{uid}/resolve-event/{event_id}', tags=['admin'], response_model=StatusResponse)
def resolve_event(uid: str, event_id: str, admin_id: str = Depends(_verify_admin_key), notes: str = Query(default='')):
    """Mark a fair-use event as resolved."""
    fair_use_db.resolve_fair_use_event(uid, event_id, admin_uid=admin_id, notes=notes)
    return {'status': 'resolved'}


@router.post('/v1/admin/fair-use/user/{uid}/reset', tags=['admin'], response_model=StatusResponse)
def reset_user_fair_use(uid: str, admin_id: str = Depends(_verify_admin_key)):
    """Reset a user's fair-use state to clean."""
    fair_use_db.reset_fair_use_state(uid, admin_uid=admin_id)
    invalidate_enforcement_cache(uid)
    return {'status': 'reset'}


@router.post('/v1/admin/fair-use/user/{uid}/set-stage', tags=['admin'], response_model=FairUseSetStageResponse)
def set_user_stage(uid: str, stage: str = Query(...), admin_id: str = Depends(_verify_admin_key)):
    """Manually set a user's enforcement stage."""
    valid_stages = {'none', 'warning', 'throttle', 'restrict'}
    if stage not in valid_stages:
        raise HTTPException(status_code=400, detail=f'Invalid stage. Must be one of: {valid_stages}')

    updates = {'stage': stage}
    if stage == 'none':
        updates['throttle_until'] = None
        updates['restrict_until'] = None

    fair_use_db.update_fair_use_state(uid, updates)
    invalidate_enforcement_cache(uid)
    return {'status': 'updated', 'stage': stage}


@router.get('/v1/admin/fair-use/case/{case_ref}', tags=['admin'], response_model=FairUseCaseLookupResponse)
def lookup_case(case_ref: str, admin_id: str = Depends(_verify_admin_key)):
    """Look up a fair-use event by case reference (for support team)."""
    # Search across all users' events for this case_ref
    query = db.collection_group('fair_use_events').where('case_ref', '==', case_ref).limit(1)
    for doc in query.stream():
        data = doc.to_dict()
        path_parts = doc.reference.path.split('/')
        if len(path_parts) >= 2:
            data['uid'] = path_parts[1]
        data['event_id'] = doc.id
        return data
    raise HTTPException(status_code=404, detail=f'Case {case_ref} not found')


SUPPORT_EMAIL = 'team@basedhardware.com'


# ---------------------------------------------------------------------------
# Public: unauthenticated case status lookup (for tracking page)
# ---------------------------------------------------------------------------


@router.get(
    '/v1/fair-use/case/{case_ref}/status',
    tags=['fair_use'],
    dependencies=[Depends(rate_limit_dependency('fair_use_case_status', requests_per_window=10, window_seconds=60))],
    response_model=PublicFairUseCaseStatusResponse,
)
def get_public_case_status(case_ref: str):
    """Public unauthenticated endpoint: look up case status by reference.

    Returns only non-sensitive info: stage, message, timestamps, support email.
    No usage data or user identity exposed.
    """
    query = db.collection_group('fair_use_events').where('case_ref', '==', case_ref).limit(1)
    for doc in query.stream():
        data = doc.to_dict()
        # Extract uid to get current enforcement stage
        path_parts = doc.reference.path.split('/')
        uid = path_parts[1] if len(path_parts) >= 2 else None

        stage = 'none'
        if uid:
            state = fair_use_db.get_fair_use_state(uid)
            stage = state.get('stage', 'none')

        created_at = data.get('created_at')
        updated_at = data.get('resolved_at') or created_at

        return {
            'case_ref': case_ref,
            'stage': stage,
            'message': _user_facing_message(stage, case_ref),
            'created_at': str(created_at) if created_at else None,
            'updated_at': str(updated_at) if updated_at else None,
            'support_email': SUPPORT_EMAIL,
        }
    raise HTTPException(status_code=404, detail='Case not found')


# ---------------------------------------------------------------------------
# Support: user-facing endpoint to see their own fair-use status
# ---------------------------------------------------------------------------


@router.get('/v1/fair-use/status', tags=['fair_use'], response_model=FairUseStatusResponse)
def get_my_fair_use_status(uid: str = Depends(get_current_user_uid)):
    """User-facing endpoint: see your own fair-use status and speech usage."""
    state = normalize_expired_restriction_state(uid, fair_use_db.get_fair_use_state(uid))
    speech = get_rolling_speech_ms(uid)

    stage = state.get('stage', 'none')
    case_ref = state.get('last_case_ref', '')

    daily_ms = speech.get('daily_ms', 0)
    three_day_ms = speech.get('three_day_ms', 0)
    weekly_ms = speech.get('weekly_ms', 0)

    # DG budget (only meaningful for restrict stage, but always returned for frontend simplicity)
    dg_budget = get_dg_budget_status(uid)

    return {
        'stage': stage,
        'case_ref': case_ref,
        'speech_hours_today': round(daily_ms / 3600000, 2),
        'speech_hours_3day': round(three_day_ms / 3600000, 2),
        'speech_hours_weekly': round(weekly_ms / 3600000, 2),
        'limits': {
            'daily_hours': round(FAIR_USE_DAILY_SPEECH_MS / 3600000, 2),
            'three_day_hours': round(FAIR_USE_3DAY_SPEECH_MS / 3600000, 2),
            'weekly_hours': round(FAIR_USE_WEEKLY_SPEECH_MS / 3600000, 2),
        },
        'usage_pct': {
            'daily': round(daily_ms / FAIR_USE_DAILY_SPEECH_MS * 100, 1) if FAIR_USE_DAILY_SPEECH_MS else 0,
            'three_day': round(three_day_ms / FAIR_USE_3DAY_SPEECH_MS * 100, 1) if FAIR_USE_3DAY_SPEECH_MS else 0,
            'weekly': round(weekly_ms / FAIR_USE_WEEKLY_SPEECH_MS * 100, 1) if FAIR_USE_WEEKLY_SPEECH_MS else 0,
        },
        'dg_budget': dg_budget,
        'message': _user_facing_message(stage, case_ref),
    }


def _user_facing_message(stage: str, case_ref: str = '') -> str:
    ref_note = f' Your case reference is {case_ref}.' if case_ref else ''
    messages = {
        'none': 'Your usage is within normal limits.',
        'warning': (
            'Your usage is higher than typical. Omi is designed for personal conversations. '
            f'If non-personal content transcription continues, your service may be adjusted.{ref_note}'
        ),
        'throttle': (
            'Your transcription quality has been temporarily reduced due to high non-personal usage. '
            'This will reset automatically. Contact support at team@basedhardware.com if you believe this is an error. '
            f'Please quote your case reference when contacting support.{ref_note}'
        ),
        'restrict': (
            'Your cloud transcription is temporarily limited. On-device transcription continues normally. '
            'Contact support at team@basedhardware.com to discuss your usage and resolve this. '
            f'Please quote your case reference when contacting support.{ref_note}'
        ),
    }
    return messages.get(stage, messages['none'])