import os
import base64
import logging
import subprocess
import tempfile
import threading
import time
import uuid
from pathlib import Path

import gradio as gr
import torch
import uvicorn
from fastapi import Depends, FastAPI, Header, HTTPException
from PIL import Image
from pydantic import BaseModel
from transformers import AutoModelForImageTextToText, AutoProcessor

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("omni-api")

MODEL_ID = os.getenv("OMNI_MODEL_ID", "LiquidAI/LFM3-VL-2B")
SERVED_NAME = os.getenv("OMNI_MODEL_NAME", "lfm3-omni-2b")
API_KEY = os.environ.get("OMNI_API_KEY", "")
if not API_KEY:
    raise RuntimeError("OMNI_API_KEY is not set. Add it to the Space secrets.")

ENABLE_VIDEO = True
VIDEO_MAX_FRAMES = 64
DEFAULT_MAX_TOKENS = 512
HARD_MAX_TOKENS = 1024
MAX_BYTES = 25 * 1024 * 1024

# ---------- model ----------

if torch.cuda.is_available():
    DEVICE, DTYPE = "cuda", torch.bfloat16
else:
    DEVICE, DTYPE = "cpu", torch.float32

log.info("Loading %s on %s (%s)...", MODEL_ID, DEVICE, DTYPE)
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID, dtype=DTYPE, device_map=DEVICE, trust_remote_code=True
).eval()
log.info("Model loaded.")

# The model handles one generation at a time.
model_lock = threading.Lock()


# ---------- helpers ----------

def to_wav(src: str, dst: str) -> bool:
    r = subprocess.run(
        ["ffmpeg", "-y", "-loglevel", "error", "-i", src,
         "-vn", "-ac", "1", "-ar", "16000", dst],
        capture_output=True,
    )
    return r.returncode == 0 and Path(dst).exists()


def video_has_audio(path: str) -> bool:
    r = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "a",
         "-show_entries", "stream=codec_type", "-of", "csv=p=0", path],
        capture_output=True, text=True,
    )
    return bool(r.stdout.strip())


def video_frame_count(path: str) -> int:
    r = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0", "-count_frames",
         "-show_entries", "stream=nb_read_frames", "-of", "csv=p=0", path],
        capture_output=True, text=True,
    )
    try:
        return int(r.stdout.strip())
    except ValueError:
        return 0


def collect_video_paths(messages: list[dict]) -> list[str]:
    paths = []
    for m in messages:
        content = m.get("content")
        if isinstance(content, list):
            for part in content:
                if part.get("type") == "video":
                    paths.append(part["video"])
    return paths


def decode_b64(data: str) -> bytes:
    try:
        raw = base64.b64decode(data, validate=True)
    except Exception:
        raise HTTPException(400, "Invalid base64 data")
    if len(raw) > MAX_BYTES:
        raise HTTPException(413, "File is larger than 25 MB")
    return raw


def decode_data_url(url: str) -> tuple[str, bytes]:
    if not url.startswith("data:") or ";base64," not in url:
        raise HTTPException(400, "Only base64 data URLs are supported")
    header, payload = url.split(",", 1)
    mime = header[5:].split(";")[0]
    return mime, decode_b64(payload)


def safe_ext(name: str, default: str) -> str:
    cleaned = "".join(c for c in name.lower() if c.isalnum())[:8]
    return cleaned or default


def convert_content(content, workdir: str, msg_idx: int):
    """Turn an OpenAI-style content value into the format the processor expects."""
    if isinstance(content, str):
        return content, False
    if not isinstance(content, list):
        raise HTTPException(400, "content must be a string or a list of parts")

    out, has_video = [], False
    for i, part in enumerate(content):
        ptype = part.get("type")
        tag = f"m{msg_idx}_p{i}"

        if ptype == "text":
            out.append({"type": "text", "text": part.get("text", "")})

        elif ptype == "image_url":
            _, data = decode_data_url(part["image_url"]["url"])
            src = os.path.join(workdir, f"{tag}.bin")
            dst = os.path.join(workdir, f"{tag}.png")
            Path(src).write_bytes(data)
            try:
                Image.open(src).convert("RGB").save(dst, "PNG")
            except Exception:
                raise HTTPException(400, "Could not decode image")
            out.append({"type": "image", "image": dst})

        elif ptype == "input_audio":
            audio = part["input_audio"]
            data = decode_b64(audio["data"])
            src = os.path.join(workdir, f"{tag}.{safe_ext(audio.get('format', 'wav'), 'wav')}")
            dst = os.path.join(workdir, f"{tag}.wav")
            Path(src).write_bytes(data)
            if not to_wav(src, dst):
                raise HTTPException(400, "Could not decode audio")
            out.append({"type": "audio", "audio": dst})

        elif ptype == "video_url":
            if not ENABLE_VIDEO:
                raise HTTPException(400, "Video is disabled on this server")
            mime, data = decode_data_url(part["video_url"]["url"])
            ext = "mp4"
            if "quicktime" in mime:
                ext = "mov"
            elif "webm" in mime:
                ext = "webm"
            elif "matroska" in mime:
                ext = "mkv"
            dst = os.path.join(workdir, f"{tag}.{ext}")
            Path(dst).write_bytes(data)

            # Clips with fewer than 2 frames (e.g. a static GIF) can't be processed as video.
            if video_frame_count(dst) < 2:
                frame = os.path.join(workdir, f"{tag}_frame.png")
                subprocess.run(
                    ["ffmpeg", "-y", "-loglevel", "error", "-i", dst,
                     "-frames:v", "1", frame],
                    capture_output=True,
                )
                if not Path(frame).exists():
                    raise HTTPException(400, "Could not read video")
                out.append({"type": "image", "image": frame})
                continue

            out.append({"type": "video", "video": dst})
            has_video = True

        else:
            raise HTTPException(400, f"Unsupported content type: {ptype}")

    return out, has_video


def convert_messages(messages: list[dict], workdir: str):
    converted, has_video = [], False
    for idx, m in enumerate(messages):
        role = m.get("role")
        if role not in ("system", "user", "assistant"):
            raise HTTPException(400, f"Unsupported role: {role}")
        content, v = convert_content(m.get("content", ""), workdir, idx)
        has_video = has_video or v
        converted.append({"role": role, "content": content})
    return converted, has_video


def run_model(messages: list[dict], has_video: bool, max_new_tokens: int) -> str:
    videos = collect_video_paths(messages) if has_video else []
    # Only ask for the soundtrack when every video actually has one.
    use_audio = bool(videos) and all(video_has_audio(v) for v in videos)
    audio_kwargs = (
        {"load_audio_from_video": True, "processor_kwargs": {"use_audio_in_video": True}}
        if use_audio else {}
    )

    # Close the reasoning block so the model answers directly.
    # The template's generation prompt opens <think>; this prefilled assistant turn
    # supplies an empty think block and the model continues from there.
    prefilled = messages + [
        {"role": "assistant", "content": "<think>\n\n</think>\n\n"}
    ]

    def build_inputs(kwargs: dict):
        return processor.apply_chat_template(
            prefilled,
            add_generation_prompt=False,
            continue_final_message=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            **kwargs,
        ).to(DEVICE)

    try:
        inputs = build_inputs(audio_kwargs)
    except Exception:
        if not videos:
            raise
        log.exception("Video preprocessing with audio failed; retrying with frames only")
        inputs = build_inputs({})

    with model_lock, torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            return_dict_in_generate=True,
        )

    start = inputs["input_ids"].shape[1]
    text = processor.decode(
        out.sequences[0, start:],
        skip_special_tokens=True,
    ).strip()
    # Safety net in case the model still emits a closing tag.
    if "</think>" in text:
        text = text.split("</think>", 1)[1].strip()
    return text


# ---------- API ----------

app = FastAPI(title="Omni API")


def require_key(authorization: str | None = Header(default=None)):
    if authorization != f"Bearer {API_KEY}":
        raise HTTPException(401, "Invalid or missing API key")


class ChatRequest(BaseModel):
    model: str | None = None
    messages: list[dict]
    max_tokens: int | None = None
    temperature: float | None = None  # accepted for compatibility; decoding is greedy


@app.get("/")
def root():
    return {"status": "ok", "model": SERVED_NAME}


@app.get("/v1/models", dependencies=[Depends(require_key)])
def list_models():
    return {
        "object": "list",
        "data": [{"id": SERVED_NAME, "object": "model", "owned_by": "liquidai"}],
    }


@app.post("/v1/chat/completions", dependencies=[Depends(require_key)])
def chat_completions(req: ChatRequest):
    if not req.messages:
        raise HTTPException(400, "messages must not be empty")

    max_new = min(req.max_tokens or DEFAULT_MAX_TOKENS, HARD_MAX_TOKENS)

    with tempfile.TemporaryDirectory() as work:
        messages, has_video = convert_messages(req.messages, work)
        try:
            text = run_model(messages, has_video, max_new)
        except HTTPException:
            raise
        except Exception:
            log.exception("Generation failed")
            raise HTTPException(500, "Model failed to generate a response")

    return {
        "id": f"chatcmpl-{uuid.uuid4().hex}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": SERVED_NAME,
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": text or "(empty response)"},
                "finish_reason": "stop",
            }
        ],
    }


# ---------- Gradio page (for humans) ----------

demo = gr.Blocks(title="Omni API")
with demo:
    gr.Markdown(
        f"## Omni API\n"
        f"Model: `{MODEL_ID}`\n\n"
        f"OpenAI-compatible endpoint: `POST /v1/chat/completions` "
        f"(requires `Authorization: Bearer <key>`)."
    )

app = gr.mount_gradio_app(app, demo, path="/ui")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)