"""Scaleway API client with async support and model selection."""

from __future__ import annotations

import json

import httpx

from migration_pipeline.config import (
    API_KEY,
    API_TIMEOUT,
    BASE_URL,
    DEFAULT_MODEL_FAMILY,
    MODEL_GLM,
    MODEL_QWEN,
)

MODEL_FAMILY_GLM = "glm"
MODEL_FAMILY_QWEN = "qwen"
_SELECTED_MODEL_FAMILY = DEFAULT_MODEL_FAMILY

_MODEL_PAYLOADS: dict[str, dict] = {
    MODEL_FAMILY_GLM: {
        "model": MODEL_GLM,
        "max_tokens": 16384,
        "temperature": 1,
        "top_p": 0.95,
        "presence_penalty": 0,
        "stream": True,
        "response_format": {"type": "text"},
    },
    MODEL_FAMILY_QWEN: {
        "model": MODEL_QWEN,
        "max_tokens": 16384,
        "temperature": 0.6,
        "top_p": 0.95,
        "presence_penalty": 0,
        "stream": True,
        "response_format": {"type": "text"},
    },
}

_MODEL_TO_FAMILY = {
    profile["model"]: family for family, profile in _MODEL_PAYLOADS.items()
}


def set_model_family(family: str) -> None:
    """Set active model family for the current pipeline run."""
    normalized = family.strip().lower()
    if normalized not in _MODEL_PAYLOADS:
        raise ValueError(f"Unsupported model family: {family}")
    global _SELECTED_MODEL_FAMILY
    _SELECTED_MODEL_FAMILY = normalized


def get_selected_model_family() -> str:
    return _SELECTED_MODEL_FAMILY


def select_model(line_count: int) -> str:
    """Return model configured for the selected family (line_count kept for compatibility)."""
    _ = line_count
    return _MODEL_PAYLOADS[_SELECTED_MODEL_FAMILY]["model"]


async def _read_streaming_content(response: httpx.Response) -> str:
    content_parts: list[str] = []
    async for raw_line in response.aiter_lines():
        line = raw_line.strip()
        if not line or line == "data: [DONE]" or not line.startswith("data: "):
            continue
        payload = line[len("data: "):]
        try:
            data = json.loads(payload)
        except ValueError:
            continue
        choices = data.get("choices", [])
        if not choices:
            continue
        delta = choices[0].get("delta", {})
        token = delta.get("content")
        if token:
            content_parts.append(token)
    return "".join(content_parts)


async def _send_messages(
    client: httpx.AsyncClient,
    model: str,
    messages: list[dict],
) -> str:
    """Core HTTP helper: POST *messages* to the chat-completions endpoint.

    Raises ``httpx.HTTPStatusError`` on non-2xx responses so the caller can
    handle retries / quarantine.
    """
    family = _MODEL_TO_FAMILY.get(model)
    if not family:
        raise ValueError(f"Unsupported model id: {model}")
    model_payload = _MODEL_PAYLOADS[family]
    payload = {**model_payload, "messages": messages, "model": model}
    headers = {
        "Authorization": "Bearer " + API_KEY,
        "Content-Type": "application/json",
    }
    if payload.get("stream"):
        async with client.stream(
            "POST",
            BASE_URL,
            json=payload,
            headers=headers,
            timeout=API_TIMEOUT,
        ) as response:
            response.raise_for_status()
            return await _read_streaming_content(response)
    response = await client.post(BASE_URL, json=payload, headers=headers, timeout=API_TIMEOUT)
    response.raise_for_status()
    return response.json()["choices"][0]["message"]["content"]


async def call_api(
    client: httpx.AsyncClient,
    model: str,
    system_prompt: str,
    user_content: str,
) -> str:
    """Stateless single-turn call (backward-compatible).

    Builds a minimal ``[system, user]`` message list and delegates to
    :func:`_send_messages`.
    """
    messages: list[dict] = [
        {"role": "system", "content": system_prompt},
        {"role": "user",   "content": user_content},
    ]
    return await _send_messages(client, model, messages)


async def call_api_multiturn(
    client: httpx.AsyncClient,
    model: str,
    messages: list[dict],
) -> str:
    """Multi-turn call: pass the full conversation history and return the next reply.

    The caller is responsible for maintaining and extending *messages*
    (appending ``{"role": "assistant", "content": reply}`` after each call).
    """
    return await _send_messages(client, model, messages)