"""
AWS Bedrock integration — streaming and non-streaming.
Uses Bearer-token auth and parses the AWS binary event-stream wire format.
"""

import struct
import json
import base64
from typing import AsyncIterator, Optional
from urllib.parse import quote

import httpx

from backend import config

_client: Optional[httpx.AsyncClient] = None


def _get_client() -> httpx.AsyncClient:
    global _client
    if _client is None or _client.is_closed:
        _client = httpx.AsyncClient(
            timeout=httpx.Timeout(connect=30.0, read=600.0, write=30.0, pool=60.0),
        )
    return _client


async def close_http_client():
    global _client
    if _client and not _client.is_closed:
        await _client.aclose()
        _client = None


# ══════════════════════════════════════════════════════════════
# AWS Event-Stream Binary Parser
# ══════════════════════════════════════════════════════════════
#
# Message layout (big-endian):
#   [total_length : 4B]
#   [headers_length : 4B]
#   [prelude_crc : 4B]
#   [headers : headers_length B]
#   [payload : total_length - headers_length - 16 B]
#   [message_crc : 4B]
# ══════════════════════════════════════════════════════════════

_MIN_MSG = 16          # 4+4+4 (prelude) + 0 (hdrs) + 0 (payload) + 4 (crc)
_PRELUDE = 12          # total_length + headers_length + prelude_crc


def _drain_messages(buf: bytes) -> tuple[list[dict], bytes]:
    """Return (parsed_events, remaining_buffer)."""
    events: list[dict] = []
    offset = 0

    while len(buf) - offset >= _MIN_MSG:
        total_length = struct.unpack("!I", buf[offset : offset + 4])[0]
        if total_length < _MIN_MSG or len(buf) - offset < total_length:
            break                               # incomplete message

        hdrs_len = struct.unpack("!I", buf[offset + 4 : offset + 8])[0]
        h_start = offset + _PRELUDE
        h_end = h_start + hdrs_len
        p_start = h_end
        p_end = offset + total_length - 4       # before message_crc

        headers = _decode_headers(buf[h_start:h_end])
        payload = buf[p_start:p_end]
        events.append({"headers": headers, "payload": payload})
        offset += total_length

    return events, buf[offset:]


def _decode_headers(data: bytes) -> dict:
    headers: dict = {}
    o = 0
    dlen = len(data)

    while o < dlen:
        if o + 1 > dlen:
            break
        name_len = data[o]; o += 1
        if o + name_len > dlen:
            break
        name = data[o : o + name_len].decode("utf-8"); o += name_len
        if o + 1 > dlen:
            break
        htype = data[o]; o += 1

        if htype == 7:                          # string
            if o + 2 > dlen: break
            vl = struct.unpack("!H", data[o : o + 2])[0]; o += 2
            if o + vl > dlen: break
            headers[name] = data[o : o + vl].decode("utf-8"); o += vl
        elif htype == 6:                        # byte-array
            if o + 2 > dlen: break
            vl = struct.unpack("!H", data[o : o + 2])[0]; o += 2
            o += vl
        elif htype in (0, 1):                   # bool true / false
            headers[name] = (htype == 0)
        elif htype == 2: o += 1                 # byte
        elif htype == 3: o += 2                 # short
        elif htype == 4:                        # int
            if o + 4 <= dlen:
                headers[name] = struct.unpack("!i", data[o : o + 4])[0]
            o += 4
        elif htype == 5: o += 8                 # long
        elif htype == 8: o += 8                 # timestamp
        elif htype == 9: o += 16                # uuid
        else:
            break                               # unknown — stop parsing

    return headers


def _unwrap_event(raw: dict) -> Optional[dict]:
    """Turn one AWS event-stream frame into an Anthropic streaming dict."""
    h = raw["headers"]
    msg_type = h.get(":message-type", "")

    # ── Exception frames ──────────────────────────
    if msg_type == "exception":
        text = raw["payload"].decode("utf-8", errors="replace")
        return {"type": "error", "error": {"message": text}}

    if msg_type != "event":
        return None

    # ── Normal event — payload is JSON with "bytes" key ───
    try:
        wrapper = json.loads(raw["payload"])
    except (json.JSONDecodeError, UnicodeDecodeError):
        return None

    b64 = wrapper.get("bytes")
    if not b64:
        return None

    try:
        inner = base64.b64decode(b64)
        return json.loads(inner)
    except Exception:
        return None


# ══════════════════════════════════════════════════════════════
# Public API
# ══════════════════════════════════════════════════════════════

async def stream_response(
    *,
    messages: list[dict],
    system_prompt: str,
    model_id: str,
    max_tokens: int = 4096,
    thinking_config: Optional[dict] = None,
) -> AsyncIterator[dict]:
    """
    Yields Anthropic streaming event dicts:
      message_start, content_block_start, content_block_delta,
      content_block_stop, message_delta, message_stop
    """
    client = _get_client()

    safe_model = quote(model_id, safe="")
    url = f"{config.BEDROCK_ENDPOINT}/model/{safe_model}/invoke-with-response-stream"

    body: dict = {
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": max_tokens,
        "system": system_prompt,
        "messages": messages,
    }

    if thinking_config and thinking_config.get("enabled"):
        budget = thinking_config.get("budget_tokens", 8000)
        body["thinking"] = {
            "type": "enabled",
            "budget_tokens": budget,
        }
        # max_tokens must cover thinking + response
        if body["max_tokens"] < budget + 1024:
            body["max_tokens"] = budget + 1024

    headers = {
        "Authorization": f"Bearer {config.BEDROCK_API_KEY}",
        "Content-Type": "application/json",
        "Accept": "application/vnd.amazon.eventstream",
    }

    async with client.stream("POST", url, json=body, headers=headers) as resp:
        if resp.status_code != 200:
            err = await resp.aread()
            raise RuntimeError(
                f"Bedrock HTTP {resp.status_code}: {err.decode('utf-8', errors='replace')[:500]}"
            )

        buf = b""
        async for chunk in resp.aiter_bytes():
            buf += chunk
            events, buf = _drain_messages(buf)
            for raw_evt in events:
                parsed = _unwrap_event(raw_evt)
                if parsed:
                    if parsed.get("type") == "error":
                        raise RuntimeError(
                            parsed.get("error", {}).get("message", "Unknown Bedrock error")
                        )
                    yield parsed


async def invoke_model_simple(
    model_id: str,
    prompt: str,
    max_tokens: int = 120,
) -> str:
    """Non-streaming invoke — used for quick tasks like title generation."""
    client = _get_client()

    safe_model = quote(model_id, safe="")
    url = f"{config.BEDROCK_ENDPOINT}/model/{safe_model}/invoke"

    body = {
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": max_tokens,
        "messages": [{"role": "user", "content": prompt}],
    }

    headers = {
        "Authorization": f"Bearer {config.BEDROCK_API_KEY}",
        "Content-Type": "application/json",
        "Accept": "application/json",
    }

    resp = await client.post(url, json=body, headers=headers)
    if resp.status_code != 200:
        raise RuntimeError(f"Bedrock invoke error {resp.status_code}: {resp.text[:300]}")

    data = resp.json()
    parts = data.get("content", [])
    return " ".join(p.get("text", "") for p in parts if p.get("type") == "text")