import logging
import os
from typing import Any

import gradio as gr
import requests


logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Hugging Face Inference Router settings
API_BASE_URL = os.getenv(
    "API_BASE_URL", "https://router.huggingface.co/v1/chat/completions"
).rstrip("/")
MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/DeepSeek-R1:fastest").strip()
HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
REQUEST_TIMEOUT_SECONDS = float(os.getenv("REQUEST_TIMEOUT_SECONDS", "20"))


def _auth_headers() -> dict[str, str]:
    if not HF_TOKEN:
        return {}
    return {"Authorization": f"Bearer {HF_TOKEN}"}


def query(payload: dict[str, Any]) -> dict[str, Any] | list[Any]:
    request_body = {
        "model": MODEL_NAME,
        "messages": [{"role": "user", "content": payload.get("inputs", "")}],
    }
    try:
        response = requests.post(
            API_BASE_URL,
            headers=_auth_headers(),
            json=request_body,
            timeout=REQUEST_TIMEOUT_SECONDS,
        )
    except requests.exceptions.Timeout:
        logger.warning("Request timed out for model=%s", MODEL_NAME)
        return {"error": "Request timed out. Please try again."}
    except requests.exceptions.RequestException as exc:
        logger.exception("Request failed: %s", exc)
        return {"error": f"Network error: {exc}"}

    if not response.ok:
        error_message = f"HTTP {response.status_code}"
        if response.status_code in (401, 403):
            logger.error("Hugging Face auth failed for model=%s", MODEL_NAME)
            return {
                "error": (
                    "Unauthorized request to Hugging Face. Check that HF_TOKEN is set "
                    "correctly in the Space secrets and that it has access to the model."
                )
            }
        if response.status_code == 404:
            logger.error("Hugging Face route not found for model=%s", MODEL_NAME)
            return {
                "error": (
                    "Hugging Face route not found. Check API_BASE_URL and MODEL_NAME, "
                    "and make sure the model is available on the router."
                )
            }
        try:
            body = response.json()
            if isinstance(body, dict) and body.get("error"):
                error_message = f"{error_message}: {body['error']}"
        except ValueError:
            # Keep the HTTP status only when body is not valid JSON.
            pass
        logger.error("Inference API error: %s", error_message)
        return {"error": error_message}

    try:
        return response.json()
    except ValueError:
        logger.error("Inference API returned non-JSON response")
        return {"error": "Model returned an invalid response format."}


def _extract_text(data: dict[str, Any] | list[Any]) -> str:
    if isinstance(data, dict):
        if data.get("error"):
            return f"Error: {data['error']}"
        choices = data.get("choices")
        if isinstance(choices, list) and choices:
            first_choice = choices[0]
            if isinstance(first_choice, dict):
                message = first_choice.get("message")
                if isinstance(message, dict):
                    content = message.get("content")
                    if isinstance(content, str) and content.strip():
                        return content.strip()
                content = first_choice.get("text")
                if isinstance(content, str) and content.strip():
                    return content.strip()
        return "No response generated."

    if isinstance(data, list) and data:
        first_item = data[0]
        if isinstance(first_item, dict):
            for key in ("generated_text", "answer", "summary_text", "text"):
                value = first_item.get(key)
                if isinstance(value, str) and value.strip():
                    return value.strip()
        if isinstance(first_item, str) and first_item.strip():
            return first_item.strip()

    return "No response generated."


def respond(message: str) -> str:
    clean_message = (message or "").strip()
    if not clean_message:
        return "Please enter a valid message."

    data = query({"inputs": clean_message})
    return _extract_text(data)


iface = gr.Interface(
    fn=respond,
    inputs=gr.Textbox(lines=3, placeholder="Ask something..."),
    outputs=gr.Textbox(label="Assistant Response"),
    title="AI Assistant",
    description="Simple AI assistant powered by the Hugging Face Inference API.",
)


if __name__ == "__main__":
    port = int(os.getenv("PORT", "7860"))
    iface.launch(server_name="0.0.0.0", server_port=port, show_error=True)