import os
import gradio as gr
from huggingface_hub import InferenceClient
from huggingface_hub.utils import HfHubHTTPError

# Modelo preferido
PREFERRED_MODEL = os.environ.get("MODEL_ID", "mistralai/Mistral-7B-Instruct-v0.2")
# Modelo de fallback atualizado
FALLBACK_MODEL = os.environ.get("FALLBACK_MODEL", "unsloth/Llama-3.2-3B-Instruct")

# token vindo do secret HF_TOKEN do Space (ou env local)
token = os.environ.get("HF_TOKEN")

def _extract_text_from_response(resp):
    if isinstance(resp, str):
        return resp
    try:
        if hasattr(resp, "generated_text"):
            return getattr(resp, "generated_text") or ""
        if hasattr(resp, "text"):
            return getattr(resp, "text") or ""
    except Exception:
        pass
    if isinstance(resp, dict):
        for key in ("generated_text", "generated_texts", "text", "output_text", "result"):
            if key in resp:
                v = resp[key]
                if isinstance(v, list) and v:
                    return v[0] if isinstance(v[0], str) else str(v[0])
                if isinstance(v, str):
                    return v
        if "choices" in resp and isinstance(resp["choices"], list) and resp["choices"]:
            first = resp["choices"][0]
            if isinstance(first, dict):
                if "message" in first and isinstance(first["message"], dict) and "content" in first["message"]:
                    maybe = first["message"]["content"]
                    if isinstance(maybe, str):
                        return maybe
                for k in ("text", "content", "generated_text"):
                    if k in first and isinstance(first[k], str):
                        return first[k]
    try:
        return str(resp)
    except Exception:
        return "<unable to decode response>"

def _call_model(model_id, prompt, max_new_tokens, temperature, top_p):
    client = InferenceClient(model=model_id, token=token)
    return client.text_generation(
        prompt,
        max_new_tokens=int(max_new_tokens),
        temperature=float(temperature),
        top_p=float(top_p),
        do_sample=True,
    )

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    if not token:
        yield "ERRO: variável HF_TOKEN não encontrada. Adicione o secret HF_TOKEN no Settings do Space."
        return

    prompt = f"{system_message}\n\n"
    for user_msg, bot_msg in history:
        if user_msg:
            prompt += f"User: {user_msg}\n"
        if bot_msg:
            prompt += f"Assistant: {bot_msg}\n"
    prompt += f"User: {message}\nAssistant:"

    try:
        out = _call_model(PREFERRED_MODEL, prompt, max_tokens, temperature, top_p)
    except HfHubHTTPError as e:
        try:
            code = e.response.status_code if e.response is not None else None
        except Exception:
            code = None

        if code == 404:
            yield f"Aviso: modelo `{PREFERRED_MODEL}` não disponível via Inference API (404). Tentando fallback para `{FALLBACK_MODEL}`..."
            try:
                out = _call_model(FALLBACK_MODEL, prompt, max_tokens, temperature, top_p)
            except Exception as e2:
                yield f"Falha no fallback para {FALLBACK_MODEL}: {e2}"
                return
        else:
            yield f"ERRO na chamada de inferência: {e}\n(verifique HF_TOKEN, permissões e se o modelo está disponível via Inference API)"
            return
    except Exception as e:
        yield f"Erro inesperado ao chamar a API: {e}"
        return

    text = _extract_text_from_response(out)
    yield text

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a helpful assistant.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.05, label="Temperature"),
        gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
    ],
    title="Chat (Mistral fallback com Llama 3.2 3B)",
)

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