import os

from cerebras.cloud.sdk import Cerebras


def get_cerebras_client() -> Cerebras:
    api_key = os.getenv("CEREBRAS_API_KEY", "")
    if not api_key:
        raise ValueError("CEREBRAS_API_KEY is not configured.")
    return Cerebras(api_key=api_key)


def run_chat_completion(user_message: str) -> dict:
    model = os.getenv("CEREBRAS_MODEL", "gpt-oss-120b")
    client = get_cerebras_client()

    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": user_message,
            }
        ],
        model=model,
    )

    content = ""
    choices = getattr(chat_completion, "choices", []) or []
    if choices and getattr(choices[0], "message", None):
        content = choices[0].message.content or ""

    return {
        "model": model,
        "content": content,
        "raw": chat_completion.model_dump(),
    }


def run_chat_completion_stream(user_message: str):
    model = os.getenv("CEREBRAS_MODEL", "gpt-oss-120b")
    client = get_cerebras_client()

    stream = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": user_message,
            }
        ],
        model=model,
        stream=True,
    )

    for event in stream:
        choices = getattr(event, "choices", None) or []
        if not choices:
            continue

        delta = getattr(choices[0], "delta", None)
        piece = ""

        if delta is not None:
            piece = getattr(delta, "content", "") or ""
            if not piece and isinstance(delta, dict):
                piece = delta.get("content", "") or ""

        if isinstance(piece, list):
            text_parts = []
            for item in piece:
                if isinstance(item, dict):
                    text_parts.append(str(item.get("text", "")))
                else:
                    text_parts.append(str(getattr(item, "text", "")))
            piece = "".join(text_parts)

        if piece:
            yield piece