from openai import OpenAI
import openai
import base64
from typing import Dict, List
import mimetypes

def encode_image(image_path: str) -> str:
    """
    Encode image file to base64 string.
    
    Args:
        image_path: Path to the image file
        
    Returns:
        Base64 encoded string of the image
    """
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def create_chat_messages(image_path: str, prompt: str) -> List[Dict]:
    """
    Create chat messages with image and prompt.
    
    Args:
        image_path: Path to the image file
        prompt: Text prompt for the model
        
    Returns:
        List of message dictionaries
    """
    mime, _ = mimetypes.guess_type(image_path)
    if mime is None:
        mime = "image/jpeg"  # fallback
    return [
        {"role": "system", "content": ""},
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:{mime};base64,{encode_image(image_path)}"
                    }
                },
                {"type": "text", "text": prompt}
            ]
        }
    ]

PROMPT = '''Extract the text, then translate the text content into Vietnamese. If it is a document, ignore the header and footer. Formulas should be represented in LaTeX format, and tables should be represented in HTML format.'''
URL = "page1.PNG"

try:
    client = OpenAI(
        api_key="VLMcafebabe",
        base_url="http://192.168.10.42:8081/v1",
        timeout=3600
    )

    messages = create_chat_messages(URL, PROMPT)

    response = client.chat.completions.create(
        model="tencent/HunyuanOCR",
        messages=messages,
        temperature=0.0,
        stream=False,
        extra_body={
            "top_k": 1,
            "repetition_penalty": 1.0
        },
    )
    print(f"Generated text: {response.choices[0].message.content}")


except openai.APIStatusError as e:
    # Đây là phần quan trọng nhất để biết lỗi gì
    print("--- LỖI TỪ SERVER VLLM ---")
    print(f"Status Code: {e.status_code}")
    print(f"Error Message: {e.message[:200]}")
    print(f"Full Response: {e.response}") # In ra chi tiết JSON trả về
except Exception as e:
    print(f"Lỗi khác: {e}")