import os
import sqlite3
import random
import numpy as np
import gradio as gr
from PIL import Image
from pillow_heif import register_heif_opener
from skimage.color import rgb2lab
from scipy.spatial import cKDTree
from huggingface_hub import hf_hub_download

# Initialize specialized image support
register_heif_opener()
Image.MAX_IMAGE_PIXELS = None 

# ================= YOUR DATASET REPO =================
DATASET_REPO_ID = "Daksh17440/satellite_images_color_clustered"
# =====================================================

print("📥 Downloading Brains from Dataset...")
DB_PATH = hf_hub_download(repo_id=DATASET_REPO_ID, filename="mosaic_library/tiles.db", repo_type="dataset")
CENTROIDS_FILE = hf_hub_download(repo_id=DATASET_REPO_ID, filename="mosaic_library/centroids.npy", repo_type="dataset")

print("⚙️ Booting Engine & Loading KD-Tree...")
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
cursor.execute("SELECT DISTINCT bucket_id FROM tiles")
active_bucket_ids = [row[0] for row in cursor.fetchall()]
bucket_inventory = {b_id: [] for b_id in active_bucket_ids}
cursor.execute("SELECT bucket_id, filepath FROM tiles")
for b_id, filepath in cursor.fetchall():
    bucket_inventory[b_id].append(filepath)
conn.close()

all_centroids = np.load(CENTROIDS_FILE)
active_centroids = all_centroids[active_bucket_ids]
color_tree = cKDTree(active_centroids)
print("✅ Engine Ready!")

def generate_mosaic(target_image_path, mode, print_w, print_dpi, tile_in, max_w, dig_tile, deep_patch, deep_out, progress=gr.Progress()):
    if target_image_path is None:
        return None, None, "❌ Please upload an image."

    target_img = Image.open(target_image_path).convert('RGB')
    input_w, input_h = target_img.size

    # ================= THE DYNAMIC MATH ENGINE =================
    if mode == "🖨️ Type 1: Printout":
        final_w_px = print_w * print_dpi
        OUTPUT_TILE_SIZE = round(tile_in * print_dpi)
        grid_columns = final_w_px / OUTPUT_TILE_SIZE
        PATCH_SIZE = max(1, round(input_w / grid_columns))
        
    elif mode == "💻 Type 2: Finer Digital":
        grid_columns = max_w // dig_tile
        OUTPUT_TILE_SIZE = int(dig_tile)
        PATCH_SIZE = max(2, round(input_w / grid_columns))
        
    else: # 🔬 Type 3: Deep Zoom
        PATCH_SIZE = int(deep_patch)
        OUTPUT_TILE_SIZE = int(deep_out)

    actual_cols = input_w // PATCH_SIZE
    actual_rows = input_h // PATCH_SIZE
    final_w = actual_cols * OUTPUT_TILE_SIZE
    final_h = actual_rows * OUTPUT_TILE_SIZE
    
    if final_w * final_h > 150_000_000:
        target_img.thumbnail((3000, 3000))
        input_w, input_h = target_img.size
        actual_cols = input_w // PATCH_SIZE
        actual_rows = input_h // PATCH_SIZE
        final_w = actual_cols * OUTPUT_TILE_SIZE
        final_h = actual_rows * OUTPUT_TILE_SIZE

    final_canvas = Image.new('RGB', (final_w, final_h))
    target_array = np.array(target_img)

    # Error Tracking Variables
    first_error = None
    failed_tiles_count = 0

    # Error Tracking Variables
    first_error = None
    failed_tiles_count = 0
    
    # NEW: In-Memory RAM Cache for lightning-fast pasting
    tile_cache = {}
    
    # ================= STITCHING LOOP =================
    for row in range(actual_rows):
        progress(row / actual_rows, desc=f"Stitching Row {row + 1} of {actual_rows}...")
        
        for col in range(actual_cols):
            y1, y2 = row * PATCH_SIZE, (row + 1) * PATCH_SIZE
            x1, x2 = col * PATCH_SIZE, (col + 1) * PATCH_SIZE
            
            patch = target_array[y1:y2, x1:x2]
            mean_rgb = patch.mean(axis=(0, 1)).reshape(1, 1, 3).astype(np.uint8)
            mean_lab = rgb2lab(mean_rgb).reshape(3)
            
            _, active_index = color_tree.query(mean_lab)
            matched_bucket_id = active_bucket_ids[active_index]
            chosen_tile_path = random.choice(bucket_inventory[matched_bucket_id])
            
            try:
                # If we have already downloaded and resized this exact tile, use it instantly!
                if chosen_tile_path in tile_cache:
                    tile_img = tile_cache[chosen_tile_path]
                
                # Otherwise, fetch it, clean the path, resize it, and save it to the cache
                else:
                    clean_path = chosen_tile_path
                    if clean_path.startswith("./"):
                        clean_path = clean_path[2:]
                    
                    if clean_path.startswith("mosaic_library/"):
                        full_repo_path = clean_path
                    else:
                        full_repo_path = f"mosaic_library/{clean_path}"
                    
                    cached_tile_path = hf_hub_download(
                        repo_id=DATASET_REPO_ID,
                        filename=full_repo_path, 
                        repo_type="dataset"
                    )
                    
                    tile_img = Image.open(cached_tile_path).convert('RGB')
                    tile_img = tile_img.resize((OUTPUT_TILE_SIZE, OUTPUT_TILE_SIZE), Image.Resampling.LANCZOS)
                    
                    # Store the finished image in RAM for next time
                    tile_cache[chosen_tile_path] = tile_img
                
                # Paste it onto the canvas
                final_canvas.paste(tile_img, (col * OUTPUT_TILE_SIZE, row * OUTPUT_TILE_SIZE))
                
            except Exception as e:
                failed_tiles_count += 1
                if first_error is None:
                    # Fallback so full_repo_path exists in the error message if it fails before assignment
                    err_path = locals().get('full_repo_path', chosen_tile_path)
                    first_error = f"{type(e).__name__}: {str(e)} | Attempted: {err_path}"

    # ================= FINAL EXPORT =================
    progress(1.0, desc="Saving High-Res Output...")
    output_filename = "final_mosaic.jpg"
    try:
        final_canvas.save(output_filename, quality=95)
    except ValueError:
        output_filename = "final_mosaic.tiff"
        final_canvas.save(output_filename, format="TIFF")

    stats_msg = f"✅ Mosaic Complete! \nGrid: {actual_cols}x{actual_rows} tiles.\nFinal Resolution: {final_w}x{final_h}px"
    
    # Append the error report to the UI if anything failed
    if failed_tiles_count > 0:
        stats_msg += f"\n\n⚠️ WARNING: {failed_tiles_count} tiles failed to load!"
        stats_msg += f"\n🔍 TRACE: {first_error}"

    return final_canvas, output_filename, stats_msg
    
# ================= GRADIO UI CONFIGURATION =================
with gr.Blocks() as demo:
    gr.Markdown("# 🌍 GeoMosaic Engine")
    
    # TOP SECTION: 3 Columns
    with gr.Row():
        with gr.Column(scale=1):
            img_input = gr.Image(type="filepath", label="Target Image")
            
        with gr.Column(scale=1):
            mode_input = gr.Radio(
                choices=["🖨️ Type 1: Printout", "💻 Type 2: Finer Digital", "🔬 Type 3: Deep Zoom"],
                value="💻 Type 2: Finer Digital",
                label="Select Mosaic Mode"
            )
            submit_btn = gr.Button("Generate Mosaic", variant="primary")
            
        with gr.Column(scale=1):
            img_output = gr.Image(type="pil", label="Web Preview (Compressed)", interactive=False)
            file_output = gr.File(label="📥 Download High-Res Mosaic", interactive=False)
            stats_output = gr.Textbox(label="Build Stats", interactive=False)

    # BOTTOM SECTION: Full Width Parameters
    gr.Markdown("---")
    gr.Markdown("### ⚙️ Calculation Parameters")
    gr.Markdown("Adjust the targets for your selected mode. The engine will calculate the exact grid dynamically.")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("**Type 1: Print Targets**")
            print_w = gr.Number(value=24, label="Target Print Width (Inches)")
            print_dpi = gr.Number(value=300, label="Printer DPI")
            tile_in = gr.Number(value=0.25, label="Tile Size (Inches)")
        
        with gr.Column():
            gr.Markdown("**Type 2: Digital Targets**")
            max_w = gr.Number(value=7680, label="Max Screen Width (px)")
            dig_tile = gr.Number(value=16, label="Digital Tile Size (px)")
            
        with gr.Column():
            gr.Markdown("**Type 3: Deep Zoom Targets**")
            deep_patch = gr.Number(value=2, label="Original Image Sample (px)")
            deep_out = gr.Number(value=64, label="Output Tile Resolution (px)")

    submit_btn.click(
        fn=generate_mosaic,
        inputs=[img_input, mode_input, print_w, print_dpi, tile_in, max_w, dig_tile, deep_patch, deep_out],
        outputs=[img_output, file_output, stats_output]
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())