#!/usr/bin/env python3
import os
import glob
import time
import streamlit as st
import fitz  # PyMuPDF
import requests
from PIL import Image
from transformers import AutoTokenizer, AutoModel
from diffusers import StableDiffusionPipeline
import cv2
import numpy as np
import logging
import asyncio
import aiofiles
from io import BytesIO

# Logging setup
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
log_records = []

class LogCaptureHandler(logging.Handler):
    def emit(self, record):
        log_records.append(record)

logger.addHandler(LogCaptureHandler())

# Page Configuration
st.set_page_config(
    page_title="AI Vision Titans 🚀",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={'About': "AI Vision Titans: PDF Snapshots, OCR, Image Gen, Line Drawings on CPU! 🌌"}
)

# Initialize st.session_state
if 'captured_files' not in st.session_state:
    st.session_state['captured_files'] = []
if 'processing' not in st.session_state:
    st.session_state['processing'] = {}

# Utility Functions
def generate_filename(sequence, ext="png"):
    timestamp = time.strftime("%d%m%Y%H%M%S")
    return f"{sequence}{timestamp}.{ext}"

def get_gallery_files(file_types):
    return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")])

def update_gallery():
    media_files = get_gallery_files(["png", "txt"])
    if media_files:
        cols = st.sidebar.columns(2)
        for idx, file in enumerate(media_files[:gallery_size * 2]):
            with cols[idx % 2]:
                if file.endswith(".png"):
                    st.image(Image.open(file), caption=file, use_container_width=True)
                elif file.endswith(".txt"):
                    with open(file, "r") as f:
                        content = f.read()
                        st.text(content[:50] + "..." if len(content) > 50 else content, help=file)

def download_pdf(url, output_path):
    try:
        response = requests.get(url, stream=True, timeout=10)
        if response.status_code == 200:
            with open(output_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            return True
    except requests.RequestException as e:
        logger.error(f"Failed to download {url}: {e}")
    return False

# Model Loaders (CPU-focused)
def load_ocr_got():
    model_id = "ucaslcl/GOT-OCR2_0"
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    model = AutoModel.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
    return tokenizer, model

def load_image_gen():
    model_id = "OFA-Sys/small-stable-diffusion-v0"  # ~300 MB
    pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32).to("cpu")
    return pipeline

def load_line_drawer():
    def edge_detection(image, style="fine"):
        img_np = np.array(image.convert("RGB"))
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
        if style == "fine":
            edges = cv2.Canny(gray, 50, 150)  # Finer lines
        else:  # "bold"
            edges = cv2.Canny(gray, 100, 200)  # Bolder lines
        return Image.fromarray(edges)
    return edge_detection

# Async Processing Functions
async def process_pdf_snapshot(pdf_path, mode="thumbnail"):
    start_time = time.time()
    status = st.empty()
    status.text(f"Processing PDF Snapshot ({mode})... (0s)")
    doc = fitz.open(pdf_path)
    output_files = []
    
    if mode == "thumbnail":
        page = doc[0]
        pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5))  # 50% scale
        output_file = generate_filename("thumbnail", "png")
        pix.save(output_file)
        output_files.append(output_file)
    elif mode == "twopage":
        for i in range(min(2, len(doc))):
            page = doc[i]
            pix = page.get_pixmap(matrix=fitz.Matrix(1.0, 1.0))  # Full scale
            output_file = generate_filename(f"twopage_{i}", "png")
            pix.save(output_file)
            output_files.append(output_file)
    
    doc.close()
    elapsed = int(time.time() - start_time)
    status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
    for file in output_files:
        if file not in st.session_state['captured_files']:
            st.session_state['captured_files'].append(file)
    update_gallery()
    return output_files

async def process_ocr(image, output_file):
    start_time = time.time()
    status = st.empty()
    status.text("Processing GOT-OCR2_0... (0s)")
    tokenizer, model = load_ocr_got()
    result = model.chat(tokenizer, image, ocr_type='ocr')
    elapsed = int(time.time() - start_time)
    status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
    async with aiofiles.open(output_file, "w") as f:
        await f.write(result)
    if output_file not in st.session_state['captured_files']:
        st.session_state['captured_files'].append(output_file)
    update_gallery()
    return result

async def process_image_gen(prompt, output_file):
    start_time = time.time()
    status = st.empty()
    status.text("Processing Image Gen... (0s)")
    pipeline = load_image_gen()
    gen_image = pipeline(prompt, num_inference_steps=20).images[0]
    elapsed = int(time.time() - start_time)
    status.text(f"Image Gen completed in {elapsed}s!")
    gen_image.save(output_file)
    if output_file not in st.session_state['captured_files']:
        st.session_state['captured_files'].append(output_file)
    update_gallery()
    return gen_image

async def process_line_drawing(image, style, output_file):
    start_time = time.time()
    status = st.empty()
    status.text(f"Processing Line Drawing ({style})... (0s)")
    edge_fn = load_line_drawer()
    line_drawing = edge_fn(image, style=style)
    elapsed = int(time.time() - start_time)
    status.text(f"Line Drawing ({style}) completed in {elapsed}s!")
    line_drawing.save(output_file)
    if output_file not in st.session_state['captured_files']:
        st.session_state['captured_files'].append(output_file)
    update_gallery()
    return line_drawing

# Main App
st.title("AI Vision Titans 🚀")

# Sidebar Gallery
st.sidebar.header("Captured Files 📜")
gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 4)
update_gallery()

st.sidebar.subheader("Action Logs 📜")
log_container = st.sidebar.empty()
with log_container:
    for record in log_records:
        st.write(f"{record.asctime} - {record.levelname} - {record.message}")

# Tabs
tab1, tab2, tab3, tab4, tab5 = st.tabs(["Camera Snap 📷", "Download PDFs 📥", "Test OCR 🔍", "Test Image Gen 🎨", "Test Line Drawings ✏️"])

with tab1:
    st.header("Camera Snap 📷")
    st.subheader("Single Capture")
    cols = st.columns(2)
    with cols[0]:
        cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0")
        if cam0_img:
            filename = generate_filename(0)
            if filename not in st.session_state['captured_files']:
                with open(filename, "wb") as f:
                    f.write(cam0_img.getvalue())
                st.image(Image.open(filename), caption=filename, use_container_width=True)
                logger.info(f"Saved snapshot from Camera 0: {filename}")
                st.session_state['captured_files'].append(filename)
                update_gallery()
    with cols[1]:
        cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
        if cam1_img:
            filename = generate_filename(1)
            if filename not in st.session_state['captured_files']:
                with open(filename, "wb") as f:
                    f.write(cam1_img.getvalue())
                st.image(Image.open(filename), caption=filename, use_container_width=True)
                logger.info(f"Saved snapshot from Camera 1: {filename}")
                st.session_state['captured_files'].append(filename)
                update_gallery()

    st.subheader("Burst Capture")
    slice_count = st.number_input("Number of Frames", min_value=1, max_value=20, value=10, key="burst_count")
    if st.button("Start Burst Capture 📸"):
        st.session_state['burst_frames'] = []
        placeholder = st.empty()
        for i in range(slice_count):
            with placeholder.container():
                st.write(f"Capturing frame {i+1}/{slice_count}...")
                img = st.camera_input(f"Frame {i}", key=f"burst_{i}_{time.time()}")
                if img:
                    filename = generate_filename(f"burst_{i}")
                    if filename not in st.session_state['captured_files']:
                        with open(filename, "wb") as f:
                            f.write(img.getvalue())
                        st.session_state['burst_frames'].append(filename)
                        logger.info(f"Saved burst frame {i}: {filename}")
                        st.image(Image.open(filename), caption=filename, use_container_width=True)
                    time.sleep(0.5)
        st.session_state['captured_files'].extend([f for f in st.session_state['burst_frames'] if f not in st.session_state['captured_files']])
        update_gallery()
        placeholder.success(f"Captured {len(st.session_state['burst_frames'])} frames!")

with tab2:
    st.header("Download PDFs 📥")
    url_input = st.text_area("Enter PDF URLs (one per line)", height=100)
    mode = st.selectbox("Snapshot Mode", ["Thumbnail", "Two-Page View"], key="download_mode")
    if st.button("Download & Snapshot 📸"):
        urls = url_input.strip().split("\n")
        for url in urls:
            if url:
                pdf_path = generate_filename("downloaded", "pdf")
                if download_pdf(url, pdf_path):
                    logger.info(f"Downloaded PDF from {url} to {pdf_path}")
                    snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode.lower().replace(" ", "")))
                    for snapshot in snapshots:
                        st.image(Image.open(snapshot), caption=snapshot, use_container_width=True)
                else:
                    st.error(f"Failed to download {url}")

with tab3:
    st.header("Test OCR 🔍")
    captured_files = get_gallery_files(["png"])
    if captured_files:
        selected_file = st.selectbox("Select Image", captured_files, key="ocr_select")
        image = Image.open(selected_file)
        st.image(image, caption="Input Image", use_container_width=True)
        if st.button("Run OCR 🚀", key="ocr_run"):
            output_file = generate_filename("ocr_output", "txt")
            st.session_state['processing']['ocr'] = True
            result = asyncio.run(process_ocr(image, output_file))
            st.text_area("OCR Result", result, height=200, key="ocr_result")
            st.success(f"OCR output saved to {output_file}")
            st.session_state['processing']['ocr'] = False
    else:
        st.warning("No images captured yet. Use Camera Snap or Download PDFs first!")

with tab4:
    st.header("Test Image Gen 🎨")
    captured_files = get_gallery_files(["png"])
    if captured_files:
        selected_file = st.selectbox("Select Image", captured_files, key="gen_select")
        image = Image.open(selected_file)
        st.image(image, caption="Reference Image", use_container_width=True)
        prompt = st.text_area("Prompt", "Generate a similar superhero image", key="gen_prompt")
        if st.button("Run Image Gen 🚀", key="gen_run"):
            output_file = generate_filename("gen_output", "png")
            st.session_state['processing']['gen'] = True
            result = asyncio.run(process_image_gen(prompt, output_file))
            st.image(result, caption="Generated Image", use_container_width=True)
            st.success(f"Image saved to {output_file}")
            st.session_state['processing']['gen'] = False
    else:
        st.warning("No images captured yet. Use Camera Snap or Download PDFs first!")

with tab5:
    st.header("Test Line Drawings ✏️")
    captured_files = get_gallery_files(["png"])
    if captured_files:
        selected_file = st.selectbox("Select Image", captured_files, key="line_select")
        image = Image.open(selected_file)
        st.image(image, caption="Input Image", use_container_width=True)
        style = st.selectbox("Line Style", ["Fine", "Bold"], key="line_style")
        if st.button("Run Line Drawing 🚀", key="line_run"):
            output_file = generate_filename(f"line_{style.lower()}", "png")
            st.session_state['processing']['line'] = True
            result = asyncio.run(process_line_drawing(image, style.lower(), output_file))
            st.image(result, caption=f"{style} Line Drawing", use_container_width=True)
            st.success(f"Line drawing saved to {output_file}")
            st.session_state['processing']['line'] = False
    else:
        st.warning("No images captured yet. Use Camera Snap or Download PDFs first!")

# Initial Gallery Update
update_gallery()