import os
import io
import time
import sys
import subprocess
import requests
from PIL import Image, ImageSequence
import gradio as gr

# llama-cpp-python import
try:
    from llama_cpp import Llama
except Exception as e:
    raise RuntimeError("llama-cpp-python import failed: " + str(e))

MODEL_DIR = "model"
MODEL_MAIN = os.path.join(MODEL_DIR, "llama-joycaption-q4_k_m.gguf")
MODEL_FALLBACK = os.path.join(MODEL_DIR, "llama-joycaption-q4_k_s.gguf")

# Candidate direct-download URLs (try in order)
CANDIDATES = [
    # Primary Q4_K_M (Jasaga then mradermacher)
    ("https://huggingface.co/Jasaga7818/llama-joycaption-beta-one-hf-llava-Q4_K_M-GGUF/resolve/main/llama-joycaption-beta-one-hf-llava-q4_k_m.gguf",
     MODEL_MAIN),
    ("https://huggingface.co/mradermacher/llama-joycaption-beta-one-hf-llava-GGUF/resolve/main/llama-joycaption-beta-one-hf-llava-q4_k_m.gguf",
     MODEL_MAIN),
    # Fallback Q4_K_S (mradermacher / Jasaga)
    ("https://huggingface.co/mradermacher/llama-joycaption-beta-one-hf-llava-GGUF/resolve/main/llama-joycaption-beta-one-hf-llava-q4_k_s.gguf",
     MODEL_FALLBACK),
    ("https://huggingface.co/Jasaga7818/llama-joycaption-beta-one-hf-llava-Q4_K_M-GGUF/resolve/main/llama-joycaption-beta-one-hf-llava-q4_k_s.gguf",
     MODEL_FALLBACK),
]

def download_curl(url: str, path: str) -> bool:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    try:
        subprocess.check_call(["curl", "-L", "-C", "-", "-o", path, url])
        return True
    except Exception:
        try:
            if os.path.exists(path):
                os.remove(path)
        except Exception:
            pass
        return False

def is_valid_gguf(path: str) -> bool:
    try:
        with open(path, "rb") as f:
            head = f.read(8)
        return head.startswith(b"GGUF")
    except Exception:
        return False

def ensure_models_downloaded():
    # If main present and valid, done.
    if os.path.exists(MODEL_MAIN) and is_valid_gguf(MODEL_MAIN):
        sys.stderr.write(f"Found valid main model: {MODEL_MAIN}\n")
        return
    # If fallback present and valid, done.
    if os.path.exists(MODEL_FALLBACK) and is_valid_gguf(MODEL_FALLBACK):
        sys.stderr.write(f"Found valid fallback model: {MODEL_FALLBACK}\n")
        return

    sys.stderr.write("Model(s) missing or invalid; attempting downloads...\n")
    for url, dest in CANDIDATES:
        sys.stderr.write(f"Downloading {url} -> {dest}\n")
        ok = download_curl(url, dest)
        if not ok:
            sys.stderr.write(f"Download failed for {url}\n")
            continue
        if is_valid_gguf(dest):
            sys.stderr.write(f"Downloaded and verified GGUF at {dest}\n")
            # if we downloaded fallback but main missing, don't copy; we'll try to load fallback later
            if dest == MODEL_MAIN:
                return
            # if dest is fallback, still continue loop to attempt main first (if available)
        else:
            sys.stderr.write(f"Downloaded file at {dest} is not a valid GGUF (header mismatch). Removing.\n")
            try:
                os.remove(dest)
            except Exception:
                pass

    sys.stderr.write("Download attempts finished.\n")

def try_load_model(path: str, n_ctx: int = 2048, n_threads: int = 2):
    try:
        sys.stderr.write(f"Initializing Llama with model {path}...\n")
        llm = Llama(model_path=path, n_ctx=n_ctx, n_threads=n_threads)
        sys.stderr.write("Model loaded successfully.\n")
        return llm
    except Exception as e:
        sys.stderr.write(f"Failed to load model {path}: {e}\n")
        return None

# Ensure at least one model file is present (download if needed)
ensure_models_downloaded()

# Prefer main, then fallback
model_to_try = None
if os.path.exists(MODEL_MAIN) and is_valid_gguf(MODEL_MAIN):
    model_to_try = MODEL_MAIN
elif os.path.exists(MODEL_FALLBACK) and is_valid_gguf(MODEL_FALLBACK):
    model_to_try = MODEL_FALLBACK
else:
    # attempt to download again and pick whatever exists
    ensure_models_downloaded()
    if os.path.exists(MODEL_MAIN) and is_valid_gguf(MODEL_MAIN):
        model_to_try = MODEL_MAIN
    elif os.path.exists(MODEL_FALLBACK) and is_valid_gguf(MODEL_FALLBACK):
        model_to_try = MODEL_FALLBACK

if model_to_try is None:
    raise FileNotFoundError("No valid GGUF model found. Place a compatible GGUF under model/ with filename\n"
                            "model/llama-joycaption-q4_k_m.gguf or model/llama-joycaption-q4_k_s.gguf.")

# Attempt to load chosen model; if load fails for magic/version, try fallback (if different)
llm = try_load_model(model_to_try, n_ctx=2048, n_threads=2)
if llm is None and model_to_try == MODEL_MAIN and os.path.exists(MODEL_FALLBACK) and is_valid_gguf(MODEL_FALLBACK):
    sys.stderr.write("Primary model failed to load; attempting fallback model.\n")
    llm = try_load_model(MODEL_FALLBACK, n_ctx=2048, n_threads=2)

if llm is None:
    # Provide clear diagnostic and exit
    sys.stderr.write("\nERROR: All model load attempts failed. Likely causes:\n"
                     " - The GGUF uses a newer GGUF version not supported by the installed llama.cpp/llama-cpp-python.\n"
                     " - The file is corrupted despite the header check.\n\n"
                     "Recommended fixes:\n"
                     " - Install a newer llama.cpp/llama-cpp-python built from main/master (supports newer GGUF versions).\n"
                     " - Or place a known-compatible GGUF (Q4_K_S from mradermacher or older GGUF) at model/llama-joycaption-q4_k_m.gguf\n"
                     " - To inspect the header run: hexdump -n4 model/llama-joycaption-q4_k_m.gguf\n")
    raise RuntimeError("Model load failed for all candidates.")

def download_bytes(url: str, timeout: int = 30) -> bytes:
    with requests.get(url, stream=True, timeout=timeout) as r:
        r.raise_for_status()
        return r.content

def load_first_frame_from_bytes(raw: bytes):
    img = Image.open(io.BytesIO(raw))
    if getattr(img, "is_animated", False):
        img = next(ImageSequence.Iterator(img))
    if img.mode != "RGB":
        img = img.convert("RGB")
    return img

def make_prompt_for_image(image_path: str, user_prompt: str = "Describe the image."):
    # JoyCaption-style multimodal GGUFs accept <img>{path}</img>
    return f"<img>{image_path}</img>\nUser: {user_prompt}\nAssistant:"

def generate_caption_from_url(url: str, prompt: str = "Describe the image."):
    if not url:
        return "No URL provided."
    try:
        raw = download_bytes(url)
    except Exception as e:
        return f"Download error: {e}"
    try:
        img = load_first_frame_from_bytes(raw)
    except Exception as e:
        return f"Image processing error: {e}"

    tmp_dir = "/tmp/joycap"
    os.makedirs(tmp_dir, exist_ok=True)
    ts = int(time.time() * 1000)
    tmp_path = os.path.join(tmp_dir, f"{ts}.jpg")
    try:
        img.save(tmp_path, format="JPEG", quality=85)
    except Exception as e:
        return f"Failed to save temp image: {e}"

    prompt_full = make_prompt_for_image(tmp_path, prompt)
    try:
        resp = llm.create(
            prompt=prompt_full,
            max_tokens=256,
            temperature=0.2,
            top_p=0.95,
            stop=["User:", "Assistant:"],
        )
        text = resp.get("choices", [{}])[0].get("text", "").strip()
        return text or "No caption generated."
    except Exception as e:
        return f"Inference error: {e}"
    finally:
        try:
            os.remove(tmp_path)
        except Exception:
            pass

iface = gr.Interface(
    fn=generate_caption_from_url,
    inputs=[
        gr.Textbox(label="Image URL", placeholder="https://example.com/photo.jpg"),
        gr.Textbox(label="Prompt (optional)", value="Describe the image."),
    ],
    outputs=gr.Textbox(label="Generated caption"),
    title="JoyCaption GGUF (Q4_K)",
    description="Runs a quantized JoyCaption GGUF locally via llama.cpp (no external API).",
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)