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

from dataclasses import dataclass
from datetime import datetime
from diffusers import StableDiffusionPipeline
from io import BytesIO
from openai import OpenAI
from PIL import Image
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
from typing import Optional

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

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="AI Vision & SFT 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': "AI Vision & SFT Titans: PDFs, OCR, Image Gen, Line Drawings, Custom Diffusion, and SFT on CPU! 🌌"
    }
)

st.session_state.setdefault('history', [])
st.session_state.setdefault('builder', None)
st.session_state.setdefault('model_loaded', False)
st.session_state.setdefault('processing', {})
st.session_state.setdefault('asset_checkboxes', {})
st.session_state.setdefault('downloaded_pdfs', {})
st.session_state.setdefault('unique_counter', 0)
st.session_state.setdefault('selected_model_type', "Causal LM")
st.session_state.setdefault('selected_model', "None")
st.session_state.setdefault('cam0_file', None)
st.session_state.setdefault('cam1_file', None)
st.session_state.setdefault('characters', [])
st.session_state.setdefault('char_form_reset', False)
if 'asset_gallery_container' not in st.session_state:
    st.session_state['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
        self.jokes = [
            "Why did the AI go to therapy? Too many layers to unpack! 😂",
            "Training complete! Time for a binary coffee break. ☕",
            "I told my neural network a joke; it couldn't stop dropping bits! 🤖",
            "I asked the AI for a pun, and it said, 'I'm punning on parallel processing!' 😄",
            "Debugging my code is like a stand-up routine—always a series of exceptions! 😆"
        ]
    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! 🎉 {random.choice(self.jokes)}")
        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(sequence, ext="png"):
    return f"{sequence}_{time.strftime('%d%m%Y%H%M%S')}.{ext}"

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

def get_download_link(file_path, mime_type="application/pdf", label="Download"):
    return f'<a href="data:{mime_type};base64,{base64.b64encode(open(file_path, "rb").read()).decode()}" 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:
        [zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.dirname(directory_path)))
         for root, _, files in os.walk(directory_path) for file in files]

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

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

def get_pdf_files():
    return sorted(glob.glob("*.pdf"))

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)
            ret = True
        else:
            ret = False
    except requests.RequestException as e:
        logger.error(f"Failed to download {url}: {e}")
        ret = False
    return ret

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 == "twopage":
            for i in range(min(2, len(doc))):
                page = doc[i]
                pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                output_file = generate_filename(f"twopage_{i}", "png")
                pix.save(output_file)
                output_files.append(output_file)
        elif mode == "allpages":
            for i in range(len(doc)):
                page = doc[i]
                pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                output_file = generate_filename(f"page_{i}", "png")
                pix.save(output_file)
                output_files.append(output_file)
        doc.close()
        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_gpt4o_ocr(image, output_file):
    start_time = time.time()
    status = st.empty()
    status.text("Processing GPT-4o OCR... (0s)")
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    messages = [{
        "role": "user",
        "content": [
            {"type": "text", "text": "Extract electronic text, and explain"},
            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}", "detail": "auto"}}
        ]
    }]
    try:
        response = client.chat.completions.create(model="gpt-4o", messages=messages, max_tokens=300)
        result = response.choices[0].message.content
        elapsed = int(time.time() - start_time)
        status.text(f"GPT-4o OCR completed in {elapsed}s!")
        async with aiofiles.open(output_file, "w") as f:
            await f.write(result)
        return result
    except Exception as e:
        status.error(f"Failed to process image with GPT-4o: {str(e)}")
        return ""

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) 
                and st.session_state['builder'].pipeline 
                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 randomize_character_content():
    intro_templates = [
        "{char} is a valiant knight who is silent and reserved, he looks handsome but aloof.",
        "{char} is a mischievous thief with a heart of gold, always sneaking around but helping those in need.",
        "{char} is a wise scholar who loves books more than people, often lost in thought.",
        "{char} is a fiery warrior with a short temper, but fiercely loyal to friends.",
        "{char} is a gentle healer who speaks softly, always carrying herbs and a warm smile."
    ]
    greeting_templates = [
        "You were startled by the sudden intrusion of a man into your home. 'I am from the knight's guild, and I have been ordered to arrest you.'",
        "A shadowy figure steps into the light. 'I heard you needed help—name’s {char}, best thief in town.'",
        "A voice calls from behind a stack of books. 'Oh, hello! I’m {char}, didn’t see you there—too many scrolls!'",
        "A booming voice echoes, 'I’m {char}, and I’m here to fight for justice—or at least a good brawl!'",
        "A soft hand touches your shoulder. 'I’m {char}, here to heal your wounds—don’t worry, I’ve got you.'"
    ]
    name = f"Character_{random.randint(1000, 9999)}"
    gender = random.choice(["Male", "Female"])
    intro = random.choice(intro_templates).format(char=name)
    greeting = random.choice(greeting_templates).format(char=name)
    return name, gender, intro, greeting

def save_character(character_data):
    characters = st.session_state.get('characters', [])
    characters.append(character_data)
    st.session_state['characters'] = characters
    with open("characters.json", "w") as f:
        json.dump(characters, f)

def load_characters():
    try:
        with open("characters.json", "r") as f:
            characters = json.load(f)
            st.session_state['characters'] = characters
    except FileNotFoundError:
        st.session_state['characters'] = []

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

tabs = st.tabs([
    "Camera Snap 📷", "Test OCR 🔍", "MD Gallery 📚", "Download PDFs 📥",  "Build Titan 🌱", 
    "Test Image Gen 🎨", "PDF Process 📄", "Image Process 🖼️", 
    "Character Editor 🧑‍🎨", "Character Gallery 🖼️"
])
(tab_camera, tab_ocr, tab_md_gallery, tab_download, tab_build, tab_imggen, tab_pdf_process, tab_image_process, tab_character_editor, tab_character_gallery) = tabs

with tab_camera:
    st.header("Camera Snap 📷")
    st.subheader("Single Capture")
    cols = st.columns(2)
    with cols[0]:
        cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0")
        if cam0_img:
            filename = generate_filename("cam0")
            if st.session_state['cam0_file'] and os.path.exists(st.session_state['cam0_file']):
                os.remove(st.session_state['cam0_file'])
            with open(filename, "wb") as f:
                f.write(cam0_img.getvalue())
            st.session_state['cam0_file'] = filename
            entry = f"Snapshot from Cam 0: {filename}"
            st.session_state['history'].append(entry)
            st.image(Image.open(filename), caption="Camera 0", use_container_width=True)
            logger.info(f"Saved snapshot from Camera 0: {filename}")
    with cols[1]:
        cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
        if cam1_img:
            filename = generate_filename("cam1")
            if st.session_state['cam1_file'] and os.path.exists(st.session_state['cam1_file']):
                os.remove(st.session_state['cam1_file'])
            with open(filename, "wb") as f:
                f.write(cam1_img.getvalue())
            st.session_state['cam1_file'] = filename
            entry = f"Snapshot from Cam 1: {filename}"
            st.session_state['history'].append(entry)
            st.image(Image.open(filename), caption="Camera 1", use_container_width=True)
            logger.info(f"Saved snapshot from Camera 1: {filename}")

with tab_download:
    st.header("Download PDFs 📥")
    if st.button("Examples 📚"):
        example_urls = [
            "https://arxiv.org/pdf/2308.03892",
            "https://arxiv.org/pdf/1912.01703",
            "https://arxiv.org/pdf/2408.11039",
            "https://arxiv.org/pdf/2109.10282",
            "https://arxiv.org/pdf/2112.10752",
            "https://arxiv.org/pdf/2308.11236",
            "https://arxiv.org/pdf/1706.03762",
            "https://arxiv.org/pdf/2006.11239",
            "https://arxiv.org/pdf/2305.11207",
            "https://arxiv.org/pdf/2106.09685",
            "https://arxiv.org/pdf/2005.11401",
            "https://arxiv.org/pdf/2106.10504"
        ]
        st.session_state['pdf_urls'] = "\n".join(example_urls)
    url_input = st.text_area("Enter PDF URLs (one per line)", value=st.session_state.get('pdf_urls', ""), height=200)
    if st.button("Robo-Download 🤖"):
        urls = url_input.strip().split("\n")
        progress_bar = st.progress(0)
        status_text = st.empty()
        total_urls = len(urls)
        existing_pdfs = get_pdf_files()
        for idx, url in enumerate(urls):
            if url:
                output_path = pdf_url_to_filename(url)
                status_text.text(f"Fetching {idx + 1}/{total_urls}: {os.path.basename(output_path)}...")
                if output_path not in existing_pdfs:
                    if download_pdf(url, output_path):
                        st.session_state['downloaded_pdfs'][url] = output_path
                        logger.info(f"Downloaded PDF from {url} to {output_path}")
                        entry = f"Downloaded PDF: {output_path}"
                        st.session_state['history'].append(entry)
                        st.session_state['asset_checkboxes'][output_path] = True
                    else:
                        st.error(f"Failed to nab {url} 😿")
                else:
                    st.info(f"Already got {os.path.basename(output_path)}! Skipping... 🐾")
                    st.session_state['downloaded_pdfs'][url] = output_path
                progress_bar.progress((idx + 1) / total_urls)
        status_text.text("Robo-Download complete! 🚀")
    mode = st.selectbox("Snapshot Mode", ["Single Page (High-Res)", "Two Pages (High-Res)", "All Pages (High-Res)"], key="download_mode")
    if st.button("Snapshot Selected 📸"):
        selected_pdfs = [path for path in get_gallery_files() if path.endswith('.pdf') and st.session_state['asset_checkboxes'].get(path, False)]
        if selected_pdfs:
            for pdf_path in selected_pdfs:
                if not os.path.exists(pdf_path):
                    st.warning(f"File not found: {pdf_path}. Skipping.")
                    continue
                mode_key = {"Single Page (High-Res)": "single",
                            "Two Pages (High-Res)": "twopage",
                            "All Pages (High-Res)": "allpages"}[mode]
                snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key))
                for snapshot in snapshots:
                    st.image(Image.open(snapshot), caption=snapshot, use_container_width=True)
                    st.session_state['asset_checkboxes'][snapshot] = True
        else:
            st.warning("No PDFs selected for snapshotting! Check some boxes in the sidebar.")

with tab_ocr:
    st.header("Test OCR 🔍")
    all_files = get_gallery_files()
    if all_files:
        if st.button("OCR All Assets 🚀"):
            full_text = "# OCR Results (GPT-4o)\n\n"
            for file in all_files:
                if file.endswith('.png'):
                    image = Image.open(file)
                else:
                    doc = fitz.open(file)
                    pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                    image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                    doc.close()
                output_file = generate_filename(f"ocr_{os.path.basename(file)}", "txt")
                result = asyncio.run(process_gpt4o_ocr(image, output_file))
                full_text += f"## {os.path.basename(file)}\n\n{result}\n\n"
                entry = f"OCR Test: {file} -> {output_file}"
                st.session_state['history'].append(entry)
            md_output_file = f"full_ocr_{int(time.time())}.md"
            with open(md_output_file, "w") as f:
                f.write(full_text)
            st.success(f"Full OCR saved to {md_output_file}")
            st.markdown(get_download_link(md_output_file, "text/markdown", "Download Full OCR Markdown"), unsafe_allow_html=True)
        selected_file = st.selectbox("Select Image or PDF", all_files, key="ocr_select")
        if selected_file:
            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()
            st.image(image, caption="Input Image", use_container_width=True)
            if st.button("Run OCR 🚀", key="ocr_run"):
                output_file = generate_filename("ocr_output", "md")
                st.session_state['processing']['ocr'] = True
                
                result = asyncio.run(process_gpt4o_ocr(image, output_file))
                
                entry = f"OCR Test: {selected_file} -> {output_file}"
                st.session_state['history'].append(entry)
                #st.text_area("OCR Result", result, height=200, key="ocr_result")
                st.code(result, language="python", line_numbers=True, wrap_lines=True, height=200)
                # 03312025 5:55AM                
                if len(result) > 50: open(output_file, "w").write(result); st.success(f"OCR output saved to {output_file}")
                else: st.warning("OCR output too short; file not saved.")
                st.session_state['processing']['ocr'] = False

            if selected_file.endswith('.pdf') and st.button("OCR All Pages 🚀", key="ocr_all_pages"):
                doc = fitz.open(selected_file)
                full_text = f"# OCR Results for {os.path.basename(selected_file)}\n\n"
                for i in range(len(doc)):
                    pix = doc[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                    image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                    output_file = generate_filename(f"ocr_page_{i}", "txt")
                    result = asyncio.run(process_gpt4o_ocr(image, output_file))
                    full_text += f"## Page {i + 1}\n\n{result}\n\n"
                    entry = f"OCR Test: {selected_file} Page {i + 1} -> {output_file}"
                    st.session_state['history'].append(entry)
                md_output_file = f"full_ocr_{os.path.basename(selected_file)}_{int(time.time())}.md"
                with open(md_output_file, "w") as f:
                    f.write(full_text)
                st.success(f"Full OCR saved to {md_output_file}")
                st.markdown(get_download_link(md_output_file, "text/markdown", "Download Full OCR Markdown"), unsafe_allow_html=True)
    else:
        st.warning("No assets in gallery yet. Use Camera Snap or Download PDFs!")

with tab_build:
    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", "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())}")
    domain = st.text_input("Target Domain", "general")
    if st.button("Download Model ⬇️"):
        config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(
            name=model_name, base_model=base_model, size="small", domain=domain
        )
        builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
        builder.load_model(base_model, config)
        builder.save_model(config.model_path)
        st.session_state['builder'] = builder
        st.session_state['model_loaded'] = True
        st.session_state['selected_model_type'] = model_type
        st.session_state['selected_model'] = config.model_path
        entry = f"Built {model_type} model: {model_name}"
        st.session_state['history'].append(entry)
        st.success(f"Model downloaded and saved to {config.model_path}! 🎉")
        st.rerun()

with tab_imggen:
    st.header("Test Image Gen 🎨")
    all_files = get_gallery_files()
    if all_files:
        selected_file = st.selectbox("Select Image or PDF", all_files, key="gen_select")
        if selected_file:
            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()
            st.image(image, caption="Reference Image", use_container_width=True)
            prompt = st.text_area("Prompt", "Generate a neon superhero version of this image", key="gen_prompt")
            if st.button("Run Image Gen 🚀", key="gen_run"):
                output_file = generate_filename("gen_output", "png")
                st.session_state['processing']['gen'] = True
                result = asyncio.run(process_image_gen(prompt, output_file))
                entry = f"Image Gen Test: {prompt} -> {output_file}"
                st.session_state['history'].append(entry)
                st.image(result, caption="Generated Image", use_container_width=True)
                st.success(f"Image saved to {output_file}")
                st.session_state['processing']['gen'] = False
    else:
        st.warning("No images or PDFs in gallery yet. Use Camera Snap or Download PDFs!")

with tab_pdf_process:
    st.header("PDF Process")
    st.subheader("Upload PDFs for GPT-based text extraction")
    gpt_models = ["gpt-4o", "gpt-4o-mini"]
    selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="pdf_gpt_model")
    detail_level = st.selectbox("Detail Level", ["auto", "low", "high"], key="pdf_detail_level")
    uploaded_pdfs = st.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True, key="pdf_process_uploader")
    view_mode = st.selectbox("View Mode", ["Single Page", "Double Page"], key="pdf_view_mode")
    if st.button("Process Uploaded PDFs", key="process_pdfs"):
        combined_text = ""
        for pdf_file in uploaded_pdfs:
            pdf_bytes = pdf_file.read()
            temp_pdf_path = f"temp_{pdf_file.name}"
            with open(temp_pdf_path, "wb") as f:
                f.write(pdf_bytes)
            try:
                doc = fitz.open(temp_pdf_path)
                st.write(f"Processing {pdf_file.name} with {len(doc)} pages")
                if view_mode == "Single Page":
                    for i, page in enumerate(doc):
                        pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                        st.image(img, caption=f"{pdf_file.name} Page {i+1}")
                        gpt_text = process_image_with_prompt(img, "Extract the electronic text from image", model=selected_gpt_model, detail=detail_level)
                        combined_text += f"\n## {pdf_file.name} - Page {i+1}\n\n{gpt_text}\n"
                else:
                    pages = list(doc)
                    for i in range(0, len(pages), 2):
                        if i+1 < len(pages):
                            pix1 = pages[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                            img1 = Image.frombytes("RGB", [pix1.width, pix1.height], pix1.samples)
                            pix2 = pages[i+1].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                            img2 = Image.frombytes("RGB", [pix2.width, pix2.height], pix2.samples)
                            total_width = img1.width + img2.width
                            max_height = max(img1.height, img2.height)
                            combined_img = Image.new("RGB", (total_width, max_height))
                            combined_img.paste(img1, (0, 0))
                            combined_img.paste(img2, (img1.width, 0))
                            st.image(combined_img, caption=f"{pdf_file.name} Pages {i+1}-{i+2}")
                            gpt_text = process_image_with_prompt(combined_img, "Extract the electronic text from image", model=selected_gpt_model,                             detail=detail_level)
                            combined_text += f"\n## {pdf_file.name} - Pages {i+1}-{i+2}\n\n{gpt_text}\n"
                        else:
                            pix = pages[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                            st.image(img, caption=f"{pdf_file.name} Page {i+1}")
                            gpt_text = process_image_with_prompt(img, "Extract the electronic text from image", model=selected_gpt_model, detail=detail_level)
                            combined_text += f"\n## {pdf_file.name} - Page {i+1}\n\n{gpt_text}\n"
                doc.close()
            except Exception as e:
                st.error(f"Error processing {pdf_file.name}: {str(e)}")
            finally:
                os.remove(temp_pdf_path)
        output_filename = generate_filename("processed_pdf", "md")
        with open(output_filename, "w", encoding="utf-8") as f:
            f.write(combined_text)
        st.success(f"PDF processing complete. MD file saved as {output_filename}")
        st.markdown(get_download_link(output_filename, "text/markdown", "Download Processed PDF MD"), unsafe_allow_html=True)

with tab_image_process:
    st.header("Image Process")
    st.subheader("Upload Images for GPT-based OCR")
    gpt_models = ["gpt-4o", "gpt-4o-mini"]
    selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="img_gpt_model")
    detail_level = st.selectbox("Detail Level", ["auto", "low", "high"], key="img_detail_level")
    prompt_img = st.text_input("Enter prompt for image processing", "Extract the electronic text from image", key="img_process_prompt")
    uploaded_images = st.file_uploader("Upload image files", type=["png", "jpg", "jpeg"], accept_multiple_files=True, key="image_process_uploader")
    if st.button("Process Uploaded Images", key="process_images"):
        combined_text = ""
        for img_file in uploaded_images:
            try:
                img = Image.open(img_file)
                st.image(img, caption=img_file.name)
                gpt_text = process_image_with_prompt(img, prompt_img, model=selected_gpt_model, detail=detail_level)
                combined_text += f"\n## {img_file.name}\n\n{gpt_text}\n"
            except Exception as e:
                st.error(f"Error processing image {img_file.name}: {str(e)}")
        output_filename = generate_filename("processed_image", "md")
        with open(output_filename, "w", encoding="utf-8") as f:
            f.write(combined_text)
        st.success(f"Image processing complete. MD file saved as {output_filename}")
        st.markdown(get_download_link(output_filename, "text/markdown", "Download Processed Image MD"), unsafe_allow_html=True)

with tab_md_gallery:
    st.header("MD Gallery and GPT Processing")
    gpt_models = ["gpt-4o", "gpt-4o-mini"]
    selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="md_gpt_model")
    md_files = sorted(glob.glob("*.md"))
    if md_files:
        st.subheader("Individual File Processing")
        cols = st.columns(2)
        for idx, md_file in enumerate(md_files):
            with cols[idx % 2]:
                st.write(md_file)
                if st.button(f"Process {md_file}", key=f"process_md_{md_file}"):
                    try:
                        with open(md_file, "r", encoding="utf-8") as f:
                            content = f.read()
                        prompt_md = "Summarize this into markdown outline with emojis and number the topics 1..12"
                        result_text = process_text_with_prompt(content, prompt_md, model=selected_gpt_model)
                        st.markdown(result_text)
                        output_filename = generate_filename(f"processed_{os.path.splitext(md_file)[0]}", "md")
                        with open(output_filename, "w", encoding="utf-8") as f:
                            f.write(result_text)
                        st.markdown(get_download_link(output_filename, "text/markdown", f"Download {output_filename}"), unsafe_allow_html=True)
                    except Exception as e:
                        st.error(f"Error processing {md_file}: {str(e)}")
        st.subheader("Knowledge Processing")
        st.write("Select MD files to combine and process:")
        selected_md = {}
        for md_file in md_files:
            selected_md[md_file] = st.checkbox(md_file, key=f"checkbox_md_{md_file}")
        Knowledge_prompt = st.text_input("1. Change into a numbered markdown outline with emojis.  2. Number the topics and make sure to include all data.", key="Knowledge_prompt")
        if st.button("Process Selected MD Files", key="process_Knowledge_md"):
            combined_content = ""
            for md_file, selected in selected_md.items():
                if selected:
                    try:
                        with open(md_file, "r", encoding="utf-8") as f:
                            combined_content += f"\n## {md_file}\n" + f.read() + "\n"
                    except Exception as e:
                        st.error(f"Error reading {md_file}: {str(e)}")
            if combined_content:
                result_text = process_text_with_prompt(combined_content, Knowledge_prompt, model=selected_gpt_model)
                st.markdown(result_text)
                output_filename = generate_filename("Knowledge_processed_md", "md")
                with open(output_filename, "w", encoding="utf-8") as f:
                    f.write(result_text)
                st.success(f"Knowledge processing complete. MD file saved as {output_filename}")
                st.markdown(get_download_link(output_filename, "text/markdown", "Download Knowledge Processed MD"), unsafe_allow_html=True)
            else:
                st.warning("No MD files selected.")
    else:
        st.warning("No MD files found.")

with tab_character_editor:
    st.header("Character Editor 🧑‍🎨")
    st.subheader("Create Your Character")

    with st.form(key="character_form"):
        if st.session_state.get('char_form_reset', False):
            default_name, default_gender, default_intro, default_greeting = randomize_character_content()
            st.session_state['char_form_reset'] = False
        else:
            default_name, default_gender, default_intro, default_greeting = randomize_character_content()

        if st.form_submit_button("Randomize Content 🎲"):
            default_name, default_gender, default_intro, default_greeting = randomize_character_content()

        name = st.text_input("Name (3-25 characters, letters, numbers, underscore, hyphen, space only)", 
                             value=default_name, 
                             max_chars=25, 
                             key="char_name")
        gender = st.radio("Gender", ["Male", "Female"], 
                          index=["Male", "Female"].index(default_gender), 
                          key="char_gender")
        intro = st.text_area("Intro (Publicly seen)", 
                             value=default_intro, 
                             max_chars=300, 
                             key="char_intro")
        greeting = st.text_area("Greeting", 
                                value=default_greeting, 
                                max_chars=300, 
                                key="char_greeting")

        if st.form_submit_button("Create Character"):
            if not name or len(name) < 3:
                st.error("Name must be 3-25 characters long.")
            elif not re.match(r'^[a-zA-Z0-9 _-]+$', name):
                st.error("Name can only contain letters, numbers, underscores, hyphens, and spaces.")
            elif not intro or not greeting:
                st.error("Intro and Greeting cannot be empty.")
            else:
                character_data = {
                    "name": name,
                    "gender": gender,
                    "intro": intro,
                    "greeting": greeting,
                    "created_at": datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                    "tags": ["OC"]
                }
                save_character(character_data)
                st.success(f"Character '{name}' created successfully!")
                st.session_state['char_form_reset'] = True
                st.rerun()

with tab_character_gallery:
    st.header("Character Gallery 🖼️")
    load_characters()
    characters = st.session_state.get('characters', [])
    if characters:
        st.subheader("Your Characters")
        cols = st.columns(3)
        for idx, char in enumerate(characters):
            with cols[idx % 3]:
                st.markdown(f"**{char['name']}**")
                st.write(f"**Gender**: {char['gender']}")
                st.write(f"**Intro**: {char['intro']}")
                st.write(f"**Greeting**: {char['greeting']}")
                st.write(f"**Created**: {char['created_at']}")
                st.write(f"**Tags**: {', '.join(char['tags'])}")
                if st.button(f"Delete {char['name']}", key=f"delete_char_{idx}"):
                    characters.pop(idx)
                    st.session_state['characters'] = characters
                    with open("characters.json", "w") as f:
                        json.dump(characters, f)
                    st.rerun()
                st.markdown("---")
    else:
        st.warning("No characters created yet. Use the Character Editor to create one!")

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]:
                st.session_state['unique_counter'] += 1
                unique_id = st.session_state['unique_counter']
                if file.endswith('.png'):
                    st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True)
                else:
                    doc = fitz.open(file)
                    pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                    st.image(img, caption=os.path.basename(file), use_container_width=True)
                    doc.close()
                checkbox_key = f"asset_{file}_{unique_id}"
                st.session_state['asset_checkboxes'][file] = st.checkbox("Use for SFT/Input", value=st.session_state['asset_checkboxes'].get(file, False), key=checkbox_key)
                mime_type = "image/png" if file.endswith('.png') else "application/pdf"
                st.markdown(get_download_link(file, mime_type, "Snag It! 📥"), unsafe_allow_html=True)
                if st.button("Zap It! 🗑️", key=f"delete_{file}_{unique_id}"):
                    os.remove(file)
                    st.session_state['asset_checkboxes'].pop(file, None)
                    st.success(f"Asset {os.path.basename(file)} vaporized! 💨")
                    st.rerun()

update_gallery()

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 is not None:
        st.sidebar.write(entry)