import base64
import mimetypes
import os
import time
import requests
import gradio as gr
from openai import OpenAI, APIError

MODEL_ID = "openrouter/free"
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", None)

# 채팅 목적에 맞지 않는 모델(검열/판정 전용, 재랭킹, 임베딩 등)을 걸러내기
# 위한 키워드입니다. 이런 모델은 ':free'이고 텍스트/이미지 입력을 지원하더라도
# 실제 대화 응답 대신 판정 라벨 등 엉뚱한 출력을 내놓기 때문에 제외합니다.
# (예: nvidia/nemotron-3.5-content-safety:free — 실제로 겪으신 사례)
_NON_CHAT_MODEL_KEYWORDS = ["content-safety", "guard", "rerank", "embed", "moderation"]

# 네트워크 오류 등으로 실시간 조회가 완전히 실패했을 때만 사용하는
# 최후의 안전망(fallback) 목록입니다. 2026-07-02 기준으로 확인된 값입니다.
_SAFETY_NET_TEXT_MODELS = [
    "meta-llama/llama-4-scout:free",
    "deepseek/deepseek-chat-v3.1:free",
]
_SAFETY_NET_VISION_MODELS = [
    "google/gemma-4-31b-it:free",
    "nvidia/nemotron-nano-12b-v2-vl:free",
]

# ── 채팅용 무료 모델 목록 캐시 (텍스트 전용 / 이미지 지원 각각) ────────────
# OpenRouter 공식 문서(openrouter.ai/docs/guides/overview/models)에 따르면
# /api/v1/models 응답은 OpenRouter 쪽에서도 edge에 캐싱되어 제공되므로,
# 우리 쪽에서도 매 요청마다 조회하지 않고 일정 시간(TTL) 동안 재사용합니다.
_models_cache: dict = {
    "text": {"models": [], "fetched_at": 0.0},
    "vision": {"models": [], "fetched_at": 0.0},
}
_CACHE_TTL_SECONDS = 60 * 60  # 1시간


def _fetch_chat_capable_free_models(require_image: bool) -> list:
    """OpenRouter 공식 API(/api/v1/models)를 호출해, 지금 이 순간
    실제 '대화 응답'을 생성하는 무료(:free) 모델 목록을 가져옵니다.
    검열/재랭킹/임베딩 등 비채팅 목적 모델은 항상 제외하며,
    require_image=True이면 이미지 입력까지 지원하는 모델만 남깁니다.
    """
    resp = requests.get("https://openrouter.ai/api/v1/models", timeout=10)
    resp.raise_for_status()
    data = resp.json().get("data", [])

    result = []
    for m in data:
        model_id = m.get("id", "")
        if not model_id.endswith(":free"):
            continue
        if any(kw in model_id.lower() for kw in _NON_CHAT_MODEL_KEYWORDS):
            continue
        modalities = m.get("architecture", {}).get("input_modalities", [])
        if require_image and "image" not in modalities:
            continue
        result.append(model_id)
    return result


def get_chat_capable_free_models(require_image: bool = False, force_refresh: bool = False) -> list:
    """캐시된 '실제 대화용' 무료 모델 목록을 반환합니다.
    캐시가 없거나(최초 호출), TTL(1시간)이 지났거나, force_refresh=True이면
    OpenRouter API를 다시 호출해 갱신합니다.
    """
    cache_key = "vision" if require_image else "text"
    cache = _models_cache[cache_key]
    safety_net = _SAFETY_NET_VISION_MODELS if require_image else _SAFETY_NET_TEXT_MODELS

    now = time.time()
    cache_expired = (now - cache["fetched_at"]) > _CACHE_TTL_SECONDS

    if force_refresh or cache_expired or not cache["models"]:
        try:
            fresh_models = _fetch_chat_capable_free_models(require_image)
            if fresh_models:
                cache["models"] = fresh_models
                cache["fetched_at"] = now
        except Exception:
            # 조회 자체가 실패하면(네트워크 오류 등), 기존 캐시가 있으면
            # 그대로 쓰고, 캐시조차 없으면 최후의 안전망 목록을 사용합니다.
            pass

    return cache["models"] or safety_net


# ── OpenRouter 클라이언트 초기화 ─────────────────────────────────────────
# Hugging Face를 전혀 거치지 않고, OpenRouter 자체 서버(openrouter.ai)로
# 직접 요청을 보냅니다. OpenAI SDK의 base_url만 OpenRouter 주소로
# 바꾸는, OpenRouter 공식 문서가 안내하는 표준 사용 방식입니다.
client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENROUTER_API_KEY,
)


def _encode_image_to_data_url(path: str) -> str:
    """로컬 이미지 파일 경로를 OpenRouter가 요구하는 base64 data URL
    형식(data:image/png;base64,...)으로 변환합니다.
    (OpenRouter 공식 문서: Image Inputs 가이드 기준)
    """
    mime_type, _ = mimetypes.guess_type(path)
    mime_type = mime_type or "image/png"
    with open(path, "rb") as f:
        encoded = base64.b64encode(f.read()).decode("utf-8")
    return f"data:{mime_type};base64,{encoded}"


def _build_content(text: str | None, files: list) -> list | str:
    """텍스트와 이미지 파일들을 OpenAI/OpenRouter 호환 content 배열로
    변환합니다. 이미지가 없으면 텍스트만 담긴 문자열을 그대로 반환합니다.
    """
    if not files:
        return text or ""

    content = []
    if text:
        content.append({"type": "text", "text": text})
    for path in files:
        try:
            content.append({
                "type": "image_url",
                "image_url": {"url": _encode_image_to_data_url(path)},
            })
        except Exception:
            # 이미지가 아니거나 읽을 수 없는 파일은 조용히 건너뜁니다.
            continue
    return content


# ── Gradio 버전/멀티모달 여부에 따라 history 형식이 다르므로 다양한 형식 지원 ──
def _normalize_history(history: list) -> list:
    """history는 Gradio 버전 및 멀티모달 여부에 따라
    - 신형(messages) 형식: [{"role": "user", "content": "..."}, ...]
    - 구형(tuples) 형식:  [["user 메시지", "bot 메시지"], ...]
    - 이미지 첨부 시: content가 파일 경로가 담긴 튜플/리스트인 경우도 있음
    다양한 형태로 들어올 수 있습니다. 어떤 형식이 와도 OpenAI/OpenRouter
    호환 messages 배열로 통일해서 반환합니다.
    """
    messages = []
    if not history:
        return messages

    for turn in history:
        if isinstance(turn, dict):
            role = turn.get("role", "user")
            content = turn.get("content", "")
        elif isinstance(turn, (list, tuple)) and len(turn) == 2:
            user_msg, bot_msg = turn
            if user_msg:
                messages.append({"role": "user", "content": user_msg})
            if bot_msg:
                messages.append({"role": "assistant", "content": bot_msg})
            continue
        else:
            role = getattr(turn, "role", "user")
            content = getattr(turn, "content", "")

        if role == "model":
            role = "assistant"

        # content가 이미지 파일 경로(튜플/리스트)인 경우, 이미지 메시지로 변환합니다.
        if isinstance(content, (tuple, list)) and content:
            path = content[0]
            if isinstance(path, str) and os.path.isfile(path):
                messages.append({
                    "role": role,
                    "content": _build_content(None, [path]),
                })
                continue

        # content가 {"path": "..."} 형태로 들어오는 경우도 처리합니다.
        if isinstance(content, dict) and "path" in content:
            messages.append({
                "role": role,
                "content": _build_content(None, [content["path"]]),
            })
            continue

        messages.append({"role": role, "content": content})

    return messages


# ── 대화 텍스트 생성 함수 (OpenRouter chat.completions 스트리밍 방식) ────────
def _call_openrouter_stream(request_model: str, messages: list, extra_body: dict):
    return client.chat.completions.create(
        model=request_model,
        messages=messages,
        max_tokens=8192,
        temperature=0.7,
        top_p=0.95,
        stream=True,  # 응답을 토큰(글자) 단위 조각으로 실시간 수신합니다.
        extra_headers={
            # OpenRouter 리더보드에 앱을 표시하고 싶을 때 사용하는
            # 선택 항목입니다. 필수는 아니므로 비워두셔도 됩니다.
            "HTTP-Referer": "https://huggingface.co/spaces",
            "X-Title": "GS AI",
        },
        extra_body=extra_body,
    )


def predict_chat(message, history: list):
    # multimodal=True일 때 message는 {"text": "...", "files": ["경로", ...]}
    # 형태의 딕셔너리로 들어옵니다. (Gradio 공식 문서 기준)
    if isinstance(message, dict):
        text = message.get("text") or ""
        files = message.get("files") or []
    else:
        text = message
        files = []

    messages = _normalize_history(history)
    messages.append({"role": "user", "content": _build_content(text, files)})

    has_image = bool(files)

    def _resolve_request_target(force_refresh: bool = False):
        # 텍스트 전용이든 이미지 첨부든, "검열/판정 전용 모델"이 실제 대화
        # 응답 자리에 뽑히는 문제를 막기 위해 항상 필터링된 캐시 목록을
        # 사용합니다. (openrouter/free를 그대로 쓰면, 채팅 목적에 맞지 않는
        # 모델이 무작위로 선택될 수 있다는 것을 실제로 확인했습니다.)
        candidate_models = get_chat_capable_free_models(
            require_image=has_image, force_refresh=force_refresh
        )
        # OpenRouter의 fallback 기능은 'models' 배열에 최대 3개까지만 허용합니다.
        # (OpenRouter 서버 응답 기준: "'models' array must have 3 items or fewer.")
        candidate_models = candidate_models[:3]
        return candidate_models[0], {"models": candidate_models}

    request_model, extra_body = _resolve_request_target()

    try:
        try:
            stream = _call_openrouter_stream(request_model, messages, extra_body)
        except APIError as e:
            # "지원 모델을 못 찾음(404)" 오류가 나면, 캐시가 낡았을 가능성이
            # 있으므로 강제로 한 번 더 새로 조회해서 딱 한 번만 재시도합니다.
            if getattr(e, "status_code", None) == 404:
                request_model, extra_body = _resolve_request_target(force_refresh=True)
                stream = _call_openrouter_stream(request_model, messages, extra_body)
            else:
                raise

        partial_output = ""
        resolved_model = request_model

        for chunk in stream:
            # 각 조각(chunk)에도 실제 응답 중인 모델명이 담겨 있어,
            # 스트리밍 도중 계속 최신값으로 갱신합니다.
            if getattr(chunk, "model", None):
                resolved_model = chunk.model

            if not chunk.choices:
                continue
            delta = chunk.choices[0].delta
            token = getattr(delta, "content", None)
            if token:
                partial_output += token
                # yield할 때마다 Gradio 화면이 실시간으로 갱신됩니다.
                yield partial_output

        # 스트리밍이 끝난 뒤, 실제 사용 모델명을 마지막에 덧붙여 최종 갱신합니다.
        yield f"{partial_output}\n\n---\n_(실제 사용 모델: `{resolved_model}`)_"

    except APIError as e:
        # 무료 모델의 분당/일일 요청 한도 초과, 이미지 미지원 모델 선택,
        # 모델 일시 중단 등 OpenRouter 쪽 사유로 오류가 나는 경우입니다.
        yield f"❌ OpenRouter API 오류: {str(e)}"
    except Exception as e:
        yield f"❌ 네트워크 통신 중 오류가 발생했습니다: {str(e)}"


# ── UI 빌드 및 API 개설 ──────────────────────────────────────────────────────
with gr.Blocks() as demo:
    gr.Markdown(
        "# 🤖 GS-AI API Server (OpenRouter Mode)\n"
        "무료 모델 중 채팅에 적합한 모델을 자동으로 골라 응답합니다.\n\n"
        "텍스트와 함께 이미지를 첨부해서 질문하실 수 있습니다. "
        "(단, 그 시점에 이용 가능한 무료 모델이 없으면 오류가 날 수 있습니다.)"
    )

    chatbot_ui = gr.Chatbot(height=480)
    textbox_ui = gr.MultimodalTextbox(
        placeholder="메시지를 입력하거나 이미지를 첨부하세요...",
        file_types=["image"],
        file_count="multiple",
    )

    gr.ChatInterface(
        fn=predict_chat,
        chatbot=chatbot_ui,
        textbox=textbox_ui,
        multimodal=True,  # message 인자가 {"text": ..., "files": [...]} 딕셔너리로 전달됩니다.
        api_name="chat",  # 외부 연동을 위한 /chat 경로 명시 활성화
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        theme=gr.themes.Soft(),
    )