from fasthtml.common import *
from monsterui.all import *
import qrcode
import io
import base64
from PIL import Image

def generate_qr_code(url, box_size=10, border=2, with_logo=False):
    c = qrcode.constants
    error_correction = c.ERROR_CORRECT_H if with_logo else c.ERROR_CORRECT_L
    qr = qrcode.QRCode(
        version=1,
        error_correction=error_correction,
        box_size=box_size,
        border=border,
    )
    qr.add_data(url)
    qr.make(fit=True)
    
    img = qr.make_image(fill_color="black", back_color="white")
    return img

def qr_to_base64(qr_img, format="PNG"):
    img_byte_arr = io.BytesIO()
    qr_img.save(img_byte_arr, format=format)
    encoded = base64.b64encode(img_byte_arr.getvalue()).decode('ascii')
    return f"data:image/{format.lower()};base64,{encoded}"

def add_logo_to_qr(qr_img, logo, logo_size_percentage=25):
    try:
        contents = logo.file.read()
        image_stream = io.BytesIO(contents)
        logo_img = Image.open(image_stream)
        logo.file.seek(0)
        
        # Calculate the size of the logo based on the QR code size
        qr_width, qr_height = qr_img.size
        logo_max_size = min(qr_width, qr_height) * logo_size_percentage // 100
        
        # Resize the logo to fit within the QR code
        logo_width, logo_height = logo_img.size
        if logo_width > logo_max_size or logo_height > logo_max_size:
            if logo_width > logo_height:
                logo_height = int(logo_height * logo_max_size / logo_width)
                logo_width = logo_max_size
            else:
                logo_width = int(logo_width * logo_max_size / logo_height)
                logo_height = logo_max_size
            logo_img = logo_img.resize((logo_width, logo_height), Image.LANCZOS)
        
        # Calculate position to place the logo (center)
        position = ((qr_width - logo_width) // 2, (qr_height - logo_height) // 2)
        
        # Create a new image with the same size as the QR code and paste the logo
        qr_with_logo = qr_img.copy().convert("RGBA")
        qr_with_logo.paste(logo_img, position, logo_img)
        
        return qr_with_logo.convert("RGB")
    except Exception as e:
        print(f"Error adding logo: {e}")
        return qr_img

# Create FastHTML app
app, rt = fast_app(hdrs = Theme.green.headers())

@rt('/')
def get():
    # Create a form for URL input
    url = LabelInput(label="URL", placeholder="Enter URL here", id="url")
    
    box_size = LabelRange(label='Box size', value='10', min=5, max=20, step=1, name="box_size")
    border_width = LabelRange(label='Border width', value='4', min=0, max=10, step=1, name="border_width")
    
    logo_upload = DivVStacked(
        # Label("Logo (Optional)", for_="logo"),
        # Upload("Logo (Optional)", id="logo"),
        LabelInput("Logo (Optional)", type="file", name="logo", id="logo", accept="image/*"),
        P("For best results, use a transparent PNG", cls="text-sm")
    )
    
    logo_size = LabelRange(label='Logo Size (%)', value='25', min=5, max=30, step=5, name="logo_size")
    
    controls = DivVStacked(DivHStacked(box_size, border_width), logo_upload, logo_size)
    
    generate_button = Button((UkIcon("rocket", cls='mr-2'), "Generate"), cls=ButtonT.primary)
        
    form = Form(
        Div(url, DivCentered(controls), DivCentered(generate_button), cls='space-y-5'),
        cls='space-x-5',      
        hx_post="/generate",
        hx_target="#qr-result",
        hx_swap="innerHTML"
    )
    
    # Container for QR code result
    result_div = Div(id="qr-result")
    
    # Main page content
    return Title("QR Code Generator"), Main(
                DivVStacked(
                    H1("QR Code Generator"),
                    Card(
                        DivHStacked(
                            result_div,
                            form,
                            cls='space-x-4'
                        ),
                        cls='space-y-4'            
                    ),
                    cls='py-6 space-y-4'
                )
            )

@rt('/generate')
def post(url: str, 
         box_size: int = 10, 
         border_width: int = 4, 
         logo_size: int | None = 25, 
         logo: any = None):
    if not url:
        return Alert(DivLAligned(UkIcon('triangle-alert'), 
                                 P("Please enter a URL.")), 
                     cls=AlertT.error)
    
    # Generate QR code
    if not logo:
        qr_img = generate_qr_code(url, box_size=box_size, border=border_width)
        img_base64 = qr_to_base64(qr_img)
        
    else:
        # Generate QR code with logo
        qr_img = generate_qr_code(url, box_size=box_size, border=border_width, with_logo=True)
        qr_img_with_logo = add_logo_to_qr(qr_img, logo, logo_size_percentage=logo_size)
        
        # Convert to base64
        img_base64 = qr_to_base64(qr_img_with_logo)
    
    # Return the QR code image with a download link
    return DivVStacked(
        Card(Img(src=img_base64, alt="QR Code", cls="qr-image")),
        Button((UkIcon("download", cls='mr-2'), "Download"), 
          href=img_base64, 
          download=f"qr-{url.replace('://', '-').replace('/', '-')}.png", 
          cls=ButtonT.secondary)
    )

serve()