import os
import base64
from openai import OpenAI
from dotenv import load_dotenv
from schemas import ExtractionResponse

load_dotenv()

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

async def extract_receipt_data(image_content: bytes, user_name: str, department: str) -> ExtractionResponse:
    base64_image = base64.b64encode(image_content).decode("utf-8")
    
    prompt = f"You are an HR data entry assistant. Extract the requested fields from the provided medical receipt image. The employee submitting this is {user_name} from {department}. If the date is missing, look for a 'Payment Date' as a fallback."
    
    completion = client.beta.chat.completions.parse(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                    },
                ],
            }
        ],
        response_format=ExtractionResponse,
    )
    
    return completion.choices[0].message.parsed