#!/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 streamlit.components.v1 import html
from PIL import Image

# 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

# JavaScript/HTML Components
camera_selector_html = """
<div>
  <h3>Camera & Audio Source Selector</h3>
  <select id="videoSource"></select>
  <select id="audioSource"></select>
  <button onclick="startStream()">Start Stream</button>
  <video id="video" autoplay playsinline style="width: 100%;"></video>
</div>
<script>
  const videoSource = document.getElementById('videoSource');
  const audioSource = document.getElementById('audioSource');
  const video = document.getElementById('video');
  let stream;

  navigator.mediaDevices.enumerateDevices().then(devices => {
    devices.forEach(device => {
      const option = document.createElement('option');
      option.value = device.deviceId;
      option.text = device.label || `${device.kind} ${device.deviceId}`;
      if (device.kind === 'videoinput') {
        videoSource.appendChild(option);
      } else if (device.kind === 'audioinput') {
        audioSource.appendChild(option);
      }
    });
  }).catch(err => console.error('Error enumerating devices:', err));

  function startStream() {
    if (stream) {
      stream.getTracks().forEach(track => track.stop());
    }
    const constraints = {
      video: { deviceId: videoSource.value ? { exact: videoSource.value } : undefined },
      audio: { deviceId: audioSource.value ? { exact: audioSource.value } : undefined }
    };
    navigator.mediaDevices.getUserMedia(constraints)
      .then(mediaStream => {
        stream = mediaStream;
        video.srcObject = stream;
        console.log('Stream started');
      })
      .catch(err => console.error('Error starting stream:', err));
  }
</script>
"""

image_capture_html = """
<div>
  <h3>Image Capture - Camera {id}</h3>
  <video id="video{id}" autoplay playsinline style="width: 100%;"></video>
  <button onclick="captureFrame{id}()">Capture Frame 📸</button>
  <canvas id="canvas{id}" style="display: none;"></canvas>
</div>
<script>
  const video{id} = document.getElementById('video{id}');
  const canvas{id} = document.getElementById('canvas{id}');
  let stream{id};

  navigator.mediaDevices.getUserMedia({ video: true })
    .then(mediaStream => {
      stream{id} = mediaStream;
      video{id}.srcObject = stream{id};
      console.log('Camera {id} stream started');
    })
    .catch(err => console.error('Error starting Camera {id}:', err));

  function captureFrame{id}() {
    canvas{id}.width = video{id}.videoWidth;
    canvas{id}.height = video{id}.videoHeight;
    const ctx = canvas{id}.getContext('2d');
    ctx.drawImage(video{id}, 0, 0, canvas{id}.width, canvas{id}.height);
    const dataUrl = canvas{id}.toDataURL('image/png');
    const filename = `{id}${new Date().toISOString().replace(/[^0-9]/g, '')}.png`;
    const link = document.createElement('a');
    link.href = dataUrl;
    link.download = filename;
    link.click();
    console.log('Captured frame:', filename);
  }
</script>
"""

video_capture_html = """
<div>
  <h3>Video Capture - Camera {id}</h3>
  <video id="video{id}" autoplay playsinline style="width: 100%;"></video>
  <button onclick="captureVideo{id}()">Capture Video 🎥</button>
  <canvas id="canvas{id}" style="display: none;"></canvas>
</div>
<script>
  const video{id} = document.getElementById('video{id}');
  const canvas{id} = document.getElementById('canvas{id}');
  let stream{id}, recorder{id};

  navigator.mediaDevices.getUserMedia({ video: true, audio: true })
    .then(mediaStream => {
      stream{id} = mediaStream;
      video{id}.srcObject = stream{id};
      recorder{id} = new MediaRecorder(stream{id});
      const chunks = [];
      recorder{id}.ondataavailable = e => chunks.push(e.data);
      recorder{id}.onstop = () => {
        const blob = new Blob(chunks, { type: 'video/mp4' });
        const filename = `{id}${new Date().toISOString().replace(/[^0-9]/g, '')}.mp4`;
        const url = URL.createObjectURL(blob);
        const link = document.createElement('a');
        link.href = url;
        link.download = filename;
        link.click();
        console.log('Captured video:', filename);
        sliceVideo{id}(blob);
      };
      console.log('Camera {id} stream started');
    })
    .catch(err => console.error('Error starting Camera {id}:', err));

  function captureVideo{id}() {
    recorder{id}.start();
    setTimeout(() => recorder{id}.stop(), 10000); // 10 seconds
    console.log('Recording started for Camera {id}');
  }

  function sliceVideo{id}(blob) {
    const video = document.createElement('video');
    video.src = URL.createObjectURL(blob);
    video.onloadedmetadata = () => {
      const ctx = canvas{id}.getContext('2d');
      canvas{id}.width = video.videoWidth;
      canvas{id}.height = video.videoHeight;
      let frameCount = 0;
      const interval = video.duration / 10;
      video.currentTime = 0;
      const captureFrame = () => {
        if (frameCount < 10) {
          ctx.drawImage(video, 0, 0, canvas{id}.width, canvas{id}.height);
          const dataUrl = canvas{id}.toDataURL('image/png');
          const filename = `{id}${new Date().toISOString().replace(/[^0-9]/g, '')}_${frameCount}.png`;
          const link = document.createElement('a');
          link.href = dataUrl;
          link.download = filename;
          link.click();
          console.log('Captured frame:', filename);
          frameCount++;
          video.currentTime += interval;
          setTimeout(captureFrame, 100);
        }
      };
      video.play().then(captureFrame);
    };
  }
</script>
"""

# Main App
st.title("SFT Tiny Titans 🚀 (Web Cam Action!)")

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

# 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, tab4 = st.tabs(["Build Titan 🌱", "Camera Snap 📷", "Fine-Tune Titans 🔧", "Test Titans 🧪"])

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 Live Feed!)")
    st.subheader("Source Configuration")
    html(camera_selector_html, height=400)
    cols = st.columns(2)
    for i in range(2):
        with cols[i]:
            html(image_capture_html.format(id=i), height=300)
            html(video_capture_html.format(id=i), height=300)
    st.subheader("Upload Captured Files")
    uploaded_files = st.file_uploader("Upload PNGs/MP4s from Downloads", type=["png", "mp4"], accept_multiple_files=True)
    if uploaded_files:
        for file in uploaded_files:
            filename = file.name
            with open(filename, "wb") as f:
                f.write(file.read())
            logger.info(f"Saved uploaded file: {filename}")
            if filename.endswith(".png"):
                st.image(Image.open(filename), caption=filename, use_container_width=True)
            elif filename.endswith(".mp4"):
                st.video(filename)
        update_gallery()

with tab3:
    st.header("Fine-Tune Titans 🔧 (Tune Fast!)")
    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)}")
        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("Upload at least 2 PNGs in Camera Snap first! ⚠️")

with tab4:
    st.header("Test Titans 🧪 (Image Agent Demo!)")
    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 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 Test 🎨 (Image Set Demo)")
            captured_images = get_gallery_files(["png"])
            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("Upload at least 2 PNGs 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