import os
import base64
import io
import logging
from openai import AsyncOpenAI
from dotenv import load_dotenv
from PIL import Image
from backend.schemas import ExtractionResponse

load_dotenv()

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

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

MAX_IMAGE_SIZE = (2000, 2000)
IMAGE_QUALITY = 85

def compress_image(image_content: bytes) -> bytes:
    """Resizes and compresses the image to reduce upload size."""
    img = Image.open(io.BytesIO(image_content))
    
    # Convert to RGB if necessary (to save as JPEG)
    if img.mode in ("RGBA", "P"):
        img = img.convert("RGB")
    
    # Resize if larger than max dimensions
    img.thumbnail(MAX_IMAGE_SIZE, Image.Resampling.LANCZOS)
    
    # Save to bytes
    output_buffer = io.BytesIO()
    img.save(output_buffer, format="JPEG", quality=IMAGE_QUALITY, optimize=True)
    return output_buffer.getvalue()

async def extract_receipt_data(image_content: bytes, user_name: str, department: str) -> ExtractionResponse:
    # 1. Compress Image
    compressed_content = compress_image(image_content)
    base64_image = base64.b64encode(compressed_content).decode("utf-8")
    
    # 2. Refined Prompt
    prompt = (
        f"You are an HR data entry assistant helping an employee in Malaysia. "
        f"Extract the requested fields from the provided medical receipt image. "
        f"The employee submitting this is {user_name} from {department}. "
        f"IMPORTANT: The currency is always Ringgit Malaysia (MYR). Extract the total amount and assume it is in MYR. "
        f"If the date is missing, look for a 'Payment Date' as a fallback. "
        f"Analyze the receipt for authenticity. If the total amount appears altered or if the provider name is missing, "
        f"set `needs_manual_review` to `true` and provide a low `confidence_score`."
    )
    
    # 3. Async Extraction
    completion = await 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,
    )
    
    result = completion.choices[0].message.parsed
    
    # 4. Logging for Demo
    if result:
        logger.info(f"Extraction complete for {user_name}. Confidence Score: {result.confidence_score}")
        if result.needs_manual_review:
            logger.warning(f"Manual review required for receipt submitted by {user_name}")
            
    return result