import aiofiles
import asyncio
import base64
import cv2
import fitz
import glob
import io
import json
import logging
import os
import pandas as pd
import pytz
import random
import re
import requests
import shutil
import streamlit as st
import sys
import time
import torch
import zipfile

from audio_recorder_streamlit import audio_recorder
from contextlib import redirect_stdout
from dataclasses import dataclass
from datetime import datetime
from diffusers import StableDiffusionPipeline
from io import BytesIO
from moviepy.editor import VideoFileClip
from openai import OpenAI
from PIL import Image
from PyPDF2 import PdfReader
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
from typing import Optional

# Initialize OpenAI client
client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'), organization=os.getenv('OPENAI_ORG_ID'))

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

# Streamlit configuration
st.set_page_config(
    page_title="AI Multimodal Titan 🚀",
    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': "AI Multimodal Titan: PDFs, OCR, Image Gen, Audio/Video Processing, Code Execution, and More! 🌌"
    }
)

# Session state initialization
for key in ['history', 'builder', 'model_loaded', 'processing', 'asset_checkboxes', 'downloaded_pdfs', 'unique_counter', 'messages']:
    st.session_state.setdefault(key, [] if key in ['history', 'messages'] else {} if key in ['asset_checkboxes', 'downloaded_pdfs', 'processing'] else None if key == 'builder' else 0 if key == 'unique_counter' else False)
st.session_state.setdefault('selected_model_type', "Causal LM")
st.session_state.setdefault('selected_model', "None")
st.session_state.setdefault('gallery_size', 2)
st.session_state.setdefault('asset_gallery_container', st.sidebar.empty())

@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
    def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
        with st.spinner(f"Loading {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
            if config:
                self.config = config
            self.model.to("cuda" if torch.cuda.is_available() else "cpu")
        st.success(f"Model loaded! 🎉")
        return self
    def save_model(self, path: str):
        with st.spinner("Saving model... 💾"):
            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}! ✅")

class DiffusionBuilder:
    def __init__(self):
        self.config = None
        self.pipeline = None
    def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None):
        with st.spinner(f"Loading diffusion model {model_path}... ⏳"):
            self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cpu")
            if config:
                self.config = config
        st.success("Diffusion model loaded! 🎨")
        return self
    def save_model(self, path: str):
        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}! ✅")
    def generate(self, prompt: str):
        return self.pipeline(prompt, num_inference_steps=20).images[0]

def generate_filename(prompt, ext="png"):
    central = pytz.timezone('US/Central')
    safe_date_time = datetime.now(central).strftime("%m%d_%H%M")
    safe_prompt = re.sub(r'[<>:"/\\|?*]', '_', prompt)[:240]
    return f"{safe_date_time}_{safe_prompt}.{ext}"

def get_download_link(file_path, mime_type="application/pdf", label="Download"):
    with open(file_path, "rb") as f:
        data = base64.b64encode(f.read()).decode()
    return f'<a href="data:{mime_type};base64,{data}" download="{os.path.basename(file_path)}">{label}</a>'

def zip_directory(directory_path, zip_path):
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, _, files in os.walk(directory_path):
            for file in files:
                zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.dirname(directory_path)))

def get_gallery_files(file_types=["png", "pdf", "md", "wav", "mp4"]):
    return sorted(list({f for ext in file_types for f in glob.glob(f"*.{ext}")}))

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"):
    start_time = time.time()
    status = st.empty()
    status.text(f"Processing PDF Snapshot ({mode})... (0s)")
    try:
        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 == "double":
            if len(doc) >= 2:
                pix1 = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                pix2 = doc[1].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                img1 = Image.frombytes("RGB", [pix1.width, pix1.height], pix1.samples)
                img2 = Image.frombytes("RGB", [pix2.width, pix2.height], pix2.samples)
                combined_img = Image.new("RGB", (pix1.width + pix2.width, max(pix1.height, pix2.height)))
                combined_img.paste(img1, (0, 0))
                combined_img.paste(img2, (pix1.width, 0))
                output_file = generate_filename("double", "png")
                combined_img.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()
        elapsed = int(time.time() - start_time)
        status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
        return output_files
    except Exception as e:
        status.error(f"Failed to process PDF: {str(e)}")
        return []

async def process_ocr(image, output_file):
    start_time = time.time()
    status = st.empty()
    status.text("Processing GOT-OCR2_0... (0s)")
    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 = generate_filename("temp", "png")
    image.save(temp_file)
    result = model.chat(tokenizer, temp_file, ocr_type='ocr')
    os.remove(temp_file)
    elapsed = int(time.time() - start_time)
    status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
    async with aiofiles.open(output_file, "w") as f:
        await f.write(result)
    return result

async def process_image_gen(prompt, output_file):
    start_time = time.time()
    status = st.empty()
    status.text("Processing Image Gen... (0s)")
    pipeline = st.session_state['builder'].pipeline if st.session_state.get('builder') and isinstance(st.session_state['builder'], DiffusionBuilder) else 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]
    elapsed = int(time.time() - start_time)
    status.text(f"Image Gen completed in {elapsed}s!")
    gen_image.save(output_file)
    return gen_image

def process_image_with_prompt(image, prompt, model="gpt-4o-mini", detail="auto"):
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}", "detail": detail}}]}]
    try:
        response = client.chat.completions.create(model=model, messages=messages, max_tokens=300)
        return response.choices[0].message.content
    except Exception as e:
        return f"Error processing image with GPT: {str(e)}"

def process_text_with_prompt(text, prompt, model="gpt-4o-mini"):
    messages = [{"role": "user", "content": f"{prompt}\n\n{text}"}]
    try:
        response = client.chat.completions.create(model=model, messages=messages, max_tokens=300)
        return response.choices[0].message.content
    except Exception as e:
        return f"Error processing text with GPT: {str(e)}"

def process_audio(audio_input, prompt):
    with open(audio_input, "rb") as file:
        transcription = client.audio.transcriptions.create(model="whisper-1", file=file)
    response = client.chat.completions.create(model="gpt-4o-mini", messages=[{"role": "user", "content": f"{prompt}\n\n{transcription.text}"}])
    return transcription.text, response.choices[0].message.content

def process_video(video_path, prompt):
    base64Frames, audio_path = process_video_frames(video_path)
    with open(video_path, "rb") as file:
        transcription = client.audio.transcriptions.create(model="whisper-1", file=file)
    messages = [{"role": "user", "content": ["These are the frames from the video.", *map(lambda x: {"type": "image_url", "image_url": {"url": f'data:image/jpg;base64,{x}', "detail": "low"}}, base64Frames), {"type": "text", "text": f"The audio transcription is: {transcription.text}\n\n{prompt}"}]}]
    response = client.chat.completions.create(model="gpt-4o-mini", messages=messages)
    return response.choices[0].message.content

def process_video_frames(video_path, seconds_per_frame=2):
    base64Frames = []
    base_video_path, _ = os.path.splitext(video_path)
    video = cv2.VideoCapture(video_path)
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = video.get(cv2.CAP_PROP_FPS)
    frames_to_skip = int(fps * seconds_per_frame)
    curr_frame = 0
    while curr_frame < total_frames - 1:
        video.set(cv2.CAP_PROP_POS_FRAMES, curr_frame)
        success, frame = video.read()
        if not success:
            break
        _, buffer = cv2.imencode(".jpg", frame)
        base64Frames.append(base64.b64encode(buffer).decode("utf-8"))
        curr_frame += frames_to_skip
    video.release()
    audio_path = f"{base_video_path}.mp3"
    try:
        clip = VideoFileClip(video_path)
        clip.audio.write_audiofile(audio_path, bitrate="32k")
        clip.audio.close()
        clip.close()
    except:
        logger.info("No audio track found in video.")
    return base64Frames, audio_path

def execute_code(code):
    buffer = io.StringIO()
    try:
        with redirect_stdout(buffer):
            exec(code, {}, {})
        return buffer.getvalue(), None
    except Exception as e:
        return None, str(e)
    finally:
        buffer.close()

# Sidebar
st.sidebar.subheader("Gallery Settings")
st.session_state['gallery_size'] = st.sidebar.slider("Gallery Size", 1, 10, st.session_state['gallery_size'], key="gallery_size_slider")

# Tabs
tabs = st.tabs(["Camera 📷", "Download 📥", "OCR 🔍", "Build 🌱", "Image Gen 🎨", "PDF 📄", "Image 🖼️", "Audio 🎵", "Video 🎥", "Code 🧑‍💻", "Gallery 📚"])
(tab_camera, tab_download, tab_ocr, tab_build, tab_imggen, tab_pdf, tab_image, tab_audio, tab_video, tab_code, tab_gallery) = tabs

with tab_camera:
    st.header("Camera Snap 📷")
    cols = st.columns(2)
    for i, cam_key in enumerate(["cam0", "cam1"]):
        with cols[i]:
            cam_img = st.camera_input(f"Take a picture - Cam {i}", key=cam_key)
            if cam_img:
                filename = generate_filename(f"cam{i}")
                with open(filename, "wb") as f:
                    f.write(cam_img.getvalue())
                st.session_state[f'cam{i}_file'] = filename
                st.session_state['history'].append(f"Snapshot from Cam {i}: {filename}")
                st.image(Image.open(filename), caption=f"Camera {i}", use_container_width=True)

with tab_download:
    st.header("Download PDFs 📥")
    url_input = st.text_area("Enter PDF URLs (one per line)", height=200)
    if st.button("Download 🤖"):
        urls = url_input.strip().split("\n")
        progress_bar = st.progress(0)
        for idx, url in enumerate(urls):
            if url:
                output_path = generate_filename(url, "pdf")
                if download_pdf(url, output_path):
                    st.session_state['downloaded_pdfs'][url] = output_path
                    st.session_state['history'].append(f"Downloaded PDF: {output_path}")
                    st.session_state['asset_checkboxes'][output_path] = True
                progress_bar.progress((idx + 1) / len(urls))

with tab_ocr:
    st.header("Test OCR 🔍")
    all_files = get_gallery_files()
    if all_files:
        selected_file = st.selectbox("Select File", all_files, key="ocr_select")
        if selected_file and st.button("Run OCR 🚀"):
            if selected_file.endswith('.png'):
                image = Image.open(selected_file)
            else:
                doc = fitz.open(selected_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("ocr_output", "txt")
            result = asyncio.run(process_ocr(image, output_file))
            st.text_area("OCR Result", result, height=200)
            st.session_state['history'].append(f"OCR Test: {selected_file} -> {output_file}")

with tab_build:
    st.header("Build Titan 🌱")
    model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
    base_model = st.selectbox("Select Model", ["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"] if model_type == "Causal LM" else ["OFA-Sys/small-stable-diffusion-v0", "stabilityai/stable-diffusion-2-base"])
    model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
    if st.button("Download Model ⬇️"):
        config = (ModelConfig 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()
        builder.load_model(base_model, config)
        builder.save_model(config.model_path)
        st.session_state['builder'] = builder
        st.session_state['model_loaded'] = True

with tab_imggen:
    st.header("Test Image Gen 🎨")
    prompt = st.text_area("Prompt", "Generate a futuristic cityscape")
    if st.button("Run Image Gen 🚀"):
        output_file = generate_filename("gen_output", "png")
        result = asyncio.run(process_image_gen(prompt, output_file))
        st.image(result, caption="Generated Image", use_container_width=True)
        st.session_state['history'].append(f"Image Gen Test: {prompt} -> {output_file}")

with tab_pdf:
    st.header("PDF Process 📄")
    uploaded_pdfs = st.file_uploader("Upload PDFs", type=["pdf"], accept_multiple_files=True)
    view_mode = st.selectbox("View Mode", ["Single Page", "Double Page"], key="pdf_view_mode")
    if st.button("Process PDFs"):
        for pdf_file in uploaded_pdfs:
            pdf_path = generate_filename(pdf_file.name, "pdf")
            with open(pdf_path, "wb") as f:
                f.write(pdf_file.read())
            snapshots = asyncio.run(process_pdf_snapshot(pdf_path, "double" if view_mode == "Double Page" else "single"))
            for snapshot in snapshots:
                st.image(Image.open(snapshot), caption=snapshot)
                text = process_image_with_prompt(Image.open(snapshot), "Extract the electronic text from image")
                st.text_area(f"Extracted Text from {snapshot}", text)
                code_prompt = f"Generate Python code based on this text:\n\n{text}"
                code = process_text_with_prompt(text, code_prompt)
                st.code(code, language="python")
                if st.button(f"Execute Code from {snapshot}"):
                    output, error = execute_code(code)
                    if error:
                        st.error(f"Error: {error}")
                    else:
                        st.success(f"Output: {output or 'No output'}")

with tab_image:
    st.header("Image Process 🖼️")
    uploaded_images = st.file_uploader("Upload Images", type=["png", "jpg"], accept_multiple_files=True)
    prompt = st.text_input("Prompt", "Extract the electronic text from image")
    if st.button("Process Images"):
        for img_file in uploaded_images:
            img = Image.open(img_file)
            st.image(img, caption=img_file.name)
            result = process_image_with_prompt(img, prompt)
            st.text_area(f"Result for {img_file.name}", result)

with tab_audio:
    st.header("Audio Process 🎵")
    audio_bytes = audio_recorder()
    if audio_bytes:
        filename = generate_filename("recording", "wav")
        with open(filename, "wb") as f:
            f.write(audio_bytes)
        st.audio(filename)
        transcript, summary = process_audio(filename, "Summarize this audio in markdown")
        st.text_area("Transcript", transcript)
        st.markdown(summary)

with tab_video:
    st.header("Video Process 🎥")
    video_input = st.file_uploader("Upload Video", type=["mp4"])
    if video_input:
        video_path = generate_filename(video_input.name, "mp4")
        with open(video_path, "wb") as f:
            f.write(video_input.read())
        st.video(video_path)
        result = process_video(video_path, "Summarize this video in markdown")
        st.markdown(result)

with tab_code:
    st.header("Code Executor 🧑‍💻")
    code_input = st.text_area("Python Code", height=400)
    if st.button("Run Code"):
        output, error = execute_code(code_input)
        if error:
            st.error(f"Error: {error}")
        else:
            st.success(f"Output: {output or 'No output'}")

with tab_gallery:
    st.header("Gallery 📚")
    all_files = get_gallery_files()
    for file in all_files:
        if file.endswith('.png'):
            st.image(Image.open(file), caption=file)
        elif file.endswith('.pdf'):
            doc = fitz.open(file)
            pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
            st.image(Image.frombytes("RGB", [pix.width, pix.height], pix.samples), caption=file)
            doc.close()
        elif file.endswith('.md'):
            with open(file, "r") as f:
                st.markdown(f.read())
        elif file.endswith('.wav'):
            st.audio(file)
        elif file.endswith('.mp4'):
            st.video(file)

# Update gallery in sidebar
def update_gallery():
    container = st.session_state['asset_gallery_container']
    container.empty()
    all_files = get_gallery_files()
    if all_files:
        container.markdown("### Asset Gallery 📸📖")
        cols = container.columns(2)
        for idx, file in enumerate(all_files[:st.session_state['gallery_size']]):
            with cols[idx % 2]:
                if file.endswith('.png'):
                    st.image(Image.open(file), caption=os.path.basename(file))
                elif file.endswith('.pdf'):
                    doc = fitz.open(file)
                    pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
                    st.image(Image.frombytes("RGB", [pix.width, pix.height], pix.samples), caption=os.path.basename(file))
                    doc.close()
                st.checkbox("Select", key=f"asset_{file}", value=st.session_state['asset_checkboxes'].get(file, False))
                st.markdown(get_download_link(file, "application/octet-stream", "Download"), unsafe_allow_html=True)
                if st.button("Delete", key=f"delete_{file}"):
                    os.remove(file)
                    st.session_state['asset_checkboxes'].pop(file, None)
                    st.experimental_rerun()

update_gallery()

# Sidebar logs and history
st.sidebar.subheader("Action Logs 📜")
for record in log_records:
    st.sidebar.write(f"{record.asctime} - {record.levelname} - {record.message}")
st.sidebar.subheader("History 📜")
for entry in st.session_state.get("history", []):
    if entry:
        st.sidebar.write(entry)