from __future__ import annotations
import json
import base64

try:
    from huggingface_hub import InferenceClient as _HFClient
    _HF_AVAILABLE = True
except ImportError:
    _HF_AVAILABLE = False

import urllib.request

HF_API_BASE = "https://api-inference.huggingface.co/models"
HF_ROUTER_BASE = "https://router.huggingface.co/hf-inference/models"

def _try_router_url(model_id: str, payload: dict, api_key: str, timeout: int = 60) -> str:
    """Try the new HF Router endpoint as fallback."""
    url = f"{HF_ROUTER_BASE}/{model_id}/v1/chat/completions"
    data = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
    req = urllib.request.Request(url, data=data, headers=headers, method="POST")
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        result = json.loads(resp.read().decode("utf-8"))
    return result["choices"][0]["message"]["content"]

def generate(model_id: str, prompt: str, system: str = "",
             api_key: str = "", max_tokens: int = 1024, temperature: float = 0.3) -> str:
    """Text generation — MiniCPM concept extraction and Tiny Aya translation."""
    if _HF_AVAILABLE and api_key:
        client = _HFClient(api_key=api_key)
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        response = client.chat_completion(
            messages=messages, model=model_id,
            max_tokens=max_tokens, temperature=temperature
        )
        return response.choices[0].message.content

    # Fallback: raw HTTP
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})
    payload = {"inputs": {"messages": messages},
               "parameters": {"max_new_tokens": max_tokens, "temperature": temperature}}
    data = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    req = urllib.request.Request(f"{HF_API_BASE}/{model_id}", data=data, headers=headers, method="POST")
    with urllib.request.urlopen(req, timeout=60) as resp:
        result = json.loads(resp.read().decode("utf-8"))
    if isinstance(result, list): return result[0].get("generated_text", str(result[0]))
    if isinstance(result, dict): return result.get("generated_text", result.get("text", str(result)))
    return str(result)

def vision_generate(model_id: str, image_b64: str, prompt: str,
                    api_key: str = "", max_tokens: int = 1024) -> str:
    """Vision+language — MiniCPM-V for OCR and visual understanding."""
    if _HF_AVAILABLE and api_key:
        client = _HFClient(api_key=api_key)
        image_data_url = f"data:image/png;base64,{image_b64}"
        messages = [{"role": "user", "content": [
            {"type": "image_url", "image_url": {"url": image_data_url}},
            {"type": "text", "text": prompt},
        ]}]
        response = client.chat_completion(
            messages=messages, model=model_id, max_tokens=max_tokens
        )
        return response.choices[0].message.content

    # Fallback: raw HTTP
    url = f"{HF_API_BASE}/{model_id}"
    payload = {"inputs": {"image": image_b64, "question": prompt},
               "parameters": {"max_new_tokens": max_tokens}}
    data = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    req = urllib.request.Request(url, data=data, headers=headers, method="POST")
    with urllib.request.urlopen(req, timeout=90) as resp:
        result = json.loads(resp.read().decode("utf-8"))
    if isinstance(result, list): return str(result[0])
    if isinstance(result, dict): return result.get("generated_text", result.get("text", str(result)))
    return str(result)

def transcribe(model_id: str, audio_bytes: bytes, api_key: str = "") -> str:
    """ASR — Whisper via HuggingFace for voice input."""
    if _HF_AVAILABLE and api_key:
        client = _HFClient(api_key=api_key)
        result = client.automatic_speech_recognition(audio=audio_bytes, model=model_id)
        return result.text if hasattr(result, "text") else str(result)

    # Fallback: raw HTTP
    url = f"{HF_API_BASE}/{model_id}"
    headers = {"Content-Type": "audio/wav"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    req = urllib.request.Request(url, data=audio_bytes, headers=headers, method="POST")
    with urllib.request.urlopen(req, timeout=60) as resp:
        result = json.loads(resp.read().decode("utf-8"))
    return result.get("text", "")