import logging
import os
from collections.abc import Iterator

import gradio as gr
from cohere import ClientV2

logger = logging.getLogger(__name__)

model_id = "tiny-aya-global"

# Initialize Cohere client
api_key = os.getenv("COHERE_API_KEY")
if not api_key:
    raise ValueError("COHERE_API_KEY environment variable is required")
client = ClientV2(api_key=api_key, client_name="hf-tiny-aya-global")


def _extract_text(content: object) -> str:
    """Extract plain text from any Cohere content shape.

    Handles plain strings, objects with a `.text` attribute,
    and lists of content blocks (e.g. [{'text': '...', 'type': 'text'}]).
    """
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = [_extract_text(block) for block in content]
        return "".join(parts)
    text = getattr(content, "text", None)
    if text is not None:
        return str(text)
    if isinstance(content, dict):
        return str(content.get("text", ""))
    return ""


def generate(
    message: str,
    history: list[dict],
    system_prompt: str = "",
    temperature: float = 0.1,
    max_new_tokens: int = 700,
) -> Iterator[str]:
    """Stream a response from the Cohere API for the given message and conversation history."""
    messages: list[dict[str, str]] = []
    system_prompt = (system_prompt or "").strip()
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})

    for item in history:
        role = item.get("role")
        content = _extract_text(item.get("content"))
        if role in ("assistant", "user") and content:
            messages.append({"role": role, "content": content})

    current_text = str(message or "").strip()
    if not current_text:
        yield ""
        return
    messages.append({"role": "user", "content": current_text})

    try:
        response = client.chat_stream(
            model=model_id,
            messages=messages,
            temperature=temperature,
            max_tokens=max_new_tokens,
        )

        output = ""
        for event in response:
            if getattr(event, "type", None) in ("content-delta", "content-start"):
                delta = getattr(event, "delta", None)
                if delta is None:
                    continue
                msg = getattr(delta, "message", None)
                if msg is None:
                    continue
                text = _extract_text(getattr(msg, "content", None))
                if text:
                    output += text
                    yield output

    except Exception:
        logger.exception("Cohere API error")
        gr.Warning("Something went wrong generating a response. Please try again.")
        yield ""


examples = [
    ["Explica en español qué significa la palabra japonesa 'ikigai' y da un ejemplo práctico."],
    ["اكتب فقرة قصيرة تصف غروب الشمس في الصحراء"],
    ["Kwa nini ni muhimu kujifunza lugha zaidi ya moja? Toa sababu tatu."],
    ["一个从未见过大海的山村孩子，第一次来到海边。用三到五句话描述他的感受。"],
    [
        "Translate the following sentence from Basque to French: "
        "'Hizkuntza-eredu handiek milioika testu erabiltzen dituzte ikasteko, "
        "baina hizkuntza txikientzat datu gutxiago dago.'"
    ],
    [
        "ਜੇਕਰ ਕੋਈ ਵਿਅਕਤੀ ਪਹਿਲੀ ਵਾਰ ਵਿਦੇਸ਼ ਜਾ ਰਿਹਾ ਹੈ, ਤਾਂ ਉਸ ਨੂੰ ਕਿਹੜੀਆਂ ਗੱਲਾਂ ਦਾ "
        "ਧਿਆਨ ਰੱਖਣਾ ਚਾਹੀਦਾ ਹੈ? ਪੰਜ ਸੁਝਾਅ ਦਿਓ।"
    ],
    [
        "Eglurwch mewn tair brawddeg pam mae bioamrywiaeth yn bwysig i ecosystemau."
    ],
    [
        "ถ้าคุณต้องการเริ่มต้นออกกำลังกายเป็นประจำ ควรเริ่มต้นอย่างไร? ให้คำแนะนำสามข้อ"
    ]
]

example_labels = [
    "Spanish — Explain 'ikigai'",
    "Arabic — Describe a desert sunset",
    "Swahili — Why learn multiple languages?",
    "Chinese — A child sees the ocean",
    "English — Basque to French translation",
    "Punjabi — Travel tips for first-timers",
    "Welsh — Why biodiversity is important to ecosystems",
    "Thai — How to start exercising regularly",
]

DESCRIPTION = (
    "**[Tiny Aya](https://huggingface.co/CohereLabs/tiny-aya-global)** is a lightweight "
    "multilingual language model by [Cohere Labs](https://cohere.com/research). "
    "Try chatting in any of 70+ supported languages!"
)

demo = gr.ChatInterface(
    fn=generate,
    chatbot=gr.Chatbot(min_height=600),
    textbox=gr.Textbox(
        autofocus=True,
        placeholder="Type a message in any language... / Escribe en cualquier idioma... / أكتب بأي لغة...",
    ),
    additional_inputs=[
        gr.Textbox(
            label="System Prompt (optional)",
            placeholder="e.g. You are a helpful multilingual assistant. Always respond in the user's language.",
            lines=3,
        ),
        gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, step=0.05, value=0.1),
        gr.Slider(label="Max New Tokens", minimum=100, maximum=2000, step=10, value=700),
    ],
    stop_btn=False,
    title="Tiny Aya",
    description=DESCRIPTION,
    examples=examples,
    example_labels=example_labels,
    run_examples_on_click=True,
    cache_examples=False,
    delete_cache=(1800, 1800),
    save_history=True,
)

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(
            primary_hue="green",
            font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
        ),
        css_paths="style.css",
    )