import os
from fastapi import FastAPI, HTTPException, Header
from pydantic import BaseModel
from groq import Groq

app = FastAPI()

# Get the key from environment variables
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
client = Groq(api_key=GROQ_API_KEY)

class ImageRequest(BaseModel):
    base64_image: str

@app.post("/explain-chart")
async def explain_chart(request: ImageRequest, authorization: str = Header(None)):
    
    if not authorization or authorization != f"Bearer {GROQ_API_KEY}":
        raise HTTPException(status_code=401, detail="Unauthorized")

    try:
        completion = client.chat.completions.create(
            # Using the vision-capable model
            model="llama-3.2-11b-vision-preview", 
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "Explain this IPL data chart. What are the key insights?"},
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:image/png;base64,{request.base64_image}"}
                        }
                    ]
                }
            ]
        )
        return {"explanation": completion.choices[0].message.content}
    except Exception as e:
        print(f"Error detail: {e}")
        raise HTTPException(status_code=500, detail=str(e))
    
# after clicking the button in the Streamlit app, it will send a POST request to this endpoint 
# with the base64 image, and the FastAPI server will process it and return the AI-generated explanation.
# this is happening in the background, so the user experience is seamless.