"""Bedrock-Client. Zweiter Modellweg über AWS Bedrock (EU-Inferenzprofile, Frankfurt).

Teil des EU/DSGVO-Umbaus (Phase 1). Die Anfragen werden vollständig von AWS in
europäischen Regionen verarbeitet, nichts geht an Anthropic. In Phase 1 bedient
dieser Weg nur werkzeuglose Aufrufe (tools=None). Aufrufe mit WebSearch/WebFetch
laufen weiter über das Claude CLI, bis die staan-Recherche-Schleife (Phase 2)
fertig ist. Das Routing sitzt zentral in claude_client.call_claude.

Kosten. Bedrock meldet nur Token. Die Umrechnung in USD passiert hier über die
Preistabelle BEDROCK_PRICING aus config.py, damit ClaudeUsage.cost_usd gefüllt
bleibt und Statistik und Credits-Abrechnung unverändert funktionieren.
"""
import asyncio
import contextvars
import logging
import time

from config import (
    BEDROCK_CACHE_READ_FACTOR,
    BEDROCK_CACHE_WRITE_FACTOR,
    BEDROCK_MAX_TOKENS,
    BEDROCK_MODEL_MAP,
    BEDROCK_PRICING,
    BEDROCK_REGION,
    BEDROCK_RETRY_BUDGET_S,
    BEDROCK_RETRY_MIN_REST_S,
    BEDROCK_RETRY_WAITS,
    CLAUDE_MODEL_STANDARD,
    CLAUDE_TIMEOUT,
)
from agents.claude_client import (
    ClaudeCliError,
    ClaudeUsage,
    _cancel_event_var,
    _sanitize_mdash,
)

logger = logging.getLogger("osint.bedrock_client")

# Entspricht dem JSON-Zwang des CLI-Wegs bei tools=None (siehe claude_client.call_claude)
_JSON_ONLY_SYSTEM = (
    "CRITICAL: You are a JSON-only output agent. "
    "Output EXCLUSIVELY a single valid JSON object. "
    "No explanatory text, no markdown fences, no continuation of previous responses. "
    "Start your response with { and end with }."
)

# botocore-Fehlercodes, gruppiert nach den error_type-Kategorien des CLI-Wegs,
# damit die Retry-Steuerung des Orchestrators unverändert greift.
_RATE_LIMIT_CODES = {
    "ThrottlingException",
    "TooManyRequestsException",
    "ServiceUnavailableException",
    "ServiceQuotaExceededException",
    "ModelNotReadyException",
}
# Innerhalb der Kategorie 'rate_limit' zwei verschiedene Ursachen, die nach
# aussen gleich behandelt werden (damit die Retry-Steuerung des Orchestrators
# unveraendert greift), im Log aber auseinandergehalten gehoeren:
# Kapazitaet heisst, AWS hat gerade keine freie Rechenzeit fuer das Modell.
# Kontingent heisst, WIR haben zu viele Anfragen je Minute gestellt.
_KAPAZITAETS_CODES = {"ServiceUnavailableException", "ModelNotReadyException"}
_KONTINGENT_CODES = {
    "ThrottlingException",
    "TooManyRequestsException",
    "ServiceQuotaExceededException",
}
_AUTH_ERROR_CODES = {
    "AccessDeniedException",
    "UnrecognizedClientException",
    "ExpiredTokenException",
    "InvalidSignatureException",
}
# Transiente botocore-Netzfehler, die als TimeoutError hochgereicht werden
# (der Orchestrator behandelt TimeoutError/ConnectionError als retry-fähig).
_TRANSIENT_EXC_NAMES = (
    "ReadTimeoutError",
    "ConnectTimeoutError",
    "EndpointConnectionError",
    "ConnectionClosedError",
)

_client = None

# Klartext je Stoerungsart, genutzt in Log, Refresh-Protokoll und Oberflaeche.
_ART_TEXT = {
    "kapazitaet": "hat keine freie Kapazität",
    "kontingent": "meldet ausgeschöpftes Kontingent",
    "verbindung": "nicht erreichbar",
}

# --- Wartebudget und Sichtbarkeit je Refresh --------------------------------
# Alle drei Variablen gelten fuer den laufenden Refresh (ContextVar, also je
# Task getrennt). Ohne gesetzten Kontext verhaelt sich der Client wie bisher,
# nur mit Wiederholung: das Budget startet dann beim Standardwert und die
# Meldungen gehen ausschliesslich ins Log.
_wartebudget_var: contextvars.ContextVar[list] = contextvars.ContextVar("bedrock_wartebudget")
_stoerungen_var: contextvars.ContextVar[list] = contextvars.ContextVar("bedrock_stoerungen")
# Rueckmeldung an die Oberflaeche waehrend einer Wartezeit. Der Orchestrator
# setzt hier eine Funktion (art, sekunden) -> None, damit die Pipeline-Anzeige
# nicht wie ein Haenger aussieht.
_wartemelder_var: contextvars.ContextVar = contextvars.ContextVar("bedrock_wartemelder")


def refresh_kontext_starten() -> None:
    """Setzt Wartebudget und Stoerungsliste fuer einen neuen Refresh zurueck."""
    _wartebudget_var.set([BEDROCK_RETRY_BUDGET_S])
    _stoerungen_var.set([])


def stoerungen() -> list[dict]:
    """Stoerungen des laufenden Refreshs, aelteste zuerst."""
    return list(_stoerungen_var.get([]) or [])


def stoerungen_zusammenfassen() -> str:
    """Einzeiler ueber die Stoerungen des Laufs, leer wenn es keine gab.

    Landet im Refresh-Protokoll, damit ein Lauf mit Aussetzern nicht wie ein
    vollstaendiger aussieht.
    """
    eintraege = stoerungen()
    if not eintraege:
        return ""
    je_art: dict[str, int] = {}
    gewartet = 0.0
    for e in eintraege:
        je_art[e["art"]] = je_art.get(e["art"], 0) + 1
        gewartet += e["wartezeit"]
    teile = [f"{anzahl}x {_ART_TEXT.get(art, art)}" for art, anzahl in sorted(je_art.items())]
    return f"KI-Dienst: {', '.join(teile)}, insgesamt {gewartet:.0f}s gewartet"


def wartemelder_setzen(melder) -> None:
    """Hinterlegt die Rueckmeldung an die Oberflaeche fuer diesen Refresh."""
    _wartemelder_var.set(melder)


def _budget_rest() -> float:
    behaelter = _wartebudget_var.get(None)
    if behaelter is None:
        behaelter = [BEDROCK_RETRY_BUDGET_S]
        _wartebudget_var.set(behaelter)
    return behaelter[0]


def _budget_verbrauchen(sekunden: float) -> None:
    behaelter = _wartebudget_var.get(None)
    if behaelter is not None:
        behaelter[0] = max(0.0, behaelter[0] - sekunden)


def merke_stoerung(art: str, modell: str, wartezeit: float) -> None:
    """Haelt eine wiederholte Stoerung fuer das Refresh-Protokoll fest."""
    liste = _stoerungen_var.get(None)
    if liste is None:
        liste = []
        _stoerungen_var.set(liste)
    liste.append({"art": art, "modell": modell, "wartezeit": wartezeit})


def _melde_wartezeit(art: str, sekunden: float) -> None:
    """Meldet eine laufende Wartezeit an die Oberflaeche, wenn eingerichtet."""
    melder = _wartemelder_var.get(None)
    if melder is None:
        return
    try:
        melder(art, sekunden)
    except Exception as e:  # Anzeige darf den Lauf nie gefaehrden
        logger.debug("Wartemeldung fehlgeschlagen: %s", e)


def _get_client():
    """Erzeugt den bedrock-runtime-Client einmalig. Lazy Import, damit die App
    auch ohne installiertes boto3 startet, solange das Backend 'cli' ist."""
    global _client
    if _client is None:
        try:
            import boto3
            from botocore.config import Config
        except ImportError as e:
            raise ClaudeCliError(
                "cli_error",
                f"boto3 ist nicht installiert, Bedrock-Backend nicht nutzbar ({e})",
            )
        _client = boto3.client(
            "bedrock-runtime",
            region_name=BEDROCK_REGION,
            config=Config(
                connect_timeout=10,
                read_timeout=CLAUDE_TIMEOUT,
                # 'adaptive' bremst sich bei Drosselung selbst ein, statt
                # stur nachzuschieben. Die Versuchszahl bleibt klein, weil das
                # geduldige Warten in call_bedrock sitzt: botocore deckt die
                # Sekunden ab, unsere Schleife die Zehnersekunden.
                retries={"max_attempts": 3, "mode": "adaptive"},
            ),
        )
    return _client


def _fehlercode(exc: Exception) -> str:
    """Botocore-Fehlercode einer Exception, leer wenn keiner vorliegt."""
    response = getattr(exc, "response", None)
    if isinstance(response, dict):
        return response.get("Error", {}).get("Code", "") or ""
    return ""


def _classify_bedrock_error(exc: Exception) -> str:
    """Ordnet einer botocore-Exception eine error_type-Kategorie zu."""
    code = _fehlercode(exc)
    if code in _RATE_LIMIT_CODES:
        return "rate_limit"
    if code in _AUTH_ERROR_CODES:
        return "auth_error"
    return "cli_error"


def stoerungsart(exc: Exception) -> str | None:
    """Benennt die Ursache einer wiederholbaren Stoerung, sonst None.

    'kapazitaet' = AWS hat keine freie Rechenzeit fuer das Modell (HTTP 503).
    'kontingent' = unsere Anfragen je Minute sind ausgeschoepft (HTTP 429).
    """
    code = _fehlercode(exc)
    if code in _KAPAZITAETS_CODES:
        return "kapazitaet"
    if code in _KONTINGENT_CODES:
        return "kontingent"
    if type(exc).__name__ in _TRANSIENT_EXC_NAMES:
        return "verbindung"
    return None


def _cost_usd(
    cli_model: str,
    input_tokens: int,
    output_tokens: int,
    cache_write: int,
    cache_read: int,
) -> float:
    """Token mal Preis. Preise in BEDROCK_PRICING sind USD je 1 Mio Token."""
    prices = BEDROCK_PRICING.get(cli_model)
    if not prices:
        logger.warning(f"Keine Preise für Modell '{cli_model}' in BEDROCK_PRICING, cost_usd bleibt 0")
        return 0.0
    per_in = prices["input"] / 1_000_000
    per_out = prices["output"] / 1_000_000
    return (
        input_tokens * per_in
        + output_tokens * per_out
        + cache_write * per_in * BEDROCK_CACHE_WRITE_FACTOR
        + cache_read * per_in * BEDROCK_CACHE_READ_FACTOR
    )


async def call_bedrock(
    prompt: str,
    model: str | None = None,
    raw_text: bool = False,
    timeout: float | None = None,
) -> tuple[str, ClaudeUsage]:
    """Ruft Claude über AWS Bedrock (Converse API) auf. Gibt (result_text, usage) zurück.

    Gleicher Vertrag wie call_claude mit tools=None. Fehler kommen als
    ClaudeCliError mit denselben error_type-Kategorien, Timeouts als
    TimeoutError, Abbrüche als CancelledError.

    Bei Kapazitäts- und Kontingentfehlern wird nach BEDROCK_RETRY_WAITS erneut
    versucht, begrenzt durch das Zeitbudget dieses Aufrufs und das
    Wartebudget des laufenden Refreshs.
    """
    cli_model = model or CLAUDE_MODEL_STANDARD
    bedrock_model = BEDROCK_MODEL_MAP.get(cli_model)
    if not bedrock_model:
        raise ClaudeCliError(
            "cli_error",
            f"Kein EU-Inferenzprofil für Modell '{cli_model}' in BEDROCK_MODEL_MAP hinterlegt",
        )
    if not bedrock_model.startswith("eu."):
        # Sperre gegen Nicht-EU-Profile. Ein globales Profil würde die
        # Verarbeitung ausserhalb Europas erlauben, genau das soll dieser Weg
        # ausschliessen.
        raise ClaudeCliError(
            "cli_error",
            f"Bedrock-Profil '{bedrock_model}' ist kein EU-Profil (eu.-Präfix fehlt)",
        )

    effective_timeout = timeout if timeout is not None else CLAUDE_TIMEOUT
    cancel_event = _cancel_event_var.get(None)
    if cancel_event and cancel_event.is_set():
        raise asyncio.CancelledError("Cancel angefordert")

    client = _get_client()

    kwargs = {
        "modelId": bedrock_model,
        "messages": [{"role": "user", "content": [{"text": prompt}]}],
        "inferenceConfig": {"maxTokens": BEDROCK_MAX_TOKENS},
    }
    if not raw_text:
        kwargs["system"] = [{"text": _JSON_ONLY_SYSTEM}]

    beginn = time.monotonic()
    versuch = 0
    while True:
        try:
            response = await _converse_einmal(
                client, kwargs, effective_timeout - (time.monotonic() - beginn), cancel_event
            )
            break
        except (asyncio.CancelledError, TimeoutError):
            raise
        except Exception as e:
            art = stoerungsart(e)
            wartezeit = _wartezeit(
                art, versuch, rest=effective_timeout - (time.monotonic() - beginn)
            )
            if wartezeit is None:
                _protokolliere_fehler(e, cli_model, versuch)
                if type(e).__name__ in _TRANSIENT_EXC_NAMES:
                    raise TimeoutError(f"Bedrock Verbindungsfehler ({type(e).__name__}). {e}") from e
                raise ClaudeCliError(
                    _classify_bedrock_error(e), f"{type(e).__name__}. {e}"
                ) from e
            versuch += 1
            _budget_verbrauchen(wartezeit)
            merke_stoerung(art, cli_model, wartezeit)
            logger.warning(
                "Bedrock %s (%s, %s): Versuch %d in %.0fs",
                _ART_TEXT.get(art, art), type(e).__name__, cli_model, versuch + 1, wartezeit,
            )
            _melde_wartezeit(art, wartezeit)
            await asyncio.sleep(wartezeit)
            if cancel_event and cancel_event.is_set():
                raise asyncio.CancelledError("Cancel angefordert")

    return _auswerten(response, cli_model, bedrock_model)


def _wartezeit(art: str | None, versuch: int, rest: float) -> float | None:
    """Wartezeit vor dem nächsten Versuch, None wenn nicht wiederholt wird.

    Drei Bedingungen müssen zusammenkommen: die Störung ist wiederholbar, es
    gibt noch eine Wartestufe, und die Wartezeit passt sowohl in das
    Zeitbudget dieses Aufrufs als auch in das Wartebudget des Refreshs.
    """
    if not art or versuch >= len(BEDROCK_RETRY_WAITS):
        return None
    wartezeit = BEDROCK_RETRY_WAITS[versuch]
    if rest < wartezeit + BEDROCK_RETRY_MIN_REST_S:
        logger.warning(
            "Bedrock: keine Wiederholung, Zeitbudget des Aufrufs reicht nicht "
            "(%.0fs übrig, %.0fs Wartezeit plus %.0fs Reserve nötig)",
            max(rest, 0), wartezeit, BEDROCK_RETRY_MIN_REST_S,
        )
        return None
    if _budget_rest() < wartezeit:
        logger.warning(
            "Bedrock: keine Wiederholung, Wartebudget des Laufs erschöpft (%.0fs übrig)",
            _budget_rest(),
        )
        return None
    return wartezeit


def _protokolliere_fehler(e: Exception, cli_model: str, versuch: int) -> None:
    """Schreibt den endgültigen Fehler mit seiner Ursache ins Log."""
    exc_name = type(e).__name__
    zusatz = f" nach {versuch + 1} Versuchen" if versuch else ""
    art = stoerungsart(e)
    if art:
        logger.warning(
            "Bedrock %s (%s, %s)%s. %s", _ART_TEXT.get(art, art), exc_name, cli_model, zusatz, e
        )
    elif _classify_bedrock_error(e) == "auth_error":
        logger.error("Bedrock Auth-Fehler (%s). %s", exc_name, e)
    else:
        logger.error("Bedrock-Fehler (%s). %s", exc_name, e)


async def _converse_einmal(client, kwargs: dict, rest_timeout: float, cancel_event):
    """Führt genau einen Converse-Aufruf aus und gibt die Antwort zurück."""
    if rest_timeout <= 0:
        raise TimeoutError("Bedrock Timeout: Zeitbudget vor dem Aufruf aufgebraucht")

    call_task = asyncio.create_task(asyncio.to_thread(client.converse, **kwargs))
    wait_tasks = [call_task]
    cancel_wait_task = None
    if cancel_event:
        cancel_wait_task = asyncio.create_task(cancel_event.wait())
        wait_tasks.append(cancel_wait_task)

    done, pending = await asyncio.wait(
        wait_tasks, timeout=rest_timeout, return_when=asyncio.FIRST_COMPLETED
    )
    for p in pending:
        p.cancel()

    if call_task not in done:
        # Der HTTP-Aufruf im Thread lässt sich nicht hart abbrechen, er läuft
        # aus und sein Ergebnis wird verworfen. Der Callback verhindert
        # "exception was never retrieved"-Warnungen.
        call_task.add_done_callback(
            lambda t: t.exception() if not t.cancelled() else None
        )
        if cancel_wait_task is not None and cancel_wait_task in done:
            raise asyncio.CancelledError("Cancel angefordert")
        raise TimeoutError(f"Bedrock Timeout nach {rest_timeout:.0f}s")

    # Fehler werden hier NICHT behandelt, das entscheidet die Wiederholung in
    # call_bedrock anhand der Stoerungsart.
    return call_task.result()


def _auswerten(response: dict, cli_model: str, bedrock_model: str) -> tuple[str, ClaudeUsage]:
    """Wandelt die Converse-Antwort in (Text, Verbrauch) um."""
    out_msg = response.get("output", {}).get("message", {})
    parts = [c.get("text", "") for c in out_msg.get("content", []) if "text" in c]
    result_text = "".join(parts).strip()

    stop_reason = response.get("stopReason", "")
    if stop_reason == "max_tokens":
        logger.warning(
            f"Bedrock-Antwort bei maxTokens={BEDROCK_MAX_TOKENS} abgeschnitten (Modell {cli_model})"
        )

    u = response.get("usage", {})
    input_tokens = int(u.get("inputTokens", 0) or 0)
    output_tokens = int(u.get("outputTokens", 0) or 0)
    cache_write = int(u.get("cacheWriteInputTokens", 0) or 0)
    cache_read = int(u.get("cacheReadInputTokens", 0) or 0)
    usage = ClaudeUsage(
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        cache_creation_tokens=cache_write,
        cache_read_tokens=cache_read,
        cost_usd=_cost_usd(cli_model, input_tokens, output_tokens, cache_write, cache_read),
        duration_ms=int(response.get("metrics", {}).get("latencyMs", 0) or 0),
    )
    logger.info(
        f"Bedrock [{cli_model} -> {bedrock_model}]: {usage.input_tokens} in / "
        f"{usage.output_tokens} out / cache {usage.cache_creation_tokens}+{usage.cache_read_tokens} / "
        f"${usage.cost_usd:.4f} / {usage.duration_ms}ms"
    )
    return _sanitize_mdash(result_text), usage
