helper/misc.py

import time
import asyncio

from helper.subscriptions import USAGE_PERIODS, usage_locks, usage_store, TIER_CONFIG, client_subject_bindings, CLIENT_BIND_TTL_SECONDS, MAX_CLIENT_ID_LENGTH
from fastapi import Request, HTTPException
from typing import Optional, Dict, Any, List
import re
import hashlib

def extract_user_text(messages: list) -> str:
    return " ".join(
        message_content_to_text(m.get("content"))
        for m in messages
        if isinstance(m, dict) and m.get("role") == "user"
    ).lower()

def get_usage_period_key(metric: str) -> str:
    now = time.gmtime()
    period = USAGE_PERIODS.get(metric, "daily")
    if period == "weekly":
        iso_year, iso_week, _ = time.strftime("%G %V %u", now).split(" ")
        return f"{iso_year}-W{iso_week}"
    return time.strftime("%Y-%m-%d", now)


def sanitize_client_id(raw_client_id: Optional[str]) -> Optional[str]:
    if not isinstance(raw_client_id, str):
        return None
    trimmed = raw_client_id.strip()
    if not trimmed or len(trimmed) > MAX_CLIENT_ID_LENGTH:
        return None
    if not re.match(r"^[A-Za-z0-9._:-]+$", trimmed):
        return None
    return trimmed


def get_usage_lock(metric: str, subject: str) -> asyncio.Lock:
    metric_locks = usage_locks.get(metric)
    if metric_locks is None:
        metric_locks = {}
        usage_locks[metric] = metric_locks
    lock = metric_locks.get(subject)
    if lock is None:
        lock = asyncio.Lock()
        metric_locks[subject] = lock
    return lock


def build_default_subject(request: Request, client_id: Optional[str]) -> str:
    if client_id:
        client_hash = hashlib.sha256(client_id.encode("utf-8")).hexdigest()[:24]
        return f"client:{client_hash}"
    host = request.client.host if request.client else "unknown"
    user_agent = request.headers.get("user-agent", "")
    ua_hash = (
        hashlib.sha256(user_agent.encode("utf-8")).hexdigest()[:12]
        if user_agent
        else "noua"
    )
    return f"anon:{host}:{ua_hash}"


def bind_client_subject(client_id: Optional[str], subject: str, plan_key: str):
    if not client_id:
        return
    client_subject_bindings[client_id] = {
        "subject": subject,
        "plan_key": plan_key,
        "expires_at": time.time() + CLIENT_BIND_TTL_SECONDS,
    }


def resolve_bound_subject(client_id: Optional[str], fallback_subject: str) -> str:
    if not client_id:
        return fallback_subject
    bound = client_subject_bindings.get(client_id)
    if not bound:
        return fallback_subject
    if bound.get("expires_at", 0) <= time.time():
        client_subject_bindings.pop(client_id, None)
        return fallback_subject
    return bound.get("subject", fallback_subject)


def normalize_prompt_value(prompt: Optional[str], field_name: str = "prompt") -> str:
    if not isinstance(prompt, str):
        raise HTTPException(status_code=400, detail=f"{field_name} is required")
    normalized = prompt.strip()
    if not normalized:
        raise HTTPException(status_code=400, detail=f"{field_name} is required")
    return normalized


def enforce_prompt_size(prompt: str, max_chars: int, max_bytes: int, context: str):
    char_len = len(prompt)
    byte_len = len(prompt.encode("utf-8"))
    if char_len > max_chars or byte_len > max_bytes:
        raise HTTPException(
            status_code=413,
            detail=(
                f"{context} is too large ({char_len} chars, {byte_len} bytes). "
                f"Max allowed is {max_chars} chars or {max_bytes} bytes."
            ),
        )


def message_content_to_text(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: List[str] = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
                continue
            if isinstance(item, dict):
                text = item.get("text")
                if isinstance(text, str):
                    parts.append(text)
        return " ".join(parts)
    return ""


def calculate_messages_size(messages: list) -> tuple[int, int]:
    total_chars = 0
    total_bytes = 0
    for message in messages:
        if not isinstance(message, dict):
            continue
        text = message_content_to_text(message.get("content"))
        if not text:
            continue
        total_chars += len(text)
        total_bytes += len(text.encode("utf-8"))
    return total_chars, total_bytes


def get_usage_snapshot_for_subject(plan_key: str, subject: str) -> Dict[str, Dict[str, Any]]:
    plan = TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"]
    plan_limits = plan.get("limits", {})
    snapshot: Dict[str, Dict[str, Any]] = {}
    for metric in usage_store.keys():
        limit = plan_limits.get(metric)
        window_key = get_usage_period_key(metric)
        entry = usage_store[metric].get(subject)
        used = 0
        if entry and entry.get("window") == window_key:
            used = max(0, int(entry.get("count", 0)))
        remaining = None if limit is None else max(0, int(limit) - used)
        snapshot[metric] = {
            "limit": limit,
            "used": used,
            "remaining": remaining,
            "window": window_key,
            "period": USAGE_PERIODS.get(metric, "daily"),
        }
    return snapshot