#!/usr/bin/env python3
import os
import glob
import base64
import streamlit as st
import pandas as pd
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.utils.data import Dataset, DataLoader
import csv
import time
from dataclasses import dataclass
from typing import Optional, Tuple
import zipfile
import math
from PIL import Image
import random
import logging
import numpy as np

# Logging setup with a custom buffer
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="SFT Tiny Titans 🚀",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={
        'Get Help': 'https://huggingface.co/awacke1',
        'Report a Bug': 'https://huggingface.co/spaces/awacke1',
        'About': "Tiny Titans: Small models, big dreams, and a sprinkle of chaos! 🌌"
    }
)

# Initialize st.session_state
if 'captured_images' not in st.session_state:
    st.session_state['captured_images'] = []
if 'nlp_builder' not in st.session_state:
    st.session_state['nlp_builder'] = None
if 'cv_builder' not in st.session_state:
    st.session_state['cv_builder'] = None
if 'nlp_loaded' not in st.session_state:
    st.session_state['nlp_loaded'] = False
if 'cv_loaded' not in st.session_state:
    st.session_state['cv_loaded'] = False
if 'active_tab' not in st.session_state:
    st.session_state['active_tab'] = "Build Titan 🌱"

# Model Configuration Classes
@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
    @property
    def model_path(self):
        return f"diffusion_models/{self.name}"

# Datasets
class SFTDataset(Dataset):
    def __init__(self, data, tokenizer, max_length=128):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        prompt = self.data[idx]["prompt"]
        response = self.data[idx]["response"]
        full_text = f"{prompt} {response}"
        full_encoding = self.tokenizer(full_text, max_length=self.max_length, padding="max_length", truncation=True, return_tensors="pt")
        prompt_encoding = self.tokenizer(prompt, max_length=self.max_length, padding=False, truncation=True, return_tensors="pt")
        input_ids = full_encoding["input_ids"].squeeze()
        attention_mask = full_encoding["attention_mask"].squeeze()
        labels = input_ids.clone()
        prompt_len = prompt_encoding["input_ids"].shape[1]
        if prompt_len < self.max_length:
            labels[:prompt_len] = -100
        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

class DiffusionDataset(Dataset):
    def __init__(self, images, texts):
        self.images = images
        self.texts = texts
    def __len__(self):
        return len(self.images)
    def __getitem__(self, idx):
        return {"image": self.images[idx], "text": self.texts[idx]}

# Model Builders
class ModelBuilder:
    def __init__(self):
        self.config = None
        self.model = None
        self.tokenizer = None
        self.sft_data = 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):
        try:
            with st.spinner(f"Loading {model_path}... ⏳ (Patience, young padawan!)"):
                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")
            st.success(f"Model loaded! 🎉 {random.choice(self.jokes)}")
            logger.info(f"Successfully loaded Causal LM model: {model_path}")
        except torch.cuda.OutOfMemoryError as e:
            st.error(f"GPU memory error loading {model_path}: {str(e)} 💥 (Out of GPU juice!)")
            logger.error(f"GPU memory error loading {model_path}: {str(e)}")
            raise
        except MemoryError as e:
            st.error(f"CPU memory error loading {model_path}: {str(e)} 💥 (RAM ran away!)")
            logger.error(f"CPU memory error loading {model_path}: {str(e)}")
            raise
        except Exception as e:
            st.error(f"Failed to load {model_path}: {str(e)} 💥 (Something broke—check the logs!)")
            logger.error(f"Failed to load {model_path}: {str(e)}")
            raise
        return self
    def fine_tune_sft(self, csv_path: str, epochs: int = 3, batch_size: int = 4):
        try:
            self.sft_data = []
            with open(csv_path, "r") as f:
                reader = csv.DictReader(f)
                for row in reader:
                    self.sft_data.append({"prompt": row["prompt"], "response": row["response"]})
            dataset = SFTDataset(self.sft_data, self.tokenizer)
            dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
            optimizer = torch.optim.AdamW(self.model.parameters(), lr=2e-5)
            self.model.train()
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            self.model.to(device)
            for epoch in range(epochs):
                with st.spinner(f"Training epoch {epoch + 1}/{epochs}... ⚙️ (The AI is lifting weights!)"):
                    total_loss = 0
                    for batch in dataloader:
                        optimizer.zero_grad()
                        input_ids = batch["input_ids"].to(device)
                        attention_mask = batch["attention_mask"].to(device)
                        labels = batch["labels"].to(device)
                        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                        loss = outputs.loss
                        loss.backward()
                        optimizer.step()
                        total_loss += loss.item()
                    st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
            st.success(f"SFT Fine-tuning completed! 🎉 {random.choice(self.jokes)}")
            logger.info(f"Successfully fine-tuned Causal LM model: {self.config.name}")
        except Exception as e:
            st.error(f"Fine-tuning failed: {str(e)} 💥 (Training hit a snag!)")
            logger.error(f"Fine-tuning failed: {str(e)}")
            raise
        return self
    def save_model(self, path: str):
        try:
            with st.spinner("Saving model... 💾 (Packing the AI’s suitcase!)"):
                os.makedirs(os.path.dirname(path), exist_ok=True)
                self.model.save_pretrained(path)
                self.tokenizer.save_pretrained(path)
            st.success(f"Model saved at {path}! ✅ May the force be with it.")
            logger.info(f"Model saved at {path}")
        except Exception as e:
            st.error(f"Failed to save model: {str(e)} 💥 (Save operation crashed!)")
            logger.error(f"Failed to save model: {str(e)}")
            raise
    def evaluate(self, prompt: str, status_container=None):
        self.model.eval()
        if status_container:
            status_container.write("Preparing to evaluate... 🧠 (Titan’s warming up its circuits!)")
        logger.info(f"Evaluating prompt: {prompt}")
        try:
            with torch.no_grad():
                inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.model.device)
                outputs = self.model.generate(**inputs, max_new_tokens=50, do_sample=True, top_p=0.95, temperature=0.7)
                result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
                logger.info(f"Generated response: {result}")
                return result
        except Exception as e:
            logger.error(f"Evaluation error: {str(e)}")
            if status_container:
                status_container.error(f"Oops! Something broke: {str(e)} 💥 (Titan tripped over a wire!)")
            return f"Error: {str(e)}"

class DiffusionBuilder:
    def __init__(self):
        self.config = None
        self.pipeline = None
    def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None):
        from diffusers import StableDiffusionPipeline
        try:
            with st.spinner(f"Loading diffusion model {model_path}... ⏳"):
                self.pipeline = StableDiffusionPipeline.from_pretrained(model_path)
                self.pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
                if config:
                    self.config = config
            st.success(f"Diffusion model loaded! 🎨")
            logger.info(f"Successfully loaded Diffusion model: {model_path}")
        except torch.cuda.OutOfMemoryError as e:
            st.error(f"GPU memory error loading {model_path}: {str(e)} 💥 (Out of GPU juice!)")
            logger.error(f"GPU memory error loading {model_path}: {str(e)}")
            raise
        except MemoryError as e:
            st.error(f"CPU memory error loading {model_path}: {str(e)} 💥 (RAM ran away!)")
            logger.error(f"CPU memory error loading {model_path}: {str(e)}")
            raise
        except Exception as e:
            st.error(f"Failed to load {model_path}: {str(e)} 💥 (Something broke—check the logs!)")
            logger.error(f"Failed to load {model_path}: {str(e)}")
            raise
        return self
    def fine_tune_sft(self, images, texts, epochs=3):
        try:
            dataset = DiffusionDataset(images, texts)
            dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
            optimizer = torch.optim.AdamW(self.pipeline.unet.parameters(), lr=1e-5)
            self.pipeline.unet.train()
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            for epoch in range(epochs):
                with st.spinner(f"Training diffusion epoch {epoch + 1}/{epochs}... ⚙️"):
                    total_loss = 0
                    for batch in dataloader:
                        optimizer.zero_grad()
                        image = batch["image"][0].to(device)
                        text = batch["text"][0]
                        latents = self.pipeline.vae.encode(torch.tensor(np.array(image)).permute(2, 0, 1).unsqueeze(0).float().to(device)).latent_dist.sample()
                        noise = torch.randn_like(latents)
                        timesteps = torch.randint(0, self.pipeline.scheduler.num_train_timesteps, (latents.shape[0],), device=latents.device)
                        noisy_latents = self.pipeline.scheduler.add_noise(latents, noise, timesteps)
                        text_embeddings = self.pipeline.text_encoder(self.pipeline.tokenizer(text, return_tensors="pt").input_ids.to(device))[0]
                        pred_noise = self.pipeline.unet(noisy_latents, timesteps, encoder_hidden_states=text_embeddings).sample
                        loss = torch.nn.functional.mse_loss(pred_noise, noise)
                        loss.backward()
                        optimizer.step()
                        total_loss += loss.item()
                    st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
            st.success("Diffusion SFT Fine-tuning completed! 🎨")
            logger.info(f"Successfully fine-tuned Diffusion model: {self.config.name}")
        except Exception as e:
            st.error(f"Fine-tuning failed: {str(e)} 💥 (Training hit a snag!)")
            logger.error(f"Fine-tuning failed: {str(e)}")
            raise
        return self
    def save_model(self, path: str):
        try:
            with st.spinner("Saving diffusion model... 💾"):
                os.makedirs(os.path.dirname(path), exist_ok=True)
                self.pipeline.save_pretrained(path)
            st.success(f"Diffusion model saved at {path}! ✅")
            logger.info(f"Diffusion model saved at {path}")
        except Exception as e:
            st.error(f"Failed to save model: {str(e)} 💥 (Save operation crashed!)")
            logger.error(f"Failed to save model: {str(e)}")
            raise
    def generate(self, prompt: str):
        try:
            return self.pipeline(prompt, num_inference_steps=50).images[0]
        except Exception as e:
            st.error(f"Image generation failed: {str(e)} 💥 (Pixel party pooper!)")
            logger.error(f"Image generation failed: {str(e)}")
            raise

# Utility Functions
def generate_filename(sequence, ext="png"):
    from datetime import datetime
    import pytz
    central = pytz.timezone('US/Central')
    dt = datetime.now(central)
    return f"{dt.strftime('%m-%d-%Y-%I-%M-%S-%p')}.{ext}"

def get_download_link(file_path, mime_type="text/plain", label="Download"):
    try:
        with open(file_path, 'rb') as f:
            data = f.read()
        b64 = base64.b64encode(data).decode()
        return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label} 📥</a>'
    except Exception as e:
        logger.error(f"Failed to generate download link for {file_path}: {str(e)}")
        return f"Error: Could not generate link for {file_path}"

def zip_files(files, zip_path):
    try:
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for file in files:
                zipf.write(file, os.path.basename(file))
        logger.info(f"Created ZIP file: {zip_path}")
    except Exception as e:
        logger.error(f"Failed to create ZIP file {zip_path}: {str(e)}")
        raise

def delete_files(files):
    try:
        for file in files:
            os.remove(file)
            logger.info(f"Deleted file: {file}")
        st.session_state['captured_images'] = [f for f in st.session_state['captured_images'] if f not in files]
    except Exception as e:
        logger.error(f"Failed to delete files: {str(e)}")
        raise

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

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

def update_gallery():
    media_files = get_gallery_files(["png"])
    if media_files:
        cols = st.sidebar.columns(2)
        for idx, file in enumerate(media_files[:gallery_size * 2]):
            with cols[idx % 2]:
                st.image(Image.open(file), caption=file, use_container_width=True)
                st.markdown(get_download_link(file, "image/png", "Download Image"), unsafe_allow_html=True)

# Mock Search Tool for RAG
def mock_search(query: str) -> str:
    if "superhero" in query.lower():
        return "Latest trends: Gold-plated Batman statues, VR superhero battles."
    return "No relevant results found."

class PartyPlannerAgent:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
    def generate(self, prompt: str) -> str:
        self.model.eval()
        with torch.no_grad():
            inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.device)
            outputs = self.model.generate(**inputs, max_new_tokens=100, do_sample=True, top_p=0.95, temperature=0.7)
            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
    def plan_party(self, task: str) -> pd.DataFrame:
        search_result = mock_search("superhero party trends")
        prompt = f"Given this context: '{search_result}'\n{task}"
        plan_text = self.generate(prompt)
        locations = {"Wayne Manor": (42.3601, -71.0589), "New York": (40.7128, -74.0060)}
        wayne_coords = locations["Wayne Manor"]
        travel_times = {loc: calculate_cargo_travel_time(coords, wayne_coords) for loc, coords in locations.items() if loc != "Wayne Manor"}
        data = [
            {"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gold-plated Batman statues"},
            {"Location": "Wayne Manor", "Travel Time (hrs)": 0.0, "Luxury Idea": "VR superhero battles"}
        ]
        return pd.DataFrame(data)

class CVPartyPlannerAgent:
    def __init__(self, pipeline):
        self.pipeline = pipeline
    def generate(self, prompt: str) -> Image.Image:
        return self.pipeline(prompt, num_inference_steps=50).images[0]
    def plan_party(self, task: str) -> pd.DataFrame:
        search_result = mock_search("superhero party trends")
        prompt = f"Given this context: '{search_result}'\n{task}"
        data = [
            {"Theme": "Batman", "Image Idea": "Gold-plated Batman statue"},
            {"Theme": "Avengers", "Image Idea": "VR superhero battle scene"}
        ]
        return pd.DataFrame(data)

def calculate_cargo_travel_time(origin_coords: Tuple[float, float], destination_coords: Tuple[float, float], cruising_speed_kmh: float = 750.0) -> float:
    def to_radians(degrees: float) -> float:
        return degrees * (math.pi / 180)
    lat1, lon1 = map(to_radians, origin_coords)
    lat2, lon2 = map(to_radians, destination_coords)
    EARTH_RADIUS_KM = 6371.0
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = (math.sin(dlat / 2) ** 2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2) ** 2)
    c = 2 * math.asin(math.sqrt(a))
    distance = EARTH_RADIUS_KM * c
    actual_distance = distance * 1.1
    flight_time = (actual_distance / cruising_speed_kmh) + 1.0
    return round(flight_time, 2)

# Main App
st.title("SFT Tiny Titans 🚀 (Small but Mighty!)")

# Sidebar Galleries and File Management
st.sidebar.header("Media Gallery 🎨")
gallery_size = st.sidebar.slider("Gallery Size 📸", 1, 10, 4, help="Adjust how many epic captures you see! 🌟")
update_gallery()

col1, col2 = st.sidebar.columns(2)
with col1:
    if st.button("Download All 📦"):
        media_files = get_gallery_files(["png"])
        if media_files:
            zip_path = f"snapshot_collection_{int(time.time())}.zip"
            zip_files(media_files, zip_path)
            st.sidebar.markdown(get_download_link(zip_path, "application/zip", "Download All Snapshots"), unsafe_allow_html=True)
            st.sidebar.success("Snapshots zipped and ready! 🎉 Grab your loot!")
        else:
            st.sidebar.warning("No snapshots to zip! 📸 Snap some pics first!")
with col2:
    if st.button("Delete All 🗑️"):
        media_files = get_gallery_files(["png"])
        if media_files:
            delete_files(media_files)
            st.sidebar.success("All snapshots vanquished! 🧹 Gallery cleared!")
            update_gallery()
        else:
            st.sidebar.warning("Nothing to delete! 📸 Snap some pics to clear later!")

# File Uploader
uploaded_files = st.sidebar.file_uploader("Upload Files 🎵🎥🖼️📝📜", type=["mp3", "mp4", "png", "jpeg", "md", "pdf", "docx"], accept_multiple_files=True)
if uploaded_files:
    for uploaded_file in uploaded_files:
        filename = uploaded_file.name
        with open(filename, "wb") as f:
            f.write(uploaded_file.getvalue())
        logger.info(f"Uploaded file: {filename}")

# Sidebar Galleries by Type
st.sidebar.subheader("Audio Gallery 🎵")
audio_files = get_gallery_files(["mp3"])
if audio_files:
    for file in audio_files[:gallery_size]:
        st.sidebar.audio(file, format="audio/mp3")
        st.sidebar.markdown(get_download_link(file, "audio/mp3", f"Download {file}"), unsafe_allow_html=True)

st.sidebar.subheader("Video Gallery 🎥")
video_files = get_gallery_files(["mp4"])
if video_files:
    for file in video_files[:gallery_size]:
        st.sidebar.video(file, format="video/mp4")
        st.sidebar.markdown(get_download_link(file, "video/mp4", f"Download {file}"), unsafe_allow_html=True)

st.sidebar.subheader("Image Gallery 🖼️")
image_files = get_gallery_files(["png", "jpeg"])
if image_files:
    cols = st.sidebar.columns(2)
    for idx, file in enumerate(image_files[:gallery_size * 2]):
        with cols[idx % 2]:
            st.image(Image.open(file), caption=file, use_container_width=True)
            st.markdown(get_download_link(file, "image/png" if file.endswith(".png") else "image/jpeg", f"Download {file}"), unsafe_allow_html=True)

st.sidebar.subheader("Markdown Gallery 📝")
md_files = get_gallery_files(["md"])
if md_files:
    for file in md_files[:gallery_size]:
        with open(file, "r") as f:
            st.sidebar.markdown(f.read())
        st.sidebar.markdown(get_download_link(file, "text/markdown", f"Download {file}"), unsafe_allow_html=True)

st.sidebar.subheader("Document Gallery 📜")
doc_files = get_gallery_files(["pdf", "docx"])
if doc_files:
    for file in doc_files[:gallery_size]:
        mime_type = "application/pdf" if file.endswith(".pdf") else "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        st.sidebar.markdown(get_download_link(file, mime_type, f"Download {file}"), unsafe_allow_html=True)

st.sidebar.subheader("Model Management 🗂️")
model_type = st.sidebar.selectbox("Model Type", ["Causal LM", "Diffusion"])
model_dirs = get_model_files("causal_lm" if model_type == "Causal LM" else "diffusion")
selected_model = st.sidebar.selectbox("Select Saved Model", ["None"] + model_dirs)
if selected_model != "None" and st.sidebar.button("Load Model 📂"):
    builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
    config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=os.path.basename(selected_model), base_model="unknown", size="small")
    try:
        builder.load_model(selected_model, config)
        if model_type == "Causal LM":
            st.session_state['nlp_builder'] = builder
            st.session_state['nlp_loaded'] = True
        else:
            st.session_state['cv_builder'] = builder
            st.session_state['cv_loaded'] = True
        st.rerun()
    except Exception as e:
        st.error(f"Model load failed: {str(e)} 💥 (Check logs for details!)")

st.sidebar.subheader("Model Status 🚦")
st.sidebar.write(f"**NLP Model**: {'Loaded' if st.session_state['nlp_loaded'] else 'Not Loaded'} {'(Active)' if st.session_state['nlp_loaded'] and isinstance(st.session_state.get('nlp_builder'), ModelBuilder) else ''}")
st.sidebar.write(f"**CV Model**: {'Loaded' if st.session_state['cv_loaded'] else 'Not Loaded'} {'(Active)' if st.session_state['cv_loaded'] and isinstance(st.session_state.get('cv_builder'), DiffusionBuilder) else ''}")

# Tabs
tabs = [
    "Build Titan 🌱", "Camera Snap 📷", 
    "Fine-Tune Titan (NLP) 🔧", "Test Titan (NLP) 🧪", "Agentic RAG Party (NLP) 🌐",
    "Fine-Tune Titan (CV) 🔧", "Test Titan (CV) 🧪", "Agentic RAG Party (CV) 🌐"
]
tab1, tab2, tab3, tab4, tab5, tab6, tab7, tab8 = st.tabs(tabs)

# Log Tab Switches
for i, tab in enumerate(tabs):
    if st.session_state['active_tab'] != tab and st.session_state.get(f'tab{i}_active', False):
        logger.info(f"Switched to tab: {tab}")
        st.session_state['active_tab'] = tab
    st.session_state[f'tab{i}_active'] = (st.session_state['active_tab'] == tab)

with tab1:
    st.header("Build Titan 🌱")
    model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
    base_model = st.selectbox("Select Tiny Model", 
        ["HuggingFaceTB/SmolLM-135M", "HuggingFaceTB/SmolLM-360M", "Qwen/Qwen1.5-0.5B-Chat"] if model_type == "Causal LM" else 
        ["stabilityai/stable-diffusion-2-base", "runwayml/stable-diffusion-v1-5"])
    model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
    domain = st.text_input("Target Domain", "general", help="Where will your Titan flex its muscles? 💪") if model_type == "Causal LM" else None
    if st.button("Download Model ⬇️"):
        config = ModelConfig(name=model_name, base_model=base_model, size="small", domain=domain) if model_type == "Causal LM" else DiffusionConfig(name=model_name, base_model=base_model, size="small")
        builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
        try:
            builder.load_model(base_model, config)
            builder.save_model(config.model_path)
            if model_type == "Causal LM":
                st.session_state['nlp_builder'] = builder
                st.session_state['nlp_loaded'] = True
            else:
                st.session_state['cv_builder'] = builder
                st.session_state['cv_loaded'] = True
            st.rerun()
        except Exception as e:
            st.error(f"Model build failed: {str(e)} 💥 (Check logs for details!)")

with tab2:
    st.header("Camera Snap 📷 (Dual Capture!)")
    slice_count = st.number_input("Image Slice Count 🎞️", min_value=1, max_value=20, value=10, help="How many snaps to dream of? (Automation’s on vacation! 😜)")
    video_length = st.number_input("Video Dream Length (seconds) 🎥", min_value=1, max_value=30, value=10, help="Imagine a vid this long—sadly, we’re stuck with pics for now! 😂")
    cols = st.columns(2)
    with cols[0]:
        st.subheader("Camera 0 🎬")
        cam0_img = st.camera_input("Snap a Shot - Cam 0 📸", key="cam0", help="Click to capture a heroic moment! 🦸‍♂️")
        if cam0_img:
            filename = generate_filename(0)
            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_images'].append(filename)
            update_gallery()
        st.info("🚨 Multi-frame capture’s on strike! Snap one at a time—your Titan’s too cool for automation glitches! 😎")
    with cols[1]:
        st.subheader("Camera 1 🎥")
        cam1_img = st.camera_input("Snap a Shot - Cam 1 📸", key="cam1", help="Grab another epic frame! 🌟")
        if cam1_img:
            filename = generate_filename(1)
            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_images'].append(filename)
            update_gallery()
        st.info("🚨 Frame bursts? Nope, manual snaps only! One click, one masterpiece! 🎨")

with tab3:  # Fine-Tune Titan (NLP)
    st.header("Fine-Tune Titan (NLP) 🔧 (Teach Your Word Wizard Some Tricks!)")
    if not st.session_state['nlp_loaded'] or not isinstance(st.session_state['nlp_builder'], ModelBuilder):
        st.warning("Please build or load an NLP Titan first! ⚠️ (No word wizard, no magic!)")
    else:
        if st.button("Generate Sample CSV 📝"):
            sample_data = [
                {"prompt": "What is AI?", "response": "AI is artificial intelligence, simulating human smarts in machines."},
                {"prompt": "Explain machine learning", "response": "Machine learning is AI’s gym where models bulk up on data."},
                {"prompt": "What is a neural network?", "response": "A neural network is a brainy AI mimicking human noggins."},
            ]
            csv_path = f"sft_data_{int(time.time())}.csv"
            with open(csv_path, "w", newline="") as f:
                writer = csv.DictWriter(f, fieldnames=["prompt", "response"])
                writer.writeheader()
                writer.writerows(sample_data)
            st.markdown(get_download_link(csv_path, "text/csv", "Download Sample CSV"), unsafe_allow_html=True)
            st.success(f"Sample CSV generated as {csv_path}! ✅ (Fresh from the data oven!)")
        uploaded_csv = st.file_uploader("Upload CSV for SFT 📜", type="csv", help="Feed your Titan some tasty prompt-response pairs! 🍽️")
        if uploaded_csv and st.button("Fine-Tune with Uploaded CSV 🔄"):
            csv_path = f"uploaded_sft_data_{int(time.time())}.csv"
            with open(csv_path, "wb") as f:
                f.write(uploaded_csv.read())
            new_model_name = f"{st.session_state['nlp_builder'].config.name}-sft-{int(time.time())}"
            new_config = ModelConfig(name=new_model_name, base_model=st.session_state['nlp_builder'].config.base_model, size="small", domain=st.session_state['nlp_builder'].config.domain)
            st.session_state['nlp_builder'].config = new_config
            with st.status("Fine-tuning NLP Titan... ⏳ (Whipping words into shape!)", expanded=True) as status:
                st.session_state['nlp_builder'].fine_tune_sft(csv_path)
                st.session_state['nlp_builder'].save_model(new_config.model_path)
                status.update(label="Fine-tuning completed! 🎉 (Wordsmith Titan unleashed!)", state="complete")
            zip_path = f"{new_config.model_path}.zip"
            zip_directory(new_config.model_path, zip_path)
            st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned NLP Titan"), unsafe_allow_html=True)

with tab4:  # Test Titan (NLP)
    st.header("Test Titan (NLP) 🧪 (Put Your Word Wizard to the Test!)")
    if not st.session_state['nlp_loaded'] or not isinstance(st.session_state['nlp_builder'], ModelBuilder):
        st.warning("Please build or load an NLP Titan first! ⚠️ (No word wizard, no test drive!)")
    else:
        if st.session_state['nlp_builder'].sft_data:
            st.write("Testing with SFT Data:")
            with st.spinner("Running SFT data tests... ⏳ (Titan’s flexing its word muscles!)"):
                for item in st.session_state['nlp_builder'].sft_data[:3]:
                    prompt = item["prompt"]
                    expected = item["response"]
                    status_container = st.empty()
                    generated = st.session_state['nlp_builder'].evaluate(prompt, status_container)
                    st.write(f"**Prompt**: {prompt}")
                    st.write(f"**Expected**: {expected}")
                    st.write(f"**Generated**: {generated} (Titan says: '{random.choice(['Bleep bloop!', 'I am groot!', '42!'])}')")
                    st.write("---")
                    status_container.empty()
        test_prompt = st.text_area("Enter Test Prompt 🗣️", "What is AI?", help="Ask your Titan anything—it’s ready to chat! 😜")
        if st.button("Run Test ▶️"):
            with st.spinner("Testing your prompt... ⏳ (Titan’s pondering deeply!)"):
                status_container = st.empty()
                result = st.session_state['nlp_builder'].evaluate(test_prompt, status_container)
                st.write(f"**Generated Response**: {result} (Titan’s wisdom unleashed!)")
                status_container.empty()

with tab5:  # Agentic RAG Party (NLP)
    st.header("Agentic RAG Party (NLP) 🌐 (Party Like It’s 2099!)")
    st.write("This demo uses your SFT-tuned NLP Titan to plan a superhero party with mock retrieval!")
    if not st.session_state['nlp_loaded'] or not isinstance(st.session_state['nlp_builder'], ModelBuilder):
        st.warning("Please build or load an NLP Titan first! ⚠️ (No word wizard, no party!)")
    else:
        if st.button("Run NLP RAG Demo 🎉"):
            with st.spinner("Loading your SFT-tuned NLP Titan... ⏳ (Titan’s suiting up!)"):
                agent = PartyPlannerAgent(st.session_state['nlp_builder'].model, st.session_state['nlp_builder'].tokenizer)
                st.write("Agent ready! 🦸‍♂️ (Time to plan an epic bash!)")
            task = """
            Plan a luxury superhero-themed party at Wayne Manor (42.3601° N, 71.0589° W). 
            Use mock search results for the latest superhero party trends, refine for luxury elements 
            (decorations, entertainment, catering), and calculate cargo travel times from key locations 
            (New York: 40.7128° N, 74.0060° W; LA: 34.0522° N, 118.2437° W; London: 51.5074° N, 0.1278° W) 
            to Wayne Manor. Create a plan with at least 6 entries in a pandas dataframe.
            """
            with st.spinner("Planning the ultimate superhero bash... ⏳ (Calling all caped crusaders!)"):
                try:
                    locations = {
                        "Wayne Manor": (42.3601, -71.0589),
                        "New York": (40.7128, -74.0060),
                        "Los Angeles": (34.0522, -118.2437),
                        "London": (51.5074, -0.1278)
                    }
                    wayne_coords = locations["Wayne Manor"]
                    travel_times = {loc: calculate_cargo_travel_time(coords, wayne_coords) for loc, coords in locations.items() if loc != "Wayne Manor"}
                    search_result = mock_search("superhero party trends")
                    prompt = f"""
                    Given this context from a search: "{search_result}"
                    Plan a luxury superhero-themed party at Wayne Manor. Suggest luxury decorations, entertainment, and catering ideas.
                    """
                    plan_text = agent.generate(prompt)
                    catchphrases = ["To the Batmobile!", "Avengers, assemble!", "I am Iron Man!", "By the power of Grayskull!"]
                    data = [
                        {"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gold-plated Batman statues", "Catchphrase": random.choice(catchphrases)},
                        {"Location": "Los Angeles", "Travel Time (hrs)": travel_times["Los Angeles"], "Luxury Idea": "Holographic Avengers displays", "Catchphrase": random.choice(catchphrases)},
                        {"Location": "London", "Travel Time (hrs)": travel_times["London"], "Luxury Idea": "Live stunt shows with Iron Man suits", "Catchphrase": random.choice(catchphrases)},
                        {"Location": "Wayne Manor", "Travel Time (hrs)": 0.0, "Luxury Idea": "VR superhero battles", "Catchphrase": random.choice(catchphrases)},
                        {"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gourmet kryptonite-green cocktails", "Catchphrase": random.choice(catchphrases)},
                        {"Location": "Los Angeles", "Travel Time (hrs)": travel_times["Los Angeles"], "Luxury Idea": "Thor’s hammer-shaped appetizers", "Catchphrase": random.choice(catchphrases)},
                    ]
                    plan_df = pd.DataFrame(data)
                    st.write("Agentic RAG Party Plan:")
                    st.dataframe(plan_df)
                    st.write("Party on, Wayne! 🦸‍♂️🎉")
                except Exception as e:
                    st.error(f"Error planning party: {str(e)} (Even Superman has kryptonite days!)")
                    logger.error(f"Error in NLP RAG demo: {str(e)}")

with tab6:  # Fine-Tune Titan (CV)
    st.header("Fine-Tune Titan (CV) 🔧 (Paint Your Titan’s Masterpiece!)")
    if not st.session_state['cv_loaded'] or not isinstance(st.session_state['cv_builder'], DiffusionBuilder):
        st.warning("Please build or load a CV Titan first! ⚠️ (No artist, no canvas!)")
    else:
        captured_images = get_gallery_files(["png"])
        if len(captured_images) >= 2:
            demo_data = [{"image": img, "text": f"Superhero {os.path.basename(img).split('.')[0]}"} for img in captured_images[:min(len(captured_images), 10)]]
            edited_data = st.data_editor(pd.DataFrame(demo_data), num_rows="dynamic", help="Craft your image-text pairs like a superhero artist! 🎨")
            if st.button("Fine-Tune with Dataset 🔄"):
                images = [Image.open(row["image"]) for _, row in edited_data.iterrows()]
                texts = [row["text"] for _, row in edited_data.iterrows()]
                new_model_name = f"{st.session_state['cv_builder'].config.name}-sft-{int(time.time())}"
                new_config = DiffusionConfig(name=new_model_name, base_model=st.session_state['cv_builder'].config.base_model, size="small")
                st.session_state['cv_builder'].config = new_config
                with st.status("Fine-tuning CV Titan... ⏳ (Brushing up those pixels!)", expanded=True) as status:
                    st.session_state['cv_builder'].fine_tune_sft(images, texts)
                    st.session_state['cv_builder'].save_model(new_config.model_path)
                    status.update(label="Fine-tuning completed! 🎉 (Pixel Titan unleashed!)", state="complete")
                zip_path = f"{new_config.model_path}.zip"
                zip_directory(new_config.model_path, zip_path)
                st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned CV Titan"), unsafe_allow_html=True)
            csv_path = f"sft_dataset_{int(time.time())}.csv"
            with open(csv_path, "w", newline="") as f:
                writer = csv.writer(f)
                writer.writerow(["image", "text"])
                for _, row in edited_data.iterrows():
                    writer.writerow([row["image"], row["text"]])
            st.markdown(get_download_link(csv_path, "text/csv", "Download SFT Dataset CSV"), unsafe_allow_html=True)

with tab7:  # Test Titan (CV)
    st.header("Test Titan (CV) 🧪 (Unleash Your Pixel Power!)")
    if not st.session_state['cv_loaded'] or not isinstance(st.session_state['cv_builder'], DiffusionBuilder):
        st.warning("Please build or load a CV Titan first! ⚠️ (No artist, no masterpiece!)")
    else:
        test_prompt = st.text_area("Enter Test Prompt 🎨", "Neon Batman", help="Dream up a wild image—your Titan’s got the brush! 🖌️")
        if st.button("Run Test ▶️"):
            with st.spinner("Painting your masterpiece... ⏳ (Titan’s mixing colors!)"):
                image = st.session_state['cv_builder'].generate(test_prompt)
                st.image(image, caption="Generated Image", use_container_width=True)

with tab8:  # Agentic RAG Party (CV)
    st.header("Agentic RAG Party (CV) 🌐 (Party with Pixels!)")
    st.write("This demo uses your SFT-tuned CV Titan to generate superhero party images with mock retrieval!")
    if not st.session_state['cv_loaded'] or not isinstance(st.session_state['cv_builder'], DiffusionBuilder):
        st.warning("Please build or load a CV Titan first! ⚠️ (No artist, no party!)")
    else:
        if st.button("Run CV RAG Demo 🎉"):
            with st.spinner("Loading your SFT-tuned CV Titan... ⏳ (Titan’s grabbing its paintbrush!)"):
                agent = CVPartyPlannerAgent(st.session_state['cv_builder'].pipeline)
                st.write("Agent ready! 🎨 (Time to paint an epic bash!)")
            task = "Generate images for a luxury superhero-themed party."
            with st.spinner("Crafting superhero party visuals... ⏳ (Pixels assemble!)"):
                try:
                    plan_df = agent.plan_party(task)
                    st.dataframe(plan_df)
                    for _, row in plan_df.iterrows():
                        image = agent.generate(row["Image Idea"])
                        st.image(image, caption=f"{row['Theme']} - {row['Image Idea']}", use_container_width=True)
                except Exception as e:
                    st.error(f"Error in CV RAG demo: {str(e)} 💥 (Pixel party crashed!)")
                    logger.error(f"Error in CV RAG demo: {str(e)}")

# Display Logs
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}")

# Initial Gallery Update
update_gallery()