import os
from pyppeteer import launch
from PIL import Image
import json
import random
from pydantic import BaseModel
from fastapi import FastAPI, HTTPException
from jinja2 import Environment, FileSystemLoader
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)}"
    })