#!/usr/bin/env python3
import os
import base64
import streamlit as st
import csv
import time
from dataclasses import dataclass
import zipfile
import logging
from PIL import Image
import numpy as np
import cv2

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

st.set_page_config(page_title="SFT Tiny Titans 🚀", page_icon="🤖", layout="wide", initial_sidebar_state="expanded")

# Model Configurations
@dataclass
class ModelConfig:
    name: str
    base_model: str
    model_type: str = "causal_lm"
    @property
    def model_path(self):
        return f"models/{self.name}"

@dataclass
class DiffusionConfig:
    name: str
    base_model: str
    @property
    def model_path(self):
        return f"diffusion_models/{self.name}"

# Lazy-loaded Builders
class ModelBuilder:
    def __init__(self):
        self.config = None
        self.model = None
        self.tokenizer = None
    def load_model(self, model_path: str, config: ModelConfig):
        try:
            from transformers import AutoModelForCausalLM, AutoTokenizer
            import torch
            logger.info(f"Loading NLP model: {model_path}")
            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
            self.config = config
            self.model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
            logger.info("NLP model loaded successfully")
        except Exception as e:
            logger.error(f"Error loading NLP model: {str(e)}")
            raise
    def fine_tune(self, csv_path):
        try:
            from torch.utils.data import Dataset, DataLoader
            import torch
            logger.info(f"Starting NLP fine-tuning with {csv_path}")
            class SFTDataset(Dataset):
                def __init__(self, data, tokenizer):
                    self.data = data
                    self.tokenizer = tokenizer
                def __len__(self):
                    return len(self.data)
                def __getitem__(self, idx):
                    prompt = self.data[idx]["prompt"]
                    response = self.data[idx]["response"]
                    inputs = self.tokenizer(f"{prompt} {response}", return_tensors="pt", padding="max_length", max_length=128, truncation=True)
                    labels = inputs["input_ids"].clone()
                    labels[0, :len(self.tokenizer(prompt)["input_ids"][0])] = -100
                    return {"input_ids": inputs["input_ids"][0], "attention_mask": inputs["attention_mask"][0], "labels": labels[0]}
            data = []
            with open(csv_path, "r") as f:
                reader = csv.DictReader(f)
                for row in reader:
                    data.append({"prompt": row["prompt"], "response": row["response"]})
            dataset = SFTDataset(data, self.tokenizer)
            dataloader = DataLoader(dataset, batch_size=2)
            optimizer = torch.optim.AdamW(self.model.parameters(), lr=2e-5)
            self.model.train()
            for _ in range(1):
                for batch in dataloader:
                    optimizer.zero_grad()
                    outputs = self.model(**{k: v.to(self.model.device) for k, v in batch.items()})
                    outputs.loss.backward()
                    optimizer.step()
            logger.info("NLP fine-tuning completed")
        except Exception as e:
            logger.error(f"Error in NLP fine-tuning: {str(e)}")
            raise
    def evaluate(self, prompt: str):
        try:
            import torch
            logger.info(f"Evaluating NLP with prompt: {prompt}")
            self.model.eval()
            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)
                result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
                logger.info(f"NLP evaluation result: {result}")
                return result
        except Exception as e:
            logger.error(f"Error in NLP evaluation: {str(e)}")
            raise

class DiffusionBuilder:
    def __init__(self):
        self.config = None
        self.pipeline = None
    def load_model(self, model_path: str, config: DiffusionConfig):
        try:
            from diffusers import StableDiffusionPipeline
            import torch
            logger.info(f"Loading diffusion model: {model_path}")
            self.pipeline = StableDiffusionPipeline.from_pretrained(model_path)
            self.pipeline.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
            self.config = config
            logger.info("Diffusion model loaded successfully")
        except Exception as e:
            logger.error(f"Error loading diffusion model: {str(e)}")
            raise
    def fine_tune(self, images, texts):
        try:
            import torch
            import numpy as np
            logger.info("Starting diffusion fine-tuning")
            optimizer = torch.optim.AdamW(self.pipeline.unet.parameters(), lr=1e-5)
            self.pipeline.unet.train()
            for _ in range(1):
                for img, text in zip(images, texts):
                    optimizer.zero_grad()
                    img_tensor = torch.tensor(np.array(img)).permute(2, 0, 1).unsqueeze(0).float().to(self.pipeline.device) / 255.0
                    latents = self.pipeline.vae.encode(img_tensor).latent_dist.sample()
                    noise = torch.randn_like(latents)
                    timesteps = torch.randint(0, self.pipeline.scheduler.num_train_timesteps, (1,), device=latents.device)
                    noisy_latents = self.pipeline.scheduler.add_noise(latents, noise, timesteps)
                    text_emb = self.pipeline.text_encoder(self.pipeline.tokenizer(text, return_tensors="pt").input_ids.to(self.pipeline.device))[0]
                    pred_noise = self.pipeline.unet(noisy_latents, timesteps, encoder_hidden_states=text_emb).sample
                    loss = torch.nn.functional.mse_loss(pred_noise, noise)
                    loss.backward()
                    optimizer.step()
            logger.info("Diffusion fine-tuning completed")
        except Exception as e:
            logger.error(f"Error in diffusion fine-tuning: {str(e)}")
            raise
    def generate(self, prompt: str):
        try:
            logger.info(f"Generating image with prompt: {prompt}")
            img = self.pipeline(prompt, num_inference_steps=20).images[0]
            logger.info("Image generated successfully")
            return img
        except Exception as e:
            logger.error(f"Error in image generation: {str(e)}")
            raise

# Utilities
def get_download_link(file_path, mime_type="text/plain", label="Download"):
    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>'

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

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

def zip_files(files, zip_name):
    with zipfile.ZipFile(zip_name, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for file in files:
            zipf.write(file, os.path.basename(file))
    return zip_name

# Main App
st.title("SFT Tiny Titans 🚀 (Camera Input Action!)")

# Sidebar Galleries
st.sidebar.header("Captured Media 🎨")
gallery_container = st.sidebar.empty()
def update_gallery():
    media_files = get_gallery_files(["png"])
    with gallery_container:
        if media_files:
            cols = st.columns(2)
            for idx, file in enumerate(media_files[:4]):
                with cols[idx % 2]:
                    st.image(Image.open(file), caption=file.split('/')[-1], use_container_width=True)

# Sidebar Model Management
st.sidebar.subheader("Model Hub 🗂️")
model_type = st.sidebar.selectbox("Model Type", ["NLP (Causal LM)", "CV (Diffusion)"])
model_options = {
    "NLP (Causal LM)": "HuggingFaceTB/SmolLM-135M",
    "CV (Diffusion)": ["CompVis/stable-diffusion-v1-4", "stabilityai/stable-diffusion-2-base", "runwayml/stable-diffusion-v1-5"]
}
selected_model = st.sidebar.selectbox("Select Model", ["None"] + ([model_options[model_type]] if "NLP" in model_type else model_options[model_type]))
if selected_model != "None" and st.sidebar.button("Load Model 📂"):
    builder = ModelBuilder() if "NLP" in model_type else DiffusionBuilder()
    config = (ModelConfig if "NLP" in model_type else DiffusionConfig)(name=f"titan_{int(time.time())}", base_model=selected_model)
    with st.spinner("Loading... ⏳"):
        try:
            builder.load_model(selected_model, config)
            st.session_state['builder'] = builder
            st.session_state['model_loaded'] = True
            st.success("Model loaded! 🎉")
        except Exception as e:
            st.error(f"Load failed: {str(e)}")

# Tabs
tab1, tab2, tab3 = st.tabs(["Build Titan 🌱", "Camera Snap 📷", "Fine-Tune & Test 🔧🧪"])

with tab1:
    st.header("Build Titan 🌱 (Quick Start!)")
    model_type = st.selectbox("Model Type", ["NLP (Causal LM)", "CV (Diffusion)"], key="build_type")
    base_model = st.selectbox("Select Model", model_options[model_type], key="build_model")
    if st.button("Download Model ⬇️"):
        config = (ModelConfig if "NLP" in model_type else DiffusionConfig)(name=f"titan_{int(time.time())}", base_model=base_model)
        builder = ModelBuilder() if "NLP" in model_type else DiffusionBuilder()
        with st.spinner("Fetching... ⏳"):
            try:
                builder.load_model(base_model, config)
                st.session_state['builder'] = builder
                st.session_state['model_loaded'] = True
                st.success("Titan up! 🎉")
            except Exception as e:
                st.error(f"Download failed: {str(e)}")

with tab2:
    st.header("Camera Snap 📷 (Dual Capture!)")
    cols = st.columns(2)
    with cols[0]:
        st.subheader("Camera 0")
        cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0")
        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}")
            if 'captured_images' not in st.session_state:
                st.session_state['captured_images'] = []
            st.session_state['captured_images'].append(filename)
            update_gallery()
    with cols[1]:
        st.subheader("Camera 1")
        cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
        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}")
            if 'captured_images' not in st.session_state:
                st.session_state['captured_images'] = []
            st.session_state['captured_images'].append(filename)
            update_gallery()

    st.subheader("Capture 10 Frames (Video Simulation)")
    cols = st.columns(2)
    with cols[0]:
        if st.button("Capture 10 Frames - Cam 0 📸"):
            st.session_state['cam0_frames'] = []
            for i in range(10):
                img = st.camera_input(f"Frame {i} - Cam 0", key=f"cam0_frame_{i}")
                if img:
                    filename = generate_filename(f"0_{i}")
                    with open(filename, "wb") as f:
                        f.write(img.getvalue())
                    st.session_state['cam0_frames'].append(filename)
                    logger.info(f"Saved frame {i} from Camera 0: {filename}")
                    time.sleep(0.5)  # Simulate video frame rate
            if 'captured_images' not in st.session_state:
                st.session_state['captured_images'] = []
            st.session_state['captured_images'].extend(st.session_state['cam0_frames'])
            update_gallery()
            for frame in st.session_state['cam0_frames']:
                st.image(Image.open(frame), caption=frame, use_container_width=True)
    with cols[1]:
        if st.button("Capture 10 Frames - Cam 1 📸"):
            st.session_state['cam1_frames'] = []
            for i in range(10):
                img = st.camera_input(f"Frame {i} - Cam 1", key=f"cam1_frame_{i}")
                if img:
                    filename = generate_filename(f"1_{i}")
                    with open(filename, "wb") as f:
                        f.write(img.getvalue())
                    st.session_state['cam1_frames'].append(filename)
                    logger.info(f"Saved frame {i} from Camera 1: {filename}")
                    time.sleep(0.5)  # Simulate video frame rate
            if 'captured_images' not in st.session_state:
                st.session_state['captured_images'] = []
            st.session_state['captured_images'].extend(st.session_state['cam1_frames'])
            update_gallery()
            for frame in st.session_state['cam1_frames']:
                st.image(Image.open(frame), caption=frame, use_container_width=True)

with tab3:
    st.header("Fine-Tune & Test 🔧🧪")
    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
        st.warning("Load a Titan first! ⚠️")
    else:
        if isinstance(st.session_state['builder'], ModelBuilder):
            st.subheader("NLP Tune 🧠")
            uploaded_csv = st.file_uploader("Upload CSV", type="csv", key="nlp_csv")
            if uploaded_csv and st.button("Tune NLP 🔄"):
                logger.info("Initiating NLP fine-tune")
                try:
                    with open("temp.csv", "wb") as f:
                        f.write(uploaded_csv.read())
                    st.session_state['builder'].fine_tune("temp.csv")
                    st.success("NLP sharpened! 🎉")
                except Exception as e:
                    st.error(f"NLP fine-tune failed: {str(e)}")
            st.subheader("NLP Test 🧠")
            prompt = st.text_area("Prompt", "What’s a superhero?", key="nlp_test")
            if st.button("Test NLP ▶️"):
                logger.info("Running NLP test")
                try:
                    result = st.session_state['builder'].evaluate(prompt)
                    st.write(f"**Answer**: {result}")
                except Exception as e:
                    st.error(f"NLP test failed: {str(e)}")
        elif isinstance(st.session_state['builder'], DiffusionBuilder):
            st.subheader("CV Tune 🎨")
            captured_images = get_gallery_files(["png"])
            if len(captured_images) >= 2:
                texts = ["Superhero Neon", "Hero Glow", "Cape Spark"][:len(captured_images)]
                if st.button("Tune CV 🔄"):
                    logger.info("Initiating CV fine-tune")
                    try:
                        images = [Image.open(img) for img in captured_images]
                        st.session_state['builder'].fine_tune(images, texts)
                        st.success("CV polished! 🎉")
                    except Exception as e:
                        st.error(f"CV fine-tune failed: {str(e)}")
            else:
                st.warning("Capture at least 2 images in Camera Snap first! ⚠️")
            st.subheader("CV Test 🎨 (Image Set Demo)")
            if len(captured_images) >= 2:
                if st.button("Run CV Demo ▶️"):
                    logger.info("Running CV image set demo")
                    try:
                        images = [Image.open(img) for img in captured_images[:10]]
                        prompts = ["Neon " + os.path.basename(img).split('.')[0] for img in captured_images[:10]]
                        generated_images = []
                        for prompt in prompts:
                            img = st.session_state['builder'].generate(prompt)
                            generated_images.append(img)
                        cols = st.columns(2)
                        for idx, (orig, gen) in enumerate(zip(images, generated_images)):
                            with cols[idx % 2]:
                                st.image(orig, caption=f"Original: {captured_images[idx]}", use_container_width=True)
                                st.image(gen, caption=f"Generated: {prompts[idx]}", use_container_width=True)
                        md_content = "# Image Set Demo\n\nScript of filenames and descriptions:\n"
                        for i, (img, prompt) in enumerate(zip(captured_images[:10], prompts)):
                            md_content += f"{i+1}. `{img}` - {prompt}\n"
                        md_filename = f"demo_metadata_{int(time.time())}.md"
                        with open(md_filename, "w") as f:
                            f.write(md_content)
                        st.markdown(get_download_link(md_filename, "text/markdown", "Download Metadata .md"), unsafe_allow_html=True)
                        logger.info("CV demo completed with metadata")
                    except Exception as e:
                        st.error(f"CV demo failed: {str(e)}")
                        logger.error(f"Error in CV demo: {str(e)}")
            else:
                st.warning("Capture at least 2 images in Camera Snap first! ⚠️")

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

update_gallery()  # Initial gallery update