"""OpenAI-specific capacity adaptation. Generic capacity code must not parse these headers."""

from __future__ import annotations

import json
import re
import time
from collections.abc import Mapping
from typing import Any

from src.core.provider_capacity.types import CapacityObservation

# Conservative bound for reservation, not a billed tokenizer. No tiktoken in this repo.
_CHARS_PER_TOKEN = 3
_MESSAGE_OVERHEAD_TOKENS = 8
_RESET_PATTERN = re.compile(
    r"^(?:(\d+)h)?(?:(\d+)m)?(?:(\d+(?:\.\d+)?)s)?$",
    re.IGNORECASE,
)


def estimate_openai_request_tokens(
    messages: list[dict[str, Any]],
    *,
    schema: Any = None,
    max_completion_tokens: int = 0,
) -> int:
    chunks: list[str] = []
    for message in messages:
        content = message.get("content")
        if isinstance(content, str):
            chunks.append(content)
        elif content is not None:
            chunks.append(json.dumps(content, ensure_ascii=False))
    if schema is not None and hasattr(schema, "model_json_schema"):
        chunks.append(json.dumps(schema.model_json_schema(), ensure_ascii=False))
    prompt_tokens = _chars_to_tokens("".join(chunks))
    prompt_tokens += _MESSAGE_OVERHEAD_TOKENS * max(len(messages), 1)
    return prompt_tokens + max(int(max_completion_tokens), 0)


def observations_from_openai_headers(
    headers: Mapping[str, Any] | None,
) -> list[CapacityObservation]:
    if not headers:
        return []
    normalized = {str(key).lower(): str(value) for key, value in headers.items()}
    now_ms = int(time.time() * 1000)
    observations: list[CapacityObservation] = []
    token_obs = _dimension_observation(
        normalized,
        dimension="tokens",
        limit_key="x-ratelimit-limit-tokens",
        remaining_key="x-ratelimit-remaining-tokens",
        reset_key="x-ratelimit-reset-tokens",
        now_ms=now_ms,
    )
    if token_obs is not None:
        observations.append(token_obs)
    request_obs = _dimension_observation(
        normalized,
        dimension="requests",
        limit_key="x-ratelimit-limit-requests",
        remaining_key="x-ratelimit-remaining-requests",
        reset_key="x-ratelimit-reset-requests",
        now_ms=now_ms,
    )
    if request_obs is not None:
        observations.append(request_obs)
    return observations


def parse_openai_reset_duration(value: str | None) -> float | None:
    if value is None or value == "":
        return None
    text = str(value).strip()
    try:
        return float(text)
    except ValueError:
        pass
    match = _RESET_PATTERN.fullmatch(text)
    if not match or not any(match.groups()):
        return None
    hours = int(match.group(1) or 0)
    minutes = int(match.group(2) or 0)
    seconds = float(match.group(3) or 0)
    return hours * 3600 + minutes * 60 + seconds


def _dimension_observation(
    headers: Mapping[str, str],
    *,
    dimension: str,
    limit_key: str,
    remaining_key: str,
    reset_key: str,
    now_ms: int,
) -> CapacityObservation | None:
    limit = _optional_int(headers.get(limit_key))
    remaining = _optional_int(headers.get(remaining_key))
    reset = parse_openai_reset_duration(headers.get(reset_key))
    if limit is None and remaining is None and reset is None:
        return None
    return CapacityObservation(
        dimension=dimension,
        limit=limit,
        remaining=remaining,
        reset_in_seconds=reset,
        observed_at_ms=now_ms,
        source="openai_headers",
    )


def _chars_to_tokens(text: str) -> int:
    return max(1, (len(text.encode("utf-8")) + _CHARS_PER_TOKEN - 1) // _CHARS_PER_TOKEN)


def _optional_int(value: str | None) -> int | None:
    if value is None or value == "":
        return None
    try:
        return int(float(value))
    except ValueError:
        return None