import os
from pyppeteer import launch
from PIL import Image
import json
import random
from pydantic import BaseModel
from fastapi import FastAPI, HTTPException
import io
from fastapi.responses import StreamingResponse
app = FastAPI()

class HtmlInput(BaseModel):
    message: str


async def convert_html_to_image(html, output_file, width=650, height=800, device_scale_factor=2):
    try:
        browser = await launch({
            'headless': True,
            'args': ['--no-sandbox', '--disable-setuid-sandbox'],
            'timeout': 60000  # Increase timeout if needed
        })
        page = await browser.newPage()
        await page.setViewport({'width': width, 'height': height, 'deviceScaleFactor': device_scale_factor})
        await page.setContent(html)
        await page.screenshot({'path': output_file})
        await browser.close()
        return output_file  # Return the path of the captured image
    except Exception as e:
        print(f"Error in convert_html_to_image: {e}")
        raise

async def optimize_image(input_file, output_file, target_size_kb=700):
    # Open the image using PIL
    with Image.open(input_file) as img:
        # Calculate the initial quality to achieve the target size
        quality = 95
        temp_output_file = output_file + ".temp.jpg"
        img.save(temp_output_file, format='JPEG', quality=quality, optimize=True)
        while os.path.getsize(temp_output_file) > target_size_kb * 1024 and quality > 0:
            # Reduce the quality and save to temporary file
            quality -= 5
            img.save(temp_output_file, format='JPEG', quality=quality, optimize=True)
        
        # Save the optimized image to the output file
        os.replace(temp_output_file, output_file)

async def mainFunction(html_content:str):
    output_image_file =  "image.jpg"
    captured_image_path = await convert_html_to_image(html_content, output_image_file)
        
        # Optimize the image to be less than 200KB
    await optimize_image(captured_image_path, output_image_file)

    return output_image_file

# Run the asynchronous main function
@app.post("/convert-html-to-image/")
async def convert_html_to_image_endpoint(html_content:HtmlInput):
    image_path = await mainFunction(html_content.message)
    
    # Check if image was generated
    if not image_path or not os.path.exists(image_path):
        raise HTTPException(status_code=500, detail="Image generation failed.")

    # Open the image file for streaming
    with open(image_path, "rb") as img_file:
        buffer = io.BytesIO(img_file.read())

    buffer.seek(0)

    # Return the image as a downloadable file
    return StreamingResponse(buffer, media_type="image/jpeg", headers={
        "Content-Disposition": f"attachment; filename={os.path.basename(image_path)}"
    })