#!/usr/bin/env python3
import os
import glob
import base64
import time
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
from diffusers import StableDiffusionPipeline
from torch.utils.data import Dataset, DataLoader
import csv
import fitz
import requests
from PIL import Image
import numpy as np
import logging
import asyncio
import aiofiles
from io import BytesIO
from dataclasses import dataclass
from typing import Optional, Tuple
import zipfile
import math
import random
import re
import gradio as gr

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())

@dataclass
class ModelConfig:
    name: str
    base_model: str
    size: str
    domain: Optional[str] = None
    model_type: str = "causal_lm"
    @property
    def model_path(self):
        return f"models/{self.name}"

@dataclass
class DiffusionConfig:
    name: str
    base_model: str
    size: str
    domain: Optional[str] = None
    @property
    def model_path(self):
        return f"diffusion_models/{self.name}"

class ModelBuilder:
    def __init__(self):
        self.config = None
        self.model = None
        self.tokenizer = None
        self.jokes = ["Why did the AI go to therapy? Too many layers to unpack! 😂", "Training complete! Time for a binary coffee break. ☕"]
    def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
        self.model = AutoModelForCausalLM.from_pretrained(model_path)
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        if config:
            self.config = config
        self.model.to("cuda" if torch.cuda.is_available() else "cpu")
        return self
    def save_model(self, path: str):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        self.model.save_pretrained(path)
        self.tokenizer.save_pretrained(path)

class DiffusionBuilder:
    def __init__(self):
        self.config = None
        self.pipeline = None
    def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None):
        self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cpu")
        if config:
            self.config = config
        return self
    def save_model(self, path: str):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        self.pipeline.save_pretrained(path)
    def generate(self, prompt: str):
        return self.pipeline(prompt, num_inference_steps=20).images[0]

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

def pdf_url_to_filename(url):
    safe_name = re.sub(r'[<>:"/\\|?*]', '_', url)
    return f"{safe_name}.pdf"

def get_gallery_files(file_types=["png", "pdf"]):
    return sorted(list(set([f for ext in file_types for f in glob.glob(f"*.{ext}")])))  # Deduplicate files

def get_model_files(model_type="causal_lm"):
    path = "models/*" if model_type == "causal_lm" else "diffusion_models/*"
    dirs = [d for d in glob.glob(path) if os.path.isdir(d)]
    return dirs if dirs else ["None"]

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

async def process_pdf_snapshot(pdf_path, mode="single"):
    doc = fitz.open(pdf_path)
    output_files = []
    if mode == "single":
        page = doc[0]
        pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
        output_file = generate_filename("single", "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(2.0, 2.0))
            output_file = generate_filename(f"twopage_{i}", "png")
            pix.save(output_file)
            output_files.append(output_file)
    elif mode == "allpages":
        for i in range(len(doc)):
            page = doc[i]
            pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
            output_file = generate_filename(f"page_{i}", "png")
            pix.save(output_file)
            output_files.append(output_file)
    doc.close()
    return output_files

async def process_ocr(image, output_file):
    tokenizer = AutoTokenizer.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True)
    model = AutoModel.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
    temp_file = f"temp_{int(time.time())}.png"
    image.save(temp_file)
    result = model.chat(tokenizer, temp_file, ocr_type='ocr')
    os.remove(temp_file)
    async with aiofiles.open(output_file, "w") as f:
        await f.write(result)
    return result

async def process_image_gen(prompt, output_file, builder):
    if builder and isinstance(builder, DiffusionBuilder) and builder.pipeline:
        pipeline = builder.pipeline
    else:
        pipeline = StableDiffusionPipeline.from_pretrained("OFA-Sys/small-stable-diffusion-v0", torch_dtype=torch.float32).to("cpu")
    gen_image = pipeline(prompt, num_inference_steps=20).images[0]
    gen_image.save(output_file)
    return gen_image

# Gradio Interface Functions
def update_gallery(history, asset_checkboxes):
    all_files = get_gallery_files()
    gallery_images = []
    for file in all_files[:5]:  # Limit to 5 for display
        if file.endswith('.png'):
            gallery_images.append(Image.open(file))
        else:
            doc = fitz.open(file)
            pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            gallery_images.append(img)
            doc.close()
    history.append(f"Gallery updated: {len(all_files)} files")
    return gallery_images, history, asset_checkboxes

def camera_snap(image, cam_id, history, asset_checkboxes, cam_files):
    if image is not None:
        filename = generate_filename(f"cam{cam_id}")
        image.save(filename)
        history.append(f"Snapshot from Cam {cam_id}: {filename}")
        asset_checkboxes[filename] = True
        cam_files[cam_id] = filename
        return f"Image saved as {filename}", Image.open(filename), history, asset_checkboxes, cam_files
    elif cam_files.get(cam_id) and os.path.exists(cam_files[cam_id]):
        return f"Showing previous capture: {cam_files[cam_id]}", Image.open(cam_files[cam_id]), history, asset_checkboxes, cam_files
    return "No image captured", None, history, asset_checkboxes, cam_files

def download_pdfs(urls, history, asset_checkboxes):
    urls = urls.strip().split("\n")
    downloaded = []
    for url in urls:
        if url:
            output_path = pdf_url_to_filename(url)
            if download_pdf(url, output_path):
                downloaded.append(output_path)
                history.append(f"Downloaded PDF: {output_path}")
                asset_checkboxes[output_path] = True
    return f"Downloaded {len(downloaded)} PDFs", history, asset_checkboxes

def upload_pdfs(pdf_files, history, asset_checkboxes):
    uploaded = []
    for pdf_file in pdf_files:
        if pdf_file:
            output_path = f"uploaded_{int(time.time())}_{pdf_file.name}"
            with open(output_path, "wb") as f:
                f.write(pdf_file.read())
            uploaded.append(output_path)
            history.append(f"Uploaded PDF: {output_path}")
            asset_checkboxes[output_path] = True
    return f"Uploaded {len(uploaded)} PDFs", history, asset_checkboxes

def snapshot_pdfs(mode, history, asset_checkboxes):
    selected_pdfs = [path for path in get_gallery_files() if path.endswith('.pdf') and asset_checkboxes.get(path, False)]
    if not selected_pdfs:
        return "No PDFs selected", [], history, asset_checkboxes
    snapshots = []
    mode_key = {"Single Page (High-Res)": "single", "Two Pages (High-Res)": "twopage", "All Pages (High-Res)": "allpages"}[mode]
    for pdf_path in selected_pdfs:
        snap_files = asyncio.run(process_pdf_snapshot(pdf_path, mode_key))
        for snap in snap_files:
            snapshots.append(Image.open(snap))
            asset_checkboxes[snap] = True
            history.append(f"Snapshot {mode_key}: {snap}")
    return f"Generated {len(snapshots)} snapshots", snapshots, history, asset_checkboxes

def process_ocr_all(history, asset_checkboxes):
    all_files = get_gallery_files()
    if not all_files:
        return "No assets to OCR", history, asset_checkboxes
    full_text = "# OCR Results\n\n"
    for file in all_files:
        if file.endswith('.png'):
            image = Image.open(file)
        else:
            doc = fitz.open(file)
            pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
            image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            doc.close()
        output_file = generate_filename(f"ocr_{os.path.basename(file)}", "txt")
        result = asyncio.run(process_ocr(image, output_file))
        full_text += f"## {os.path.basename(file)}\n\n{result}\n\n"
        history.append(f"OCR Test: {file} -> {output_file}")
    md_output_file = f"full_ocr_{int(time.time())}.md"
    with open(md_output_file, "w") as f:
        f.write(full_text)
    return f"Full OCR saved to {md_output_file}", history, asset_checkboxes

def process_ocr_single(file_path, history, asset_checkboxes):
    if not file_path:
        return "No file selected", None, "", history, asset_checkboxes
    if file_path.endswith('.png'):
        image = Image.open(file_path)
    else:
        doc = fitz.open(file_path)
        pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
        image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        doc.close()
    output_file = generate_filename("ocr_output", "txt")
    result = asyncio.run(process_ocr(image, output_file))
    history.append(f"OCR Test: {file_path} -> {output_file}")
    return f"OCR output saved to {output_file}", image, result, history, asset_checkboxes

def build_model(model_type, base_model, model_name, domain, history):
    config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=model_name, base_model=base_model, size="small", domain=domain)
    builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
    builder.load_model(base_model, config)
    builder.save_model(config.model_path)
    history.append(f"Built {model_type} model: {model_name}")
    return builder, f"Model saved to {config.model_path}", history

def image_gen(prompt, file_path, builder, history, asset_checkboxes):
    if not file_path:
        return "No file selected", None, history, asset_checkboxes
    if file_path.endswith('.png'):
        image = Image.open(file_path)
    else:
        doc = fitz.open(file_path)
        pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
        image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        doc.close()
    output_file = generate_filename("gen_output", "png")
    gen_image = asyncio.run(process_image_gen(prompt, output_file, builder))
    history.append(f"Image Gen Test: {prompt} -> {output_file}")
    asset_checkboxes[output_file] = True
    return f"Image saved to {output_file}", gen_image, history, asset_checkboxes

# Gradio UI
with gr.Blocks(title="AI Vision & SFT Titans 🚀") as demo:
    gr.Markdown("# AI Vision & SFT Titans 🚀")
    history = gr.State(value=[])
    builder = gr.State(value=None)
    asset_checkboxes = gr.State(value={})
    cam_files = gr.State(value={})

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## Captured Files 📜")
            gallery_output = gr.Gallery(label="Asset Gallery", columns=2, height="auto")
            gr.Button("Update Gallery").click(update_gallery, inputs=[history, asset_checkboxes], outputs=[gallery_output, history, asset_checkboxes])
            gr.Markdown("## History 📜")
            history_output = gr.Textbox(label="History", lines=5, interactive=False)
            gr.Markdown("## Action Logs 📜")
            log_output = gr.Textbox(label="Logs", value="\n".join([f"{r.asctime} - {r.levelname} - {r.message}" for r in log_records]), lines=5, interactive=False)

        with gr.Column(scale=3):
            with gr.Tabs():
                with gr.TabItem("Camera Snap 📷"):
                    with gr.Row():
                        cam0_input = gr.Image(type="pil", label="Camera 0")
                        cam1_input = gr.Image(type="pil", label="Camera 1")
                    with gr.Row():
                        cam0_output = gr.Textbox(label="Cam 0 Status")
                        cam1_output = gr.Textbox(label="Cam 1 Status")
                    with gr.Row():
                        cam0_image = gr.Image(label="Cam 0 Preview")
                        cam1_image = gr.Image(label="Cam 1 Preview")
                    gr.Button("Capture Cam 0").click(camera_snap, inputs=[cam0_input, gr.State(value=0), history, asset_checkboxes, cam_files], outputs=[cam0_output, cam0_image, history, asset_checkboxes, cam_files])
                    gr.Button("Capture Cam 1").click(camera_snap, inputs=[cam1_input, gr.State(value=1), history, asset_checkboxes, cam_files], outputs=[cam1_output, cam1_image, history, asset_checkboxes, cam_files])

                with gr.TabItem("Download PDFs 📥"):
                    url_input = gr.Textbox(label="Enter PDF URLs (one per line)", lines=5)
                    pdf_upload = gr.File(label="Upload PDFs", file_count="multiple", type="binary")
                    pdf_output = gr.Textbox(label="Status")
                    snapshot_mode = gr.Dropdown(["Single Page (High-Res)", "Two Pages (High-Res)", "All Pages (High-Res)"], label="Snapshot Mode")
                    snapshot_output = gr.Textbox(label="Snapshot Status")
                    snapshot_images = gr.Gallery(label="Snapshots", columns=2, height="auto")
                    gr.Button("Download URLs").click(download_pdfs, inputs=[url_input, history, asset_checkboxes], outputs=[pdf_output, history, asset_checkboxes])
                    gr.Button("Upload PDFs").click(upload_pdfs, inputs=[pdf_upload, history, asset_checkboxes], outputs=[pdf_output, history, asset_checkboxes])
                    gr.Button("Snapshot Selected").click(snapshot_pdfs, inputs=[snapshot_mode, history, asset_checkboxes], outputs=[snapshot_output, snapshot_images, history, asset_checkboxes])

                with gr.TabItem("Test OCR 🔍"):
                    all_files = gr.Dropdown(choices=get_gallery_files(), label="Select File")
                    ocr_output = gr.Textbox(label="Status")
                    ocr_image = gr.Image(label="Input Image")
                    ocr_result = gr.Textbox(label="OCR Result", lines=5)
                    gr.Button("OCR All Assets").click(process_ocr_all, inputs=[history, asset_checkboxes], outputs=[ocr_output, history, asset_checkboxes])
                    gr.Button("OCR Selected").click(process_ocr_single, inputs=[all_files, history, asset_checkboxes], outputs=[ocr_output, ocr_image, ocr_result, history, asset_checkboxes])

                with gr.TabItem("Build Titan 🌱"):
                    model_type = gr.Dropdown(["Causal LM", "Diffusion"], label="Model Type")
                    base_model = gr.Dropdown(
                        choices=["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"],
                        label="Base Model",
                        value="HuggingFaceTB/SmolLM-135M"
                    )
                    model_name = gr.Textbox(label="Model Name", value=f"tiny-titan-{int(time.time())}")
                    domain = gr.Textbox(label="Target Domain", value="general")
                    build_output = gr.Textbox(label="Status")
                    gr.Button("Build").click(build_model, inputs=[model_type, base_model, model_name, domain, history], outputs=[builder, build_output, history])

                with gr.TabItem("Test Image Gen 🎨"):
                    gen_file = gr.Dropdown(choices=get_gallery_files(), label="Select Reference File")
                    gen_prompt = gr.Textbox(label="Prompt", value="Generate a neon superhero version of this image")
                    gen_output = gr.Textbox(label="Status")
                    gen_image = gr.Image(label="Generated Image")
                    gr.Button("Generate").click(image_gen, inputs=[gen_prompt, gen_file, builder, history, asset_checkboxes], outputs=[gen_output, gen_image, history, asset_checkboxes])

    # Update history output on every interaction
    demo.load(lambda h: "\n".join(h[-5:]), inputs=[history], outputs=[history_output])

demo.launch()