import os
from pyppeteer import launch
from PIL import Image
import json
import random
from pydantic import BaseModel
from fastapi import FastAPI
from jinja2 import Environment, FileSystemLoader
import io
app = FastAPI()

class HtmlInput(BaseModel):
    message: str



async def convert_html_to_image(html, output_file, width=612, height=800, device_scale_factor=2):
    browser = await launch(
        headless=True,
        args=[
            '--no-sandbox',
            '--disable-setuid-sandbox',
            '--disable-gpu',
            '--disable-dev-shm-usage',  # Try to disable /dev/shm usage
            '--disable-software-rasterizer',
            '--disable-accelerated-2d-canvas',
            '--disable-infobars',
            '--disable-extensions',
            '--disable-features=NetworkService',
            '--disable-features=VizDisplayCompositor',
            '--window-size=612,800'
        ]
    )
    page = await browser.newPage()
    await page.setViewport({'width': width, 'height': height, 'deviceScaleFactor': device_scale_factor})
    await page.setContent(html)
    # Capture only the viewport without extra white space
    await page.screenshot({'path': output_file})
    await browser.close()
    return output_file  # Return the path of the captured image

# Function to tilt the content randomly
def tilt_content(html_content):
    tilt_angle = random.randint(-3, 2)  # Adjust the range of tilting angle as needed
    tilted_content = f'<div style="transform: rotate({tilt_angle}deg);">{html_content}</div>'
    return tilted_content

def apply_random_css(html_content):
    selected_css = random.choice([".txORqLQ { font-size: 13px; font-weight: bold; fill: #000080; color:#000080; }",
                                  ".txORqLQ { font-size: 12px; font-weight: bold; fill: #000080; color:#000080; }",
                                  ".txORqLQ { font: bold 10px 'Courier New',Courier,monospace; }",
                                  ".txORqLQ { font: bold 10.5px 'Courier New',Courier,monospace; }"])
    offset_content=html_content
    offset_content = f'<style>{selected_css}</style>' + offset_content
    return offset_content

def apply_random_css_positioning(html_content):
    selected_css = random.choice([".address { position: relative;bottom: 48.6rem;left: 16.8rem;width: 11rem; margin-top:0.3rem}",
                                  ".address { position: relative;bottom: 47.8rem;left: 10.9rem;width: 30rem; margin-top:0.3rem}"
                                  ])
    offset_content=html_content
    offset_content = f'<style>{selected_css}</style>' + offset_content
    return offset_content

def apply_position_offset(html_content):
    # Random offsets in x and y directions
    offset_x = random.randint(-20, 20)  # Adjust the range of x offset as needed
    offset_y = random.randint(-20, 20)  # Adjust the range of y offset as needed
    
    # Apply the offsets to the content
    offset_content = f'<div style="position: relative; left: {offset_x}px; top: {offset_y}px;">{html_content}</div>'
    return offset_content

async def optimize_image(input_file, output_file, target_size_kb=200):
    # 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):
    print("in main function")
    # Create img folder if not exists
    if not os.path.exists('img'):
        os.makedirs('img')
    if not os.path.exists('json'):
        os.makedirs('json')

    # Write the HTML string to the file
    html_file = io.StringIO(html_content)
    print("html file created")
    template_loader = FileSystemLoader(searchpath=".")
    template_env = Environment(loader=template_loader)
    template = template_env.from_string(html_file.getvalue())
    print("Template rendered successfully")

    # Load JSON data from files
    with open("modified_data.json", "r", encoding="utf-8") as universal_file:
        universal_json_data = json.load(universal_file)
    

    # Define the start number
    start_number = 17800

    # Loop over each record in the universal JSON
    for index, record in enumerate(universal_json_data):
        # Increment index by start_number
        index += start_number
        
        # Generate json file

        # Load the template

        json_output_file = os.path.join('json', f"{index:05d}.json")
        with open(json_output_file, "w", encoding="utf-8") as json_file:
            json.dump(record, json_file, indent=4, ensure_ascii=False)
        html_output = template.render(record)
        html_output= apply_random_css(html_output)
        html_output=apply_random_css_positioning(html_output)


        # Apply random position offset
        # html_output = apply_position_offset(html_output)

        # Tilt the content randomly
        # html_output = tilt_content(html_output)

               
        # Convert HTML to image
        output_image_file = os.path.join('img', f"{index:05d}.jpg")
        captured_image_path = await convert_html_to_image(html_output, output_image_file)
        
        # Optimize the image to be less than 200KB
        await optimize_image(captured_image_path, output_image_file)

# Run the asynchronous main function
@app.post("/convert-html-to-image/")
async def convert_html_to_image_endpoint(html_content:HtmlInput):
    await mainFunction(html_content.message)