"""Gradio Space wrapper for the Copyleft Cultivars Qwen3 v5clean champion.

The model weights live in a separate Hugging Face model repository. Set
MODEL_ID to that repository before starting the Space; the default is the
proposed public model-repository name used by the deployment guide.
"""

from __future__ import annotations

import os
from typing import Any

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


MODEL_ID = os.getenv("MODEL_ID", "CopyleftCultivars/qwen3-v5clean")
MODEL_REVISION = os.getenv("MODEL_REVISION")
HF_TOKEN = os.getenv("HF_TOKEN")
SYSTEM_PROMPT = """You are the Copyleft Cultivars natural-farming assistant.

Give practical, careful answers about Korean Natural Farming, regenerative
agriculture, soil biology, composting, fermentation, crop care, and farm
inputs. Explain assumptions and units. Do not invent measurements, sources,
tool results, or local conditions. When a question depends on local weather,
soil tests, regulations, or a professional diagnosis, say what information is
missing and recommend an appropriate local source or professional.
"""


def _load_model() -> tuple[Any, Any]:
    """Load the pinned model and tokenizer once when the Space starts."""
    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    common_kwargs: dict[str, Any] = {
        "token": HF_TOKEN,
        "torch_dtype": dtype,
        "device_map": "auto",
    }
    if MODEL_REVISION:
        common_kwargs["revision"] = MODEL_REVISION

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN, revision=MODEL_REVISION)
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **common_kwargs)
    model.eval()
    return tokenizer, model


tokenizer, model = _load_model()
MODEL_DEVICE = next(model.parameters()).device


def _message_content(value: Any) -> str:
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        text = value.get("text")
        if isinstance(text, str):
            return text
    return ""


def _normalise_history(history: list[Any] | None) -> list[dict[str, str]]:
    """Accept Gradio message history and ignore unsupported multimodal entries."""
    normalised: list[dict[str, str]] = []
    for entry in history or []:
        if isinstance(entry, dict):
            role = entry.get("role")
            content = _message_content(entry.get("content"))
            if role in {"user", "assistant"} and content:
                normalised.append({"role": role, "content": content})
            continue
        if isinstance(entry, (list, tuple)) and len(entry) >= 2:
            user_text = _message_content(entry[0])
            assistant_text = _message_content(entry[1])
            if user_text:
                normalised.append({"role": "user", "content": user_text})
            if assistant_text:
                normalised.append({"role": "assistant", "content": assistant_text})
    return normalised


def _tokenize_messages(messages: list[dict[str, str]], enable_thinking: bool) -> Any:
    """Render Qwen chat messages, retaining compatibility with older Transformers."""
    template_kwargs = {
        "tokenize": True,
        "add_generation_prompt": True,
        "return_tensors": "pt",
    }
    try:
        return tokenizer.apply_chat_template(
            messages,
            enable_thinking=enable_thinking,
            **template_kwargs,
        )
    except TypeError:
        return tokenizer.apply_chat_template(messages, **template_kwargs)


def chat(
    message: str,
    history: list[Any] | None,
    enable_thinking: bool,
    max_new_tokens: int,
    temperature: float,
) -> str:
    """Generate one answer from the v5clean champion."""
    user_message = (message or "").strip()
    if not user_message:
        return "Please enter a question."

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages.extend(_normalise_history(history))
    messages.append({"role": "user", "content": user_message})

    rendered = _tokenize_messages(messages, enable_thinking)
    if isinstance(rendered, torch.Tensor):
        model_inputs = {"input_ids": rendered}
    else:
        model_inputs = dict(rendered)
    model_inputs = {
        key: value.to(MODEL_DEVICE) if hasattr(value, "to") else value
        for key, value in model_inputs.items()
    }
    prompt_length = model_inputs["input_ids"].shape[-1]

    generation_kwargs: dict[str, Any] = {
        **model_inputs,
        "max_new_tokens": int(max_new_tokens),
        "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
        "eos_token_id": tokenizer.eos_token_id,
        "do_sample": temperature > 0.01,
    }
    if temperature > 0.01:
        generation_kwargs["temperature"] = float(temperature)
        generation_kwargs["top_p"] = 0.9

    with torch.inference_mode():
        generated = model.generate(**generation_kwargs)
    answer_tokens = generated[0, prompt_length:]
    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True).strip()
    return answer or "The model returned an empty response. Please try again."


thinking_toggle = gr.Checkbox(
    value=False,
    label="Enable reasoning",
    info="Show the model's reasoning mode when supported by the tokenizer.",
)
max_tokens_slider = gr.Slider(
    minimum=64,
    maximum=2048,
    value=512,
    step=64,
    label="Maximum new tokens",
)
temperature_slider = gr.Slider(
    minimum=0.0,
    maximum=1.2,
    value=0.6,
    step=0.1,
    label="Temperature",
)

demo = gr.ChatInterface(
    fn=chat,
    type="messages",
    additional_inputs=[thinking_toggle, max_tokens_slider, temperature_slider],
    title="Copyleft Cultivars · Qwen3 v5clean",
    description=(
        "A research preview of the current English v5clean champion for "
        "natural-farming questions. Verify recommendations locally before use."
    ),
    examples=[
        [
            "How can I prepare a small-batch fermented plant juice safely?",
            False,
            512,
            0.6,
        ],
        [
            "What information should I collect before diagnosing poor soil drainage?",
            False,
            512,
            0.6,
        ],
        [
            "How do Korean Natural Farming inputs differ from ordinary compost tea?",
            False,
            512,
            0.6,
        ],
    ],
)


if __name__ == "__main__":
    demo.launch()