import os
import base64
import json
from io import BytesIO
from typing import List, Optional, Union, Dict, Any

from fastapi import FastAPI, File, UploadFile, HTTPException, Depends, Security, status, Form
from fastapi.security import APIKeyHeader
from pydantic import BaseModel, Field
import requests
import uvicorn
from PIL import Image

app = FastAPI()

# Environment variables for authentication
HF_TOKEN = os.getenv("HF_TOKEN")
API_KEY = os.getenv("API_KEY")

if not HF_TOKEN:
    raise ValueError("Hugging Face API token (HF_TOKEN) is missing!")

if not API_KEY:
    raise ValueError("API key (API_KEY) is missing!")

# API Key security scheme
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=True)

# Authentication dependency
async def get_api_key(api_key: str = Security(api_key_header)):
    if api_key != API_KEY:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid API Key"
        )
    return api_key

# Models
class Message(BaseModel):
    role: str
    content: Union[str, List[Dict[str, Any]]]
    
    def to_dict(self):
        return {"role": self.role, "content": self.content}

class ChatItem(BaseModel):
    messages: List[Message]
    temperature: float = 0.7
    max_tokens: int = 1024
    top_p: float = 0.9
    frequency_penalty: float = 0.0
    presence_penalty: float = 0.0

class ImageUploadRequest(BaseModel):
    prompt: str
    system_prompt: Optional[str] = ""
    temperature: float = 0.7
    max_tokens: int = 1024
    top_p: float = 0.9

def encode_image_to_base64(image_file):
    """Convert an image file to base64 string"""
    img = Image.open(image_file)
    buffer = BytesIO()
    img.save(buffer, format=img.format or "JPEG")
    img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return img_str

def generate_with_phi_vision(messages_data):
    """Generate response using Phi-3.5-Vision model"""
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/json"
    }
    
    api_url = "https://api-inference.huggingface.co/models/microsoft/Phi-3.5-vision-instruct"
    
    # Convert pydantic models to dictionaries if necessary
    message_dicts = []
    for msg in messages_data:
        if hasattr(msg, 'to_dict'):
            message_dicts.append(msg.to_dict())
        else:
            message_dicts.append(msg)
    
    payload = {"inputs": message_dicts}
    
    try:
        response = requests.post(api_url, json=payload, headers=headers)
        
        if response.status_code != 200:
            # Handle the "model is loading" scenario
            if "is currently loading" in response.text:
                error_data = response.json()
                raise HTTPException(
                    status_code=503,
                    detail=f"Model is currently loading. Estimated time: {error_data.get('estimated_time', 'unknown')} seconds"
                )
            else:
                raise HTTPException(
                    status_code=response.status_code,
                    detail=f"Hugging Face API error: {response.text}"
                )
        
        return response.json()
    except requests.exceptions.RequestException as e:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail=f"Error connecting to Hugging Face API: {str(e)}"
        )

@app.post("/generate/")
async def generate_text(item: ChatItem, api_key: str = Depends(get_api_key)):
    """Generate text response using Phi-3.5-Vision model"""
    try:
        # Convert Message objects to dictionaries
        messages_dict = [msg.to_dict() for msg in item.messages]
        response = generate_with_phi_vision(messages_dict)
        return {"response": response}
    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/generate_with_image/")
async def generate_with_image(
    prompt: str = Form(...),
    system_prompt: str = Form(""),
    image: UploadFile = File(...),
    temperature: float = Form(0.7),
    max_tokens: int = Form(1024),
    top_p: float = Form(0.9),
    api_key: str = Depends(get_api_key)
):
    """Generate response from Phi-3.5-Vision based on an image and text prompt"""
    try:
        # Convert image to base64
        img_base64 = encode_image_to_base64(image.file)
        
        # Determine content type
        content_type = image.content_type or "image/jpeg"
        
        # Construct the messages in the format expected by the model
        messages = [
            {
                "role": "system",
                "content": system_prompt if system_prompt else "You are a helpful vision assistant."
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:{content_type};base64,{img_base64}"}}
                ]
            }
        ]
        
        # Call the model with retry logic
        response = generate_with_phi_vision(messages)
        return {"response": response}
    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/test/")
async def test_api(api_key: str = Depends(get_api_key)):
    """Test if the API is working"""
    return {"message": "API is working!"}

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