from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image
import io
import json
from utils import get_text
from json_flatten import flatten

app = FastAPI(
    title="DOCUMANTICAI API",
    description="""
    This API allows you to upload an image and get a formatted response with details and image information.
    """
)

@app.post("/upload")
async def upload_image(fields:str, model:str, file: UploadFile = File(...)):
    """
    ### Endpoint Description:
    Extract form data from an uploaded image and return the extracted data in JSON format.

    #### Request Parameters:
    - `file`: The image file to extract data from. (Required)

    #### Response:
    ### Notes:
    - The image should be in a supported format (e.g., PNG, JPEG).
    - The data extracted will vary depending on the image content.
    """
    
    try:
        # Load the uploaded image
        image = Image.open(io.BytesIO(await file.read()))
        
        # Example: Get image details
        image_details = {
            "filename": file.filename,
            "format": image.format,
            "size": image.size,  # (width, height)
            "mode": image.mode
        }
        response = get_text(image,image_details['filename'], model, fields)
        # Step 1: Convert the escaped JSON string to a proper dictionary        
        # Step 2: Convert the response to a proper dictionary
        response = json.loads(response)
        
        # Step 3: Convert fields and values into key-value pairs
        if 'fields' in response and 'values' in response:
            response = dict(zip(response['fields'], response['values']))
        
        # response flattening
        response = flatten(response)

        # Process image (example: return metadata)
        return JSONResponse(content={"response": response, "details": image_details})
    
    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=400)
    
@app.post("/list_models")
async def list_models():
    """
    ### Endpoint Description:
    List available models for text generation.

    #### Response:
    - A list of available models for text generation.
    """
    return JSONResponse(content={"models": ["gpt-4o-mini", "gpt-4o", "deepseek-chat", "claude-3-5-sonnet-20241022", "llama_llm_d","llama_llm_o"]})