# Install system dependencies first
!apt-get update
!apt-get install -y \
    chromium-browser \
    chromium-chromedriver \
    libnss3 \
    libxss1 \
    libatk-bridge2.0-0 \
    libgtk-3-0 \
    libgbm-dev

import nest_asyncio
nest_asyncio.apply()

import os
import gradio as gr
import logging
from roboflow import Roboflow
from PIL import Image, ImageDraw
import cv2
import numpy as np
from math import atan2, degrees
import asyncio
from pyppeteer import launch

# Configure logging
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler("debug.log"),
        logging.StreamHandler()
    ]
)

# Roboflow and model configuration
ROBOFLOW_API_KEY = "KUP9w62eUcD5PrrRMJsV"  # Replace with your API key if needed
PROJECT_NAME = "model_verification_project"
VERSION_NUMBER = 2

# ----------------------------
# Asynchronous function to generate handwriting image via Pyppeteer
# ----------------------------
async def _generate_handwriting_image(text_prompt, screenshot_path):
    try:
        # Launch Chromium with the correct path
        browser = await launch(
            headless=True,
            executablePath="/usr/bin/chromium-browser",  # Explicit path to Chromium
            args=[
                '--no-sandbox',
                '--disable-setuid-sandbox',
                '--disable-dev-shm-usage',
                '--disable-gpu',
                '--single-process',
                '--no-zygote',
                '--window-size=1920,1080'
            ]
        )
        page = await browser.newPage()
        
        # Navigate to Calligraphr
        await page.goto('https://www.calligraphr.com/en/font/', {
            'waitUntil': 'networkidle2',
            'timeout': 60000  # 60 seconds timeout
        })
        
        # Wait for the text input field
        await page.waitForSelector('#text-input', {'timeout': 30000})
        
        # Type the text prompt
        await page.type('#text-input', text_prompt)
        
        # Wait for rendering
        await asyncio.sleep(5)
        
        # Take a screenshot
        await page.screenshot({
            'path': screenshot_path,
            'clip': {'x': 100, 'y': 200, 'width': 600, 'height': 150}
        })
        return screenshot_path

    except Exception as e:
        logging.error(f"Pyppeteer error: {str(e)}")
        return None

    finally:
        # Close the browser
        if 'browser' in locals():
            await browser.close()

def generate_handwriting_image(text_prompt, screenshot_path="/tmp/handwriting.png"):
    """
    Synchronous wrapper around the async Pyppeteer call.
    """
    try:
        loop = asyncio.get_event_loop()
        result = loop.run_until_complete(_generate_handwriting_image(text_prompt, screenshot_path))
        return result
    except Exception as e:
        logging.error(f"Error generating handwriting image: {e}")
        return None

# ----------------------------
# Main processing function
# ----------------------------
def process_image(image, text):
    try:
        # Initialize Roboflow
        rf = Roboflow(api_key=ROBOFLOW_API_KEY)
        logging.debug("Initialized Roboflow API.")
        project = rf.workspace().project(PROJECT_NAME)
        logging.debug("Accessed project in Roboflow.")
        model = project.version(VERSION_NUMBER).model
        logging.debug("Loaded model from Roboflow.")

        # Save input image temporarily
        input_image_path = "/tmp/input_image.jpg"
        image.save(input_image_path)
        logging.debug(f"Input image saved to {input_image_path}.")

        # Perform inference
        logging.debug("Performing inference on the image...")
        prediction = model.predict(input_image_path, confidence=70, overlap=50).json()
        logging.debug(f"Inference result: {prediction}")

        pil_image = image.convert("RGBA")
        logging.debug("Converted image to RGBA mode.")

        # Iterate over detected objects (assumed white paper)
        for obj in prediction['predictions']:
            # Paper dimensions
            white_paper_width = obj['width']
            white_paper_height = obj['height']

            # Padding
            padding_x = int(white_paper_width * 0.1)
            padding_y = int(white_paper_height * 0.1)
            box_width = white_paper_width - 2 * padding_x
            box_height = white_paper_height - 2 * padding_y
            logging.debug(f"Padded white paper dimensions: width={box_width}, height={box_height}.")

            # Calculate padded coordinates
            x1_padded = int(obj['x'] - white_paper_width / 2 + padding_x)
            y1_padded = int(obj['y'] - white_paper_height / 2 + padding_y)
            x2_padded = int(obj['x'] + white_paper_width / 2 - padding_x)
            y2_padded = int(obj['y'] + white_paper_height / 2 - padding_y)

            # Detect paper angle
            angle = detect_paper_angle(np.array(image), (x1_padded, y1_padded, x2_padded, y2_padded))
            logging.debug(f"Detected paper angle: {angle} degrees.")

            # (Optional) debug bounding box
            debug_layer = pil_image.copy()
            debug_draw = ImageDraw.Draw(debug_layer)
            debug_draw.rectangle([(x1_padded, y1_padded), (x2_padded, y2_padded)], outline="red", width=3)
            debug_layer.save("/tmp/debug_bounding_box.png")
            logging.debug("Saved bounding box debug image to /tmp/debug_bounding_box.png.")

            # Generate handwriting image
            handwriting_path = generate_handwriting_image(text, "/tmp/handwriting.png")
            if not handwriting_path:
                logging.error("Handwriting image generation failed.")
                continue

            handwriting_img = Image.open(handwriting_path).convert("RGBA")
            handwriting_img = handwriting_img.resize((box_width, box_height), Image.ANTIALIAS)
            rotated_handwriting = handwriting_img.rotate(-angle, resample=Image.BICUBIC, expand=True)

            # Composite the handwriting
            text_layer = Image.new("RGBA", pil_image.size, (255, 255, 255, 0))
            paste_x = int(obj['x'] - rotated_handwriting.size[0] / 2)
            paste_y = int(obj['y'] - rotated_handwriting.size[1] / 2)
            text_layer.paste(rotated_handwriting, (paste_x, paste_y), rotated_handwriting)
            pil_image = Image.alpha_composite(pil_image, text_layer)
            logging.debug("Handwriting layer composited onto the original image.")

        # Save output
        output_image_path = "/tmp/output_image.png"
        pil_image.convert("RGB").save(output_image_path)
        logging.debug(f"Output image saved to {output_image_path}.")
        return output_image_path

    except Exception as e:
        logging.error(f"Error during image processing: {e}")
        return None

# ----------------------------
# Gradio interface
# ----------------------------
interface = gr.Interface(
    fn=gradio_inference,
    inputs=[
        gr.Image(type="pil", label="Upload an Image"),
        gr.Textbox(label="Enter Text to Overlay")
    ],
    outputs=[
        gr.Image(label="Processed Image Preview"),
        gr.File(label="Download Processed Image"),
        gr.Textbox(label="Status")
    ],
    title="Roboflow Detection with Handwriting Overlay",
    description="Upload an image and enter text to overlay. The Roboflow model detects the white paper area, and a handwriting image is generated via Calligraphr using Pyppeteer. The output image is composited accordingly.",
    allow_flagging="never"
)

if __name__ == "__main__":
    interface.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        enable_queue=True
    )