import sys
import os
import json
import base64
from huggingface_hub import hf_hub_download

def download_models():
    print("Checking model files...", flush=True)
    repo_id = "xtuner/llava-phi-3-mini-gguf"
    model_file = "llava-phi-3-mini-int4.gguf"
    projector_file = "llava-phi-3-mini-mmproj-f16.gguf"
    
    # Download using hf_hub_download which is fast and supports caching
    model_path = hf_hub_download(repo_id=repo_id, filename=model_file)
    projector_path = hf_hub_download(repo_id=repo_id, filename=projector_file)
    
    print(f"Model loaded from cache: {model_path}", flush=True)
    print(f"Projector loaded from cache: {projector_path}", flush=True)
    return model_path, projector_path

def describe_image(image_path, prompt="Describe this image in detail."):
    model_path, projector_path = download_models()
    
    print("Loading llama-cpp-python model and projector...", flush=True)
    from llama_cpp import Llama
    from llama_cpp.llama_chat_format import Llava15ChatHandler
    
    chat_handler = Llava15ChatHandler(clip_model_path=projector_path)
    
    # Initialize the CPU-friendly quantized Phi-3 LLaVA model
    llm = Llama(
        model_path=model_path,
        chat_handler=chat_handler,
        n_ctx=2048,
        n_threads=4, # Use 4 CPU threads for fast inference
        verbose=False
    )
    
    print("Reading image and encoding in base64...", flush=True)
    with open(image_path, "rb") as f:
        img_bytes = f.read()
        base64_image = base64.b64encode(img_bytes).decode("utf-8")
        
    data_url = f"data:image/jpeg;base64,{base64_image}"
    
    print("Running multi-modal inference...", flush=True)
    res = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": "You are an assistant who describes images in detail."},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": data_url}}
                ]
            }
        ]
    )
    
    content = res["choices"][0]["message"]["content"]
    return content

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python describe.py <image_path> [prompt]")
        sys.exit(1)
        
    img_path = sys.argv[1]
    prompt = sys.argv[2] if len(sys.argv) > 2 else "Describe this image in detail."
    
    try:
        desc = describe_image(img_path, prompt)
        print("---RESULT_START---")
        print(desc)
        print("---RESULT_END---")
    except Exception as e:
        print(f"ERROR: {str(e)}", file=sys.stderr)
        sys.exit(1)