app.py

# --- Combined Imports ------------------------------------
import io
import os
import re
import base64
import glob
import logging
import random
import shutil
import time
import zipfile
import json
import asyncio
import aiofiles

from datetime import datetime
from collections import Counter
from dataclasses import dataclass, field
from io import BytesIO
from typing import Optional, List, Dict, Any

import pandas as pd
import pytz
import streamlit as st
from PIL import Image, ImageDraw, UnidentifiedImageError # Added ImageDraw and UnidentifiedImageError
from reportlab.pdfgen import canvas
from reportlab.lib.utils import ImageReader
from reportlab.lib.pagesizes import letter # Default page size
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as PlatypusImage, PageBreak, Preformatted
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import inch
from reportlab.lib.enums import TA_CENTER, TA_LEFT # For text alignment
import fitz # PyMuPDF

# --- Hugging Face Imports ---
from huggingface_hub import InferenceClient, HfApi, list_models
from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError # Import specific exceptions

# --- App Configuration -----------------------------------
st.set_page_config(
    page_title="Vision & Layout Titans (HF) 🚀🖼️",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={
        'Get Help': 'https://huggingface.co/docs',
        'Report a Bug': None, # Replace with your bug report link if desired
        'About': "Combined App: PDF Layout Generator + Hugging Face Powered AI Tools 🌌"
    }
)


# Conditional imports for optional/heavy libraries
try:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModelForVision2Seq, AutoModelForImageToWaveform, pipeline
    # Add more AutoModel classes as needed for different tasks (Vision, OCR, etc.)
    _transformers_available = True
except ImportError:
    _transformers_available = False
    # Place warning inside main app area if sidebar isn't ready
    # st.sidebar.warning("AI/ML libraries (torch, transformers) not found. Local model features disabled.")

try:
    from diffusers import StableDiffusionPipeline
    _diffusers_available = True
except ImportError:
    _diffusers_available = False
    # Don't show warning if transformers also missing, handled above
    # if _transformers_available:
        # st.sidebar.warning("Diffusers library not found. Diffusion model features disabled.")


import requests # Keep requests import

# --- 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):
        # Limit stored logs to avoid memory issues
        if len(log_records) > 200:
             log_records.pop(0)
        log_records.append(record)
logger.addHandler(LogCaptureHandler())

# --- Environment Variables & Constants -------------------
HF_TOKEN = os.getenv("HF_TOKEN")
DEFAULT_PROVIDER = "hf-inference"
# Model List (curated, similar to Gradio example) - can be updated
FEATURED_MODELS_LIST = [
    "meta-llama/Meta-Llama-3.1-8B-Instruct", # Updated Llama model
    "mistralai/Mistral-7B-Instruct-v0.3",
    "google/gemma-2-9b-it", # Added Gemma 2
    "Qwen/Qwen2-7B-Instruct", # Added Qwen2
    "microsoft/Phi-3-mini-4k-instruct",
    "HuggingFaceH4/zephyr-7b-beta",
    "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO", # Larger Mixture of Experts
    # Add a smaller option
    "HuggingFaceTB/SmolLM-1.7B-Instruct"
]
# Add common vision models if planning local loading
VISION_MODELS_LIST = [
    "Salesforce/blip-image-captioning-large",
    "microsoft/trocr-large-handwritten", # OCR model
    "llava-hf/llava-1.5-7b-hf", # Vision Language Model
    "google/vit-base-patch16-224", # Basic Vision Transformer
]
DIFFUSION_MODELS_LIST = [
    "stabilityai/stable-diffusion-xl-base-1.0", # Common SDXL
    "runwayml/stable-diffusion-v1-5", # Classic SD 1.5
    "OFA-Sys/small-stable-diffusion-v0", # Tiny diffusion
]


# --- Session State Initialization (Combined & Updated) ---
# Combined PDF Generator specific (replaces layout specific)
st.session_state.setdefault('combined_pdf_sources', []) # List of dicts {'filepath': path, 'type': type}

# General App State
st.session_state.setdefault('history', [])
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('cam0_file', None)
st.session_state.setdefault('cam1_file', None)
st.session_state.setdefault('characters', [])
st.session_state.setdefault('char_form_reset_key', 0) # For character form reset
# Removed gallery_size state - no longer used
# st.session_state.setdefault('gallery_size', 10)

# --- Hugging Face & Local Model State ---
st.session_state.setdefault('hf_inference_client', None) # Store initialized client
st.session_state.setdefault('hf_provider', DEFAULT_PROVIDER)
st.session_state.setdefault('hf_custom_key', "")
st.session_state.setdefault('hf_selected_api_model', FEATURED_MODELS_LIST[0]) # Default API model
st.session_state.setdefault('hf_custom_api_model', "") # User override for API model

# Local Model Management
st.session_state.setdefault('local_models', {}) # Dict to store loaded models: {'path': {'model': obj, 'tokenizer/proc': obj, 'type': 'causal/vision/etc'}}
st.session_state.setdefault('selected_local_model_path', None) # Path of the currently active local model

# Inference Parameters (shared for API and local where applicable)
st.session_state.setdefault('gen_max_tokens', 512)
st.session_state.setdefault('gen_temperature', 0.7)
st.session_state.setdefault('gen_top_p', 0.95)
st.session_state.setdefault('gen_frequency_penalty', 0.0) # Corresponds to repetition_penalty=1.0
st.session_state.setdefault('gen_seed', -1) # -1 for random

# Removed asset_gallery_container - render directly in sidebar
# if 'asset_gallery_container' not in st.session_state:
#     st.session_state['asset_gallery_container'] = st.sidebar.empty()

# --- Dataclasses (Refined for Local Models) -------------
@dataclass
class LocalModelConfig:
    name: str                   # User-defined local name
    hf_id: str                  # Hugging Face model ID used for download
    model_type: str             # 'causal', 'vision', 'diffusion', 'ocr', etc.
    size_category: str = "unknown" # e.g., 'small', 'medium', 'large'
    domain: Optional[str] = None
    local_path: str = field(init=False) # Path where it's saved

    def __post_init__(self):
        # Define local path based on type and name
        type_folder = f"{self.model_type}_models"
        safe_name = re.sub(r'[^\w\-]+', '_', self.name) # Sanitize name for path
        self.local_path = os.path.join(type_folder, safe_name)

    def get_full_path(self):
        return os.path.abspath(self.local_path)

# (Keep DiffusionConfig if still using diffusers library separately)
@dataclass
class DiffusionConfig: # Kept for clarity in diffusion tab if needed
    name: str
    base_model: str
    size: str
    domain: Optional[str] = None
    @property
    def model_path(self):
        # Ensure diffusion models are saved in their own distinct top-level folder
        return f"diffusion_models/{re.sub(r'[^w-]+', '_', self.name)}"


# --- Helper Functions (Combined and refined) -------------
def generate_filename(sequence, ext="png"):
    timestamp = time.strftime('%Y%m%d_%H%M%S')
    safe_sequence = re.sub(r'[^\w\-]+', '_', str(sequence))
    return f"{safe_sequence}_{timestamp}.{ext}"

def pdf_url_to_filename(url):
    name = re.sub(r'^https?://', '', url)
    name = re.sub(r'[<>:"/\\|?*]', '_', name)
    return name[:100] + ".pdf" # Limit length

def get_download_link(file_path, mime_type="application/octet-stream", label="Download"):
    if not os.path.exists(file_path): return f"{label} (File not found)"
    try:
        with open(file_path, "rb") as f: file_bytes = f.read()
        b64 = base64.b64encode(file_bytes).decode()
        return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label}</a>'
    except Exception as e:
        logger.error(f"Error creating download link for {file_path}: {e}")
        return f"{label} (Error)"

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:
                file_path = os.path.join(root, file)
                zipf.write(file_path, os.path.relpath(file_path, os.path.dirname(directory_path)))

def get_local_model_paths(model_type="causal"):
    """Gets paths of locally saved models of a specific type."""
    pattern = f"{model_type}_models/*"
    dirs = [d for d in glob.glob(pattern) if os.path.isdir(d)]
    return dirs

def get_gallery_files(file_types=("png", "pdf", "jpg", "jpeg", "md", "txt")):
    """Gets all files with specified extensions in the current directory."""
    all_files = set()
    for ext in file_types:
        # Ensure the glob pattern correctly targets files in the script's directory
        all_files.update(glob.glob(f"./*.{ext.lower()}")) # Use ./* for current dir
        all_files.update(glob.glob(f"./*.{ext.upper()}"))
    # Convert to list and remove potential './' prefix for cleaner display
    return sorted([os.path.normpath(f) for f in all_files])

def get_pdf_files():
    # Use get_gallery_files to find PDFs
    return get_gallery_files(['pdf'])

def download_pdf(url, output_path):
    try:
        headers = {'User-Agent': 'Mozilla/5.0'}
        response = requests.get(url, stream=True, timeout=20, headers=headers)
        response.raise_for_status()
        with open(output_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192): f.write(chunk)
        logger.info(f"Successfully downloaded {url} to {output_path}")
        return True
    except requests.exceptions.RequestException as e:
        logger.error(f"Failed to download {url}: {e}")
        if os.path.exists(output_path):
            try:
                os.remove(output_path)
                logger.info(f"Removed partially downloaded file: {output_path}")
            except OSError as remove_error:
                logger.error(f"Error removing partial file {output_path}: {remove_error}")
            except Exception as general_remove_error:
                logger.error(f"General error removing partial file {output_path}: {general_remove_error}")
        return False
    except Exception as e:
        logger.error(f"An unexpected error occurred during download of {url}: {e}")
        if os.path.exists(output_path):
             try:
                  os.remove(output_path)
                  logger.info(f"Removed file after unexpected error: {output_path}")
             except OSError as remove_error:
                  logger.error(f"Error removing file after unexpected error {output_path}: {remove_error}")
             except Exception as general_remove_error:
                  logger.error(f"General error removing file after unexpected error {output_path}: {general_remove_error}")
        return False

async def process_pdf_snapshot(pdf_path, mode="single", resolution_factor=2.0):
    start_time = time.time()
    # Use a placeholder within the main app area for status during async operations
    status_placeholder = st.empty()
    status_placeholder.text(f"Processing PDF Snapshot ({mode}, Res: {resolution_factor}x)... (0s)")
    output_files = []
    try:
        doc = fitz.open(pdf_path)
        matrix = fitz.Matrix(resolution_factor, resolution_factor)
        num_pages_to_process = 0
        if mode == "single": num_pages_to_process = min(1, len(doc))
        elif mode == "twopage": num_pages_to_process = min(2, len(doc))
        elif mode == "allpages": num_pages_to_process = len(doc)

        for i in range(num_pages_to_process):
            page_start_time = time.time()
            page = doc.load_page(i) # Use load_page for efficiency
            pix = page.get_pixmap(matrix=matrix)
            base_name = os.path.splitext(os.path.basename(pdf_path))[0]
            output_file = generate_filename(f"{base_name}_pg{i+1}_{mode}", "png")

            # Ensure output path is valid before saving
            output_dir = os.path.dirname(output_file) or "."
            if not os.path.exists(output_dir): os.makedirs(output_dir)

            await asyncio.to_thread(pix.save, output_file)
            output_files.append(output_file)
            elapsed_page = int(time.time() - page_start_time)
            status_placeholder.text(f"Processing PDF Snapshot ({mode}, Res: {resolution_factor}x)... Page {i+1}/{num_pages_to_process} done ({elapsed_page}s)")
            await asyncio.sleep(0.01)

        doc.close()
        elapsed = int(time.time() - start_time)
        status_placeholder.success(f"PDF Snapshot ({mode}, {len(output_files)} files) completed in {elapsed}s!")
        return output_files
    except Exception as e:
        logger.error(f"Failed to process PDF snapshot for {pdf_path}: {e}", exc_info=True) # Add traceback
        status_placeholder.error(f"Failed to process PDF {os.path.basename(pdf_path)}: {e}")
        # Clean up any files created before the error
        for f in output_files:
            if os.path.exists(f):
                try: os.remove(f)
                except: pass
        return []


# --- HF Inference Client Management ---
def get_hf_client() -> Optional[InferenceClient]:
    """Gets or initializes the Hugging Face Inference Client based on session state."""
    provider = st.session_state.hf_provider
    custom_key = st.session_state.hf_custom_key.strip()
    token_to_use = custom_key if custom_key else HF_TOKEN

    if not token_to_use and provider != "hf-inference":
        # Don't show error here, let caller handle it if client is needed
        # st.error(f"Provider '{provider}' requires a Hugging Face API token...")
        return None
    if provider == "hf-inference" and not token_to_use:
         logger.warning("Using hf-inference provider without a token. Rate limits may apply.")
         token_to_use = None # Explicitly set to None for public inference API

    # Check if client needs re-initialization
    current_client = st.session_state.get('hf_inference_client')
    needs_reinit = True
    if current_client:
         # Compare provider and token status more carefully
         current_token = getattr(current_client, '_token', None) # Access internal token if exists
         current_provider = getattr(current_client, 'provider', None) # Access provider if exists

         token_matches = (token_to_use == current_token)
         provider_matches = (provider == current_provider)

         if token_matches and provider_matches:
              needs_reinit = False

    if needs_reinit:
        try:
            logger.info(f"Initializing InferenceClient for provider: {provider}. Token source: {'Custom Key' if custom_key else ('HF_TOKEN' if HF_TOKEN else 'None')}")
            st.session_state.hf_inference_client = InferenceClient(model=None, token=token_to_use, provider=provider) # Init without model initially
            # Store provider on client instance if possible (check InferenceClient structure or assume it's handled internally)
            setattr(st.session_state.hf_inference_client, 'provider', provider) # Explicitly store provider for re-init check
            setattr(st.session_state.hf_inference_client, '_token', token_to_use) # Explicitly store token for re-init check
            logger.info("InferenceClient initialized successfully.")
        except Exception as e:
            st.error(f"Failed to initialize Hugging Face client for provider {provider}: {e}")
            logger.error(f"InferenceClient initialization failed: {e}")
            st.session_state.hf_inference_client = None

    return st.session_state.hf_inference_client

# --- HF/Local Model Processing Functions ---
def process_text_hf(text: str, prompt: str, use_api: bool) -> str:
    """Processes text using either HF Inference API or a loaded local model."""
    status_placeholder = st.empty()
    start_time = time.time()
    result_text = ""

    params = {
        "max_new_tokens": st.session_state.gen_max_tokens,
        "temperature": st.session_state.gen_temperature,
        "top_p": st.session_state.gen_top_p,
        "repetition_penalty": st.session_state.gen_frequency_penalty, # Keep user value, adjust name below if needed
    }
    seed = st.session_state.gen_seed
    if seed != -1: params["seed"] = seed

    system_prompt = "You are a helpful assistant. Process the following text based on the user's request."
    full_prompt = f"{prompt}\n\n---\n\n{text}"
    messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": full_prompt}]

    if use_api:
        status_placeholder.info("Processing text using Hugging Face API...")
        client = get_hf_client()
        if not client: return "Error: Hugging Face client not configured/available."
        model_id = st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model
        if not model_id: return "Error: No Hugging Face API model specified."
        status_placeholder.info(f"Using API Model: {model_id}")
        try:
            # Ensure repetition_penalty is passed correctly if supported
            api_params = {
                "max_tokens": params['max_new_tokens'],
                "temperature": params['temperature'],
                "top_p": params['top_p'],
                "repetition_penalty": params.get('repetition_penalty') # Check if API uses this name
            }
            if 'seed' in params: api_params['seed'] = params['seed']

            response = client.chat_completion(model=model_id, messages=messages, **api_params)
            result_text = response.choices[0].message.content or ""
            logger.info(f"HF API text processing successful for model {model_id}.")
        except Exception as e:
            logger.error(f"HF API text processing failed for model {model_id}: {e}", exc_info=True)
            result_text = f"Error during Hugging Face API inference: {str(e)}"
    else:
        status_placeholder.info("Processing text using local model...")
        if not _transformers_available: return "Error: Transformers library not available."
        model_path = st.session_state.get('selected_local_model_path')
        if not model_path or model_path not in st.session_state.get('local_models', {}): return "Error: No suitable local model selected/loaded."
        local_model_data = st.session_state['local_models'][model_path]
        if local_model_data.get('type') != 'causal': return f"Error: Loaded model '{os.path.basename(model_path)}' is not a Causal LM."
        status_placeholder.info(f"Using Local Model: {os.path.basename(model_path)}")
        model = local_model_data.get('model')
        tokenizer = local_model_data.get('tokenizer')
        if not model or not tokenizer: return f"Error: Model/tokenizer not found for {os.path.basename(model_path)}."
        try:
            try: prompt_for_model = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            except: logger.warning(f"Chat template failed for {model_path}. Using basic format."); prompt_for_model = f"System: {system_prompt}\nUser: {full_prompt}\nAssistant:"
            inputs = tokenizer(prompt_for_model, return_tensors="pt", padding=True, truncation=True, max_length=2048).to(model.device) # Increased context slightly
            generate_params = {
                "max_new_tokens": params['max_new_tokens'],
                "temperature": params['temperature'],
                "top_p": params['top_p'],
                "repetition_penalty": params.get('repetition_penalty', 1.0),
                "do_sample": True if params['temperature'] > 0.01 else False, # Sample if temp > 0.01
                "pad_token_id": tokenizer.eos_token_id
            }
            with torch.no_grad(): outputs = model.generate(**inputs, **generate_params)
            input_length = inputs['input_ids'].shape[1]; generated_ids = outputs[0][input_length:]
            result_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
            logger.info(f"Local text processing successful for model {model_path}.")
        except Exception as e:
            logger.error(f"Local text processing failed for model {model_path}: {e}", exc_info=True)
            result_text = f"Error during local model inference: {str(e)}"

    elapsed = int(time.time() - start_time)
    status_placeholder.success(f"Text processing completed in {elapsed}s.")
    return result_text

def process_image_hf(image: Image.Image, prompt: str, use_api: bool) -> str:
    """Processes an image using either HF Inference API or a local model."""
    status_placeholder = st.empty()
    start_time = time.time()
    result_text = "[Image processing requires specific Vision model implementation]"

    if use_api:
        status_placeholder.info("Processing image using Hugging Face API (Image-to-Text)...")
        client = get_hf_client()
        if not client: return "Error: HF client not configured."
        buffered = BytesIO(); image.save(buffered, format="PNG"); img_bytes = buffered.getvalue()
        try:
            captioning_model_id = "Salesforce/blip-image-captioning-large" # Default captioner
            vqa_model_id = "llava-hf/llava-1.5-7b-hf" # Default VQA - MAY REQUIRE DIFFERENT CLIENT CALL
            # Decide whether to use captioning or VQA based on prompt? Simple approach: captioning.
            status_placeholder.info(f"Using API Image-to-Text Model: {captioning_model_id}")
            response_list = client.image_to_text(data=img_bytes, model=captioning_model_id)
            if response_list and 'generated_text' in response_list[0]:
                result_text = f"API Caption: {response_list[0]['generated_text']}\n(Prompt '{prompt}' likely ignored by this API endpoint)"
                logger.info(f"HF API image captioning successful for model {captioning_model_id}.")
            else: result_text = "Error: Unexpected response format from image-to-text API."; logger.warning(f"Unexpected API response: {response_list}")
        except Exception as e: logger.error(f"HF API image processing failed: {e}"); result_text = f"Error during HF API image inference: {str(e)}"
    else:
        status_placeholder.info("Processing image using local model...")
        if not _transformers_available: return "Error: Transformers library needed."
        model_path = st.session_state.get('selected_local_model_path')
        if not model_path or model_path not in st.session_state.get('local_models', {}): return "Error: No suitable local model selected/loaded."
        local_model_data = st.session_state['local_models'][model_path]
        model_type = local_model_data.get('type')
        if model_type not in ['vision', 'ocr']: return f"Error: Loaded model '{os.path.basename(model_path)}' is not a Vision/OCR type."
        status_placeholder.warning(f"Local {model_type} Model ({os.path.basename(model_path)}): Processing logic depends on specific model. Placeholder active.")
        # --- ADD SPECIFIC LOCAL VISION/OCR MODEL LOGIC HERE ---
        # This section needs code tailored to the loaded model's processor/generate methods
        # Example placeholder:
        processor = local_model_data.get('processor')
        model = local_model_data.get('model')
        if processor and model:
             result_text = f"[Local {model_type} model processing needs implementation for {os.path.basename(model_path)}. Prompt: '{prompt}']"
        else:
             result_text = f"Error: Missing model or processor for local {model_type} model {os.path.basename(model_path)}."
        # --- END OF PLACEHOLDER ---

    elapsed = int(time.time() - start_time)
    status_placeholder.success(f"Image processing attempt completed in {elapsed}s.")
    return result_text

async def process_hf_ocr(image: Image.Image, output_file: str, use_api: bool) -> str:
    """ Performs OCR using the process_image_hf function framework. """
    # Simple prompt for OCR task
    ocr_prompt = "Extract text content from this image."
    result = process_image_hf(image, ocr_prompt, use_api=use_api) # Pass use_api flag

    # Save the result if it looks like text (basic check)
    if result and not result.startswith("Error") and not result.startswith("["):
        try:
            async with aiofiles.open(output_file, "w", encoding='utf-8') as f:
                await f.write(result)
            logger.info(f"HF OCR result saved to {output_file}")
        except IOError as e:
             logger.error(f"Failed to save HF OCR output to {output_file}: {e}")
             result += f"\n[Error saving file: {e}]" # Append error to result if save fails

    # --- CORRECTED BLOCK ---
    elif os.path.exists(output_file):
        # Remove file if processing failed or was just a placeholder message
        try:
            os.remove(output_file)
        except OSError:
            # Log error or just ignore if removal fails
            logger.warning(f"Could not remove potentially empty/failed OCR file: {output_file}")
            pass # Ignore removal error
        except Exception as e_rem: # Catch any other error during removal
             logger.warning(f"Error removing OCR file {output_file}: {e_rem}")
             pass
    # --- END CORRECTION ---

    return result

# --- Character Functions (Keep from previous) -----------
def randomize_character_content():
    intro_templates = ["{char} is a valiant knight...", "{char} is a mischievous thief...", "{char} is a wise scholar...", "{char} is a fiery warrior...", "{char} is a gentle healer..."]
    greeting_templates = ["'I am from the knight's guild...'", "'I heard you needed help—name’s {char}...", "'Oh, hello! I’m {char}, didn’t see you there...'", "'I’m {char}, and I’m here to fight...'", "'I’m {char}, here to heal...'"]
    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', []);
    if any(c['name'] == character_data['name'] for c in characters): st.error(f"Character name '{character_data['name']}' already exists."); return False
    characters.append(character_data); st.session_state['characters'] = characters
    try:
        with open("characters.json", "w", encoding='utf-8') as f: json.dump(characters, f, indent=2); logger.info(f"Saved character: {character_data['name']}"); return True
    except IOError as e: logger.error(f"Failed to save characters.json: {e}"); st.error(f"Failed to save character file: {e}"); return False

def load_characters():
    if not os.path.exists("characters.json"): st.session_state['characters'] = []; return
    try:
        with open("characters.json", "r", encoding='utf-8') as f: characters = json.load(f)
        if isinstance(characters, list): st.session_state['characters'] = characters; logger.info(f"Loaded {len(characters)} characters.")
        else: st.session_state['characters'] = []; logger.warning("characters.json is not a list, resetting."); os.remove("characters.json")
    except (json.JSONDecodeError, IOError) as e:
        logger.error(f"Failed to load or decode characters.json: {e}"); st.error(f"Error loading character file: {e}. Starting fresh."); st.session_state['characters'] = []
        try: corrupt_filename = f"characters_corrupt_{int(time.time())}.json"; shutil.copy("characters.json", corrupt_filename); logger.info(f"Backed up corrupted character file to {corrupt_filename}"); os.remove("characters.json")
        except Exception as backup_e: logger.error(f"Could not backup corrupted character file: {backup_e}")

# --- Utility: Clean stems (Keep from previous) ----------
def clean_stem(fn: str) -> str:
    name = os.path.splitext(os.path.basename(fn))[0]; name = name.replace('-', ' ').replace('_', ' ')
    return name.strip().title()

# --- PDF Creation Functions ---
# Original image-only PDF function (might be removed or kept as an option)
def make_image_sized_pdf(sources):
    # ... (kept same as previous version for now) ...
    if not sources: st.warning("No image sources provided for PDF generation."); return None
    buf = io.BytesIO(); c = canvas.Canvas(buf, pagesize=letter)
    try:
        for idx, src in enumerate(sources, start=1):
            status_placeholder = st.empty(); status_placeholder.info(f"Adding page {idx}/{len(sources)}: {os.path.basename(str(src))}...")
            try:
                filename = f'page_{idx}'
                if isinstance(src, str):
                    if not os.path.exists(src): logger.warning(f"Image file not found: {src}. Skipping."); status_placeholder.warning(f"Skipping missing file: {os.path.basename(src)}"); continue
                    img_obj = Image.open(src); filename = os.path.basename(src)
                elif hasattr(src, 'name'): # Handle uploaded file object
                    src.seek(0); img_obj = Image.open(src); filename = getattr(src, 'name', f'uploaded_image_{idx}'); src.seek(0)
                else: continue # Skip unknown source type
                with img_obj:
                    iw, ih = img_obj.size
                    if iw <= 0 or ih <= 0: logger.warning(f"Invalid image dimensions ({iw}x{ih}) for {filename}. Skipping."); status_placeholder.warning(f"Skipping invalid image: {filename}"); continue
                    cap_h = 30; pw, ph = iw, ih + cap_h; c.setPageSize((pw, ph)); img_reader = ImageReader(img_obj)
                    c.drawImage(img_reader, 0, cap_h, width=iw, height=ih, preserveAspectRatio=True, anchor='c', mask='auto')
                    caption = clean_stem(filename); c.setFont('Helvetica', 12); c.setFillColorRGB(0, 0, 0); c.drawCentredString(pw / 2, cap_h / 2 + 3, caption)
                    c.setFont('Helvetica', 8); c.setFillColorRGB(0.5, 0.5, 0.5); c.drawRightString(pw - 10, 8, f"Page {idx}")
                    c.showPage(); status_placeholder.success(f"Added page {idx}/{len(sources)}: {filename}")
            except (IOError, OSError, UnidentifiedImageError) as img_err: logger.error(f"Error processing image {src}: {img_err}"); status_placeholder.error(f"Error adding page {idx}: {img_err}")
            except Exception as e: logger.error(f"Unexpected error adding page {idx} ({src}): {e}"); status_placeholder.error(f"Unexpected error on page {idx}: {e}")
        c.save(); buf.seek(0)
        if buf.getbuffer().nbytes < 100: st.error("PDF generation resulted in an empty file."); return None
        return buf.getvalue()
    except Exception as e: logger.error(f"Fatal error during PDF generation: {e}"); st.error(f"PDF Generation Failed: {e}"); return None

# --- NEW Combined PDF Generation Function ---
def make_combined_pdf(ordered_sources_info: List[Dict]) -> Optional[bytes]:
    if not ordered_sources_info:
        st.warning("No items selected for combined PDF generation.")
        return None

    buf = io.BytesIO()
    c = canvas.Canvas(buf, pagesize=letter)
    styles = getSampleStyleSheet()
    total_pages_generated = 0

    # Add page number function
    def draw_page_number(canvas, page_num, page_width, page_height):
        canvas.saveState()
        canvas.setFont('Helvetica', 8)
        canvas.setFillColorRGB(0.5, 0.5, 0.5)
        canvas.drawRightString(page_width - inch/2, inch/2, f"Page {page_num}")
        canvas.restoreState()

    for idx, item_info in enumerate(ordered_sources_info):
        filepath = item_info.get('filepath')
        file_type = item_info.get('type')
        filename = item_info.get('filename', f"item_{idx+1}")
        item_caption = clean_stem(filename)

        if not filepath: logger.warning(f"Skipping item {idx+1} due to missing filepath."); continue
        is_file_object = not isinstance(filepath, str)
        status_placeholder = st.empty()
        status_placeholder.info(f"Processing item {idx+1}/{len(ordered_sources_info)}: {filename} ({file_type})...")

        try:
            # --- IMAGE Processing ---
            if file_type == 'Image':
                if is_file_object: filepath.seek(0)
                try:
                    img_obj = Image.open(filepath)
                    with img_obj:
                        iw, ih = img_obj.size
                        if iw <= 0 or ih <= 0: raise ValueError("Invalid image dimensions")
                        cap_h = 30; pw, ph = iw, ih + cap_h
                        c.setPageSize((pw, ph)); img_reader = ImageReader(img_obj)
                        c.drawImage(img_reader, 0, cap_h, width=iw, height=ih, preserveAspectRatio=True, anchor='c', mask='auto')
                        c.setFont('Helvetica', 12); c.setFillColorRGB(0, 0, 0); c.drawCentredString(pw / 2, cap_h / 2 + 3, item_caption)
                        total_pages_generated += 1; draw_page_number(c, total_pages_generated, pw, ph)
                        c.showPage()
                finally:
                    if is_file_object: filepath.seek(0)

            # --- PDF Processing ---
            elif file_type == 'PDF':
                src_doc = None
                try:
                    if is_file_object: filepath.seek(0); pdf_bytes = filepath.read(); src_doc = fitz.open("pdf", pdf_bytes)
                    else: src_doc = fitz.open(filepath)
                    if len(src_doc) == 0: st.warning(f"Skipping empty PDF: {filename}"); continue
                    for i, page in enumerate(src_doc):
                        page_rect = page.rect; pw, ph = page_rect.width, page_rect.height
                        if pw <= 0 or ph <= 0: continue
                        c.setPageSize((pw, ph))
                        pix = page.get_pixmap(dpi=150) # Render as image
                        if pix.width > 0 and pix.height > 0:
                             img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); img_reader = ImageReader(img)
                             c.drawImage(img_reader, 0, 0, width=pw, height=ph)
                        else: c.setFont('Helvetica', 10); c.setFillColorRGB(1,0,0); c.drawCentredString(pw/2, ph/2, f"Failed to render page {i+1} preview")
                        overlay_text = f"{item_caption} (p{i+1})"; c.setFont('Helvetica', 8); c.setFillColorRGB(0, 0, 0, alpha=0.6); c.drawString(10, 10, overlay_text)
                        total_pages_generated += 1; draw_page_number(c, total_pages_generated, pw, ph)
                        c.showPage()
                finally:
                     if src_doc: src_doc.close()
                     if is_file_object: filepath.seek(0)

            # --- TEXT/MARKDOWN Processing ---
            elif file_type == 'Text':
                if is_file_object:
                    filepath.seek(0)
                    try: text_content = filepath.read().decode('utf-8')
                    except: text_content = filepath.read().decode('latin-1', errors='replace')
                else:
                    with open(filepath, 'r', encoding='utf-8', errors='ignore') as f: text_content = f.read()

                temp_buf = io.BytesIO()
                temp_doc = SimpleDocTemplate(temp_buf, pagesize=letter, leftMargin=inch, rightMargin=inch, topMargin=inch, bottomMargin=inch)
                story = [Paragraph(f"Content from: {item_caption}", styles['h2']), Spacer(1, 0.2*inch)]
                # Use Preformatted for simple text dump
                story.append(Preformatted(text_content, styles['Code']))
                temp_doc.build(story)
                temp_buf.seek(0)

                text_pdf = fitz.open("pdf", temp_buf.read())
                for i, page in enumerate(text_pdf):
                     page_rect = page.rect; pw, ph = page_rect.width, page_rect.height
                     c.setPageSize((pw, ph)); pix = page.get_pixmap(dpi=150)
                     if pix.width > 0 and pix.height > 0:
                          img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); img_reader = ImageReader(img)
                          c.drawImage(img_reader, 0, 0, width=pw, height=ph)
                     else: c.setFont('Helvetica', 10); c.setFillColorRGB(1,0,0); c.drawCentredString(pw/2, ph/2, f"Failed to render text page {i+1}")
                     total_pages_generated += 1; draw_page_number(c, total_pages_generated, pw, ph)
                     c.showPage()
                text_pdf.close()

            else: # Unknown type
                logger.warning(f"Unsupported file type for PDF combination: {filename} ({file_type})")
                c.setPageSize(letter); c.setFont('Helvetica-Bold', 14); c.setFillColorRGB(0.7, 0.7, 0); c.drawCentredString(letter[0] / 2, letter[1] / 2 + 20, f"Unsupported File: {filename}")
                c.setFont('Helvetica', 10); c.drawCentredString(letter[0] / 2, letter[1] / 2 - 20, f"Type: {file_type}. Cannot include.")
                total_pages_generated += 1; draw_page_number(c, total_pages_generated, letter[0], letter[1])
                c.showPage()

        except Exception as item_err:
            logger.error(f"Error processing item {filename} for PDF: {item_err}", exc_info=True)
            try: # Add error page
                 c.setPageSize(letter); c.setFont('Helvetica-Bold', 14); c.setFillColorRGB(1, 0, 0); c.drawCentredString(letter[0] / 2, letter[1] / 2 + 20, f"Error processing: {filename}")
                 c.setFont('Helvetica', 10); c.drawCentredString(letter[0] / 2, letter[1] / 2 - 20, f"{str(item_err)[:100]}"); total_pages_generated += 1; draw_page_number(c, total_pages_generated, letter[0], letter[1]); c.showPage()
            except: logger.error(f"Failed to add error page for {filename}")
        finally:
             status_placeholder.empty()

    if total_pages_generated == 0: st.error("No pages were successfully added."); return None
    try:
         c.save(); buf.seek(0)
         if buf.getbuffer().nbytes < 100: st.error("Combined PDF generation resulted empty."); return None
         return buf.getvalue()
    except Exception as e: logger.error(f"Fatal error during final PDF save: {e}"); st.error(f"PDF Save Failed: {e}"); return None


# --- Sidebar Gallery Update Function (MODIFIED for Sort, PDF Preview Fix, Delete Fix) ---
def get_sort_key(filename):
    ext = os.path.splitext(filename)[1].lower()
    if ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff']: priority = 1
    elif ext in ['.md', '.txt']: priority = 2
    elif ext == '.pdf': priority = 3
    else: priority = 4
    return (priority, filename.lower())

def update_gallery():
    st.sidebar.markdown("### Asset Gallery 📸📖")
    all_files_unsorted = get_gallery_files()
    all_files = sorted(all_files_unsorted, key=get_sort_key) # Apply sorting

    if not all_files: st.sidebar.info("No assets found."); return
    st.sidebar.caption(f"Found {len(all_files)} assets:")

    for idx, file in enumerate(all_files):
        st.session_state['unique_counter'] += 1
        unique_id = st.session_state['unique_counter']
        item_key_base = f"gallery_item_{os.path.basename(file)}_{unique_id}"
        basename = os.path.basename(file)
        st.sidebar.markdown(f"**{basename}**")

        try:
            file_ext = os.path.splitext(file)[1].lower()
            preview_failed = False
            # Previews with better error handling
            if file_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff']:
                 try:
                     with st.sidebar.expander("Preview", expanded=False): st.image(Image.open(file), use_container_width=True)
                 except Exception as img_err: st.sidebar.warning(f"Img preview failed: {img_err}"); preview_failed = True
            elif file_ext == '.pdf':
                 try:
                     with st.sidebar.expander("Preview (Page 1)", expanded=False):
                          doc = fitz.open(file)
                          if len(doc) > 0:
                              pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
                              if pix.width > 0 and pix.height > 0: img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); st.image(img, use_container_width=True)
                              else: st.warning("Failed to render PDF page."); preview_failed = True
                          else: st.warning("Empty PDF")
                          doc.close()
                 except Exception as pdf_err: st.sidebar.warning(f"PDF preview failed: {pdf_err}"); logger.warning(f"PDF preview error {file}: {pdf_err}"); preview_failed = True
            elif file_ext in ['.md', '.txt']:
                 try:
                     with st.sidebar.expander("Preview (Start)", expanded=False):
                          with open(file, 'r', encoding='utf-8', errors='ignore') as f: content_preview = f.read(200)
                          st.code(content_preview + "...", language='markdown' if file_ext == '.md' else 'text')
                 except Exception as txt_err: st.sidebar.warning(f"Text preview failed: {txt_err}"); preview_failed = True

            # Actions
            action_cols = st.sidebar.columns(3)
            with action_cols[0]:
                 checkbox_key = f"cb_{item_key_base}"
                 st.session_state.setdefault('asset_checkboxes', {})
                 st.session_state['asset_checkboxes'][file] = st.checkbox("Select", value=st.session_state['asset_checkboxes'].get(file, False), key=checkbox_key)
            with action_cols[1]:
                 mime_map = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.pdf': 'application/pdf', '.txt': 'text/plain', '.md': 'text/markdown'}
                 mime_type = mime_map.get(file_ext, "application/octet-stream"); dl_key = f"dl_{item_key_base}"
                 try:
                     with open(file, "rb") as fp: st.download_button(label="📥", data=fp, file_name=basename, mime=mime_type, key=dl_key, help="Download")
                 except Exception as dl_e: st.error(f"DL Err: {dl_e}")
            with action_cols[2]:
                delete_key = f"del_{item_key_base}"
                if st.button("🗑️", key=delete_key, help=f"Delete {basename}"):
                    delete_success = False
                    try:
                        os.remove(file)
                        st.session_state['asset_checkboxes'].pop(file, None)
                        if file in st.session_state.get('layout_snapshots', []): st.session_state['layout_snapshots'].remove(file) # Remove if also in old list
                        logger.info(f"Deleted asset: {file}")
                        st.toast(f"Deleted {basename}!", icon="✅")
                        delete_success = True
                    except OSError as e: logger.error(f"Error deleting file {file}: {e}"); st.error(f"Could not delete {basename}: {e}")
                    except Exception as e: logger.error(f"Unexpected error deleting file {file}: {e}"); st.error(f"Could not delete {basename}: {e}")
                    # Rerun to refresh the gallery list after attempting delete
                    st.rerun()

        except FileNotFoundError: st.sidebar.error(f"File vanished: {basename}"); st.session_state['asset_checkboxes'].pop(file, None)
        except Exception as e: st.sidebar.error(f"Display Error: {basename}"); logger.error(f"Error displaying asset {file}: {e}")
        st.sidebar.markdown("---")

# --- UI Elements -----------------------------------------
# Sidebar Structure
st.sidebar.subheader("🤖 Hugging Face Settings")
# ... (HF API, Local Model, Params Expanders - code unchanged) ...
with st.sidebar.expander("API Inference Settings", expanded=False):
    st.session_state.hf_custom_key = st.text_input("Custom HF Token (BYOK)", value=st.session_state.get('hf_custom_key', ""), type="password", key="hf_custom_key_input", help="Enter your Hugging Face API token. Overrides HF_TOKEN env var.")
    token_status = "Custom Key Set" if st.session_state.hf_custom_key else ("Default HF_TOKEN Set" if HF_TOKEN else "No Token Set"); st.caption(f"Token Status: {token_status}")
    providers_list = ["hf-inference", "cerebras", "together", "sambanova", "novita", "cohere", "fireworks-ai", "hyperbolic", "nebius"]
    st.session_state.hf_provider = st.selectbox("Inference Provider", options=providers_list, index=providers_list.index(st.session_state.get('hf_provider', DEFAULT_PROVIDER)), key="hf_provider_select", help="Select the backend provider. Some require specific API keys.")
    if not st.session_state.hf_custom_key and not HF_TOKEN and st.session_state.hf_provider != "hf-inference": st.warning(f"Provider '{st.session_state.hf_provider}' may require a token.")
    st.session_state.hf_custom_api_model = st.text_input("Custom API Model ID", value=st.session_state.get('hf_custom_api_model', ""), key="hf_custom_model_input", placeholder="e.g., google/gemma-2-9b-it", help="Overrides the featured model selection below if provided.")
    effective_api_model = st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model
    st.session_state.hf_selected_api_model = st.selectbox("Featured API Model", options=FEATURED_MODELS_LIST, index=FEATURED_MODELS_LIST.index(st.session_state.get('hf_selected_api_model', FEATURED_MODELS_LIST[0])), key="hf_featured_model_select", help="Select a common model. Ignored if Custom API Model ID is set.")
    st.caption(f"Effective API Model: {effective_api_model}")
with st.sidebar.expander("Local Model Selection", expanded=True):
    if not _transformers_available: st.warning("Transformers library not found.")
    else:
        local_model_options = ["None"] + list(st.session_state.get('local_models', {}).keys())
        current_selection = st.session_state.get('selected_local_model_path'); current_selection = current_selection if current_selection in local_model_options else "None"
        selected_path = st.selectbox("Active Local Model", options=local_model_options, index=local_model_options.index(current_selection), format_func=lambda x: os.path.basename(x) if x != "None" else "None", key="local_model_selector", help="Select a loaded local model.")
        st.session_state.selected_local_model_path = selected_path if selected_path != "None" else None
        if st.session_state.selected_local_model_path:
             model_info = st.session_state.local_models[st.session_state.selected_local_model_path]
             st.caption(f"Type: {model_info.get('type', '?')} | Device: {model_info.get('model').device if model_info.get('model') else 'N/A'}")
        else: st.caption("No local model selected.")
with st.sidebar.expander("Generation Parameters", expanded=False):
    st.session_state.gen_max_tokens = st.slider("Max New Tokens", 1, 4096, st.session_state.get('gen_max_tokens', 512), step=1, key="param_max_tokens")
    st.session_state.gen_temperature = st.slider("Temperature", 0.01, 2.0, st.session_state.get('gen_temperature', 0.7), step=0.01, key="param_temp")
    st.session_state.gen_top_p = st.slider("Top-P", 0.01, 1.0, st.session_state.get('gen_top_p', 0.95), step=0.01, key="param_top_p")
    st.session_state.gen_frequency_penalty = st.slider("Repetition Penalty", 1.0, 2.0, st.session_state.get('gen_frequency_penalty', 0.0)+1.0, step=0.05, key="param_repetition", help="1.0 means no penalty.")
    st.session_state.gen_seed = st.slider("Seed", -1, 65535, st.session_state.get('gen_seed', -1), step=1, key="param_seed", help="-1 for random.")

st.sidebar.markdown("---")
# Gallery is rendered later by calling update_gallery()

# --- App Title & Main Area ---
st.title("Vision & Layout Titans (HF) 🚀🖼️📄")
st.markdown("Combined App: PDF Layout Generator + Hugging Face Powered AI Tools")

# Warning for missing libraries in main area if sidebar not ready
if not _transformers_available:
     st.warning("AI/ML libraries (torch, transformers) not found. Local model features disabled.")
elif not _diffusers_available:
     st.warning("Diffusers library not found. Diffusion model features disabled.")


# --- Main Application Tabs ---
tabs_to_create = [
    "Combined PDF Generator 📄", # Renamed Tab 0
    "Camera Snap 📷",
    "Download PDFs 📥",
    "Build Titan (Local Models) 🌱",
    "PDF Page Process (HF) 📄", # Clarified name
    "Image Process (HF) 🖼️",
    "Text Process (HF) 📝",
    "Test OCR (HF) 🔍",
    "Test Image Gen (Diffusers) 🎨",
    "Character Editor 🧑‍🎨",
    "Character Gallery 🖼️",
]
tabs = st.tabs(tabs_to_create)

# --- Tab Implementations ---

# --- Tab 1: Combined PDF Generator (OVERHAULED) ---
with tabs[0]:
    st.header("Combined PDF Generator 📄➕🖼️➕...")
    st.markdown("Select assets (Images, PDFs, Text/MD) from the sidebar gallery, reorder them, and generate a combined PDF.")

    # --- Get Selected Files ---
    selected_files_paths = [
        f for f, selected in st.session_state.get('asset_checkboxes', {}).items()
        if selected and os.path.exists(f) # Ensure file still exists
    ]

    if not selected_files_paths:
        st.info("👈 Select one or more assets from the sidebar gallery using the checkboxes.")
    else:
        st.info(f"{len(selected_files_paths)} assets selected from gallery.")

        # --- Populate DataFrame for Reordering ---
        combined_records = []
        for idx, filepath in enumerate(selected_files_paths):
            filename = os.path.basename(filepath)
            ext = os.path.splitext(filename)[1].lower()
            file_type = "Unknown"
            if ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff']: file_type = "Image"
            elif ext == '.pdf': file_type = "PDF"
            elif ext in ['.md', '.txt']: file_type = "Text"

            combined_records.append({
                "filename": filename,
                "filepath": filepath, # Keep the path
                "type": file_type,
                "order": idx, # Initial order based on selection
            })

        combined_df_initial = pd.DataFrame(combined_records)

        st.markdown("#### Reorder Selected Assets for PDF")
        st.caption("Edit the 'Order' column or drag rows to set the sequence for the combined PDF.")

        edited_combined_df = st.data_editor(
             combined_df_initial,
             column_config={
                  "filename": st.column_config.TextColumn("Filename", disabled=True),
                  "filepath": None, # Hide filepath column
                  "type": st.column_config.TextColumn("Type", disabled=True),
                  "order": st.column_config.NumberColumn(
                       "Order",
                       min_value=0,
                       # max_value=len(combined_df_initial)-1, # Max can cause issues if rows added/removed by user selection change
                       step=1,
                       required=True,
                  ),
             },
             hide_index=True,
             use_container_width=True,
             num_rows="dynamic", # Allow drag-and-drop reordering
             key="combined_pdf_editor"
        )

        # Sort by the edited 'order' column
        ordered_combined_df = edited_combined_df.sort_values('order').reset_index(drop=True)

        # Prepare list of dicts for the PDF generation function
        ordered_sources_info_for_pdf = ordered_combined_df[['filepath', 'type', 'filename']].to_dict('records')

        # --- Generate & Download ---
        st.subheader("Generate Combined PDF")
        if st.button("🖋️ Generate Combined PDF", key="generate_combined_pdf_btn"):
            if not ordered_sources_info_for_pdf:
                st.warning("No items available after reordering.")
            else:
                with st.spinner("Generating combined PDF... This might take a while."):
                    combined_pdf_bytes = make_combined_pdf(ordered_sources_info_for_pdf)

                if combined_pdf_bytes:
                    # Create filename
                    now = datetime.now(pytz.timezone("US/Central"))
                    prefix = now.strftime("%Y%m%d-%H%M%p")
                    first_item_name = clean_stem(ordered_sources_info_for_pdf[0].get('filename','combined'))
                    combined_pdf_fname = f"{prefix}_Combined_{first_item_name}.pdf"
                    combined_pdf_fname = re.sub(r'[^\w\-\.\_]', '_', combined_pdf_fname) # Sanitize

                    st.success(f"✅ Combined PDF ready: **{combined_pdf_fname}**")
                    st.download_button(
                        "⬇️ Download Combined PDF",
                        data=combined_pdf_bytes,
                        file_name=combined_pdf_fname,
                        mime="application/pdf",
                        key="download_combined_pdf_btn"
                    )
                    # Add preview (optional, might be slow for large combined PDFs)
                    # ... (preview logic similar to other tabs if desired) ...
                else:
                    st.error("Combined PDF generation failed. Check logs or input files.")


# --- Tab 2: Camera Snap ---
with tabs[1]:
    st.header("Camera Snap 📷")
    st.subheader("Single Capture (Adds to General Gallery)")
    cols = st.columns(2)
    with cols[0]:
        cam0_img = st.camera_input("Take a picture - Cam 0", key="main_cam0")
        if cam0_img:
            filename = generate_filename("cam0_snap");
            if st.session_state.get('cam0_file') and os.path.exists(st.session_state['cam0_file']): 
                try: 
                    os.remove(st.session_state['cam0_file']) 
                except OSError: 
                    pass
            try:
                with open(filename, "wb") as f: f.write(cam0_img.getvalue())
                st.session_state['cam0_file'] = filename; st.session_state['history'].append(f"Snapshot from Cam 0: {filename}"); st.image(Image.open(filename), caption="Camera 0 Snap", use_container_width=True); logger.info(f"Saved snapshot from Camera 0: {filename}"); st.success(f"Saved {filename}")
                update_gallery(); # Refresh sidebar without rerun
            except Exception as e: st.error(f"Failed to save Cam 0 snap: {e}"); logger.error(f"Failed to save Cam 0 snap {filename}: {e}")
    with cols[1]:
        cam1_img = st.camera_input("Take a picture - Cam 1", key="main_cam1")
        if cam1_img:
            filename = generate_filename("cam1_snap")
            if st.session_state.get('cam1_file') and os.path.exists(st.session_state['cam1_file']): 
                try: 
                    os.remove(st.session_state['cam1_file']) 
                except OSError: 
                    pass
            try:
                with open(filename, "wb") as f: f.write(cam1_img.getvalue())
                st.session_state['cam1_file'] = filename; st.session_state['history'].append(f"Snapshot from Cam 1: {filename}"); st.image(Image.open(filename), caption="Camera 1 Snap", use_container_width=True); logger.info(f"Saved snapshot from Camera 1: {filename}"); st.success(f"Saved {filename}")
                update_gallery(); # Refresh sidebar without rerun
            except Exception as e: st.error(f"Failed to save Cam 1 snap: {e}"); logger.error(f"Failed to save Cam 1 snap {filename}: {e}")


# --- Tab 3: Download PDFs ---
with tabs[2]:
    st.header("Download PDFs 📥")
    st.markdown("Download PDFs from URLs and optionally create image snapshots.")
    if st.button("Load Example arXiv URLs 📚", key="load_examples"):
        example_urls = ["https://arxiv.org/pdf/2308.03892", "https://arxiv.org/pdf/1706.03762", "https://arxiv.org/pdf/2402.17764", "https://www.clickdimensions.com/links/ACCERL/"]
        st.session_state['pdf_urls_input'] = "\n".join(example_urls)
    url_input = st.text_area("Enter PDF URLs (one per line)", value=st.session_state.get('pdf_urls_input', ""), height=150, key="pdf_urls_textarea")
    if st.button("Robo-Download PDFs 🤖", key="download_pdfs_button"):
        urls = [url.strip() for url in url_input.strip().split("\n") if url.strip()]
        if not urls: st.warning("Please enter at least one URL.")
        else:
            progress_bar = st.progress(0); status_text = st.empty(); total_urls = len(urls); download_count = 0; existing_pdfs = get_pdf_files()
            for idx, url in enumerate(urls):
                output_path = pdf_url_to_filename(url); status_text.text(f"Processing {idx + 1}/{total_urls}: {os.path.basename(output_path)}..."); progress_bar.progress((idx + 1) / total_urls)
                if os.path.exists(output_path): # Check existence properly
                    st.info(f"Already exists: {os.path.basename(output_path)}")
                    st.session_state['downloaded_pdfs'][url] = output_path
                    # Ensure checkbox state is preserved or reset if needed
                    st.session_state['asset_checkboxes'][output_path] = st.session_state['asset_checkboxes'].get(output_path, False)
                else:
                    if download_pdf(url, output_path):
                         st.session_state['downloaded_pdfs'][url] = output_path; logger.info(f"Downloaded PDF from {url} to {output_path}"); st.session_state['history'].append(f"Downloaded PDF: {output_path}"); st.session_state['asset_checkboxes'][output_path] = False; download_count += 1; existing_pdfs.append(output_path)
                    else: st.error(f"Failed to download: {url}")
            status_text.success(f"Download process complete! Successfully downloaded {download_count} new PDFs.")
            if download_count > 0: update_gallery(); # Refresh sidebar without rerun

    st.subheader("Create Snapshots from Gallery PDFs")
    snapshot_mode = st.selectbox("Snapshot Mode", ["First Page (High-Res)", "First Two Pages (High-Res)", "All Pages (High-Res)", "First Page (Low-Res Preview)"], key="pdf_snapshot_mode")
    resolution_map = {"First Page (High-Res)": 2.0, "First Two Pages (High-Res)": 2.0, "All Pages (High-Res)": 2.0, "First Page (Low-Res Preview)": 1.0}
    mode_key_map = {"First Page (High-Res)": "single", "First Two Pages (High-Res)": "twopage", "All Pages (High-Res)": "allpages", "First Page (Low-Res Preview)": "single"}
    resolution = resolution_map[snapshot_mode]; mode_key = mode_key_map[snapshot_mode]
    if st.button("Snapshot Selected PDFs 📸", key="snapshot_selected_pdfs"):
        selected_pdfs = [path for path in get_gallery_files(['pdf']) if st.session_state['asset_checkboxes'].get(path, False)]
        if not selected_pdfs: st.warning("No PDFs selected in the sidebar gallery!")
        else:
            st.info(f"Starting snapshot process for {len(selected_pdfs)} selected PDF(s)..."); snapshot_count = 0; total_snapshots_generated = 0
            for pdf_path in selected_pdfs:
                if not os.path.exists(pdf_path): st.warning(f"File not found: {pdf_path}. Skipping."); continue
                new_snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key, resolution))
                if new_snapshots:
                    snapshot_count += 1; total_snapshots_generated += len(new_snapshots)
                    st.write(f"Snapshots for {os.path.basename(pdf_path)}:"); cols = st.columns(3)
                    for i, snap_path in enumerate(new_snapshots):
                         with cols[i % 3]:
                              try: st.image(Image.open(snap_path), caption=os.path.basename(snap_path), use_container_width=True)
                              except Exception as snap_img_err: st.warning(f"Cannot display snap {os.path.basename(snap_path)}: {snap_img_err}")
                         st.session_state['asset_checkboxes'][snap_path] = False # Add to gallery
            if total_snapshots_generated > 0: st.success(f"Generated {total_snapshots_generated} snapshots from {snapshot_count} PDFs."); update_gallery(); # Refresh sidebar without rerun
            else: st.warning("No snapshots were generated. Check logs or PDF files.")


# --- Tab 4: Build Titan (Local Models) ---
with tabs[3]:
    st.header("Build Titan (Local Models) 🌱")
    st.markdown("Download and save models from Hugging Face Hub for local use.")
    if not _transformers_available:
        st.error("Transformers library not available. Cannot download or load local models.")
    else:
        build_model_type = st.selectbox("Select Model Type", ["Causal LM", "Vision/Multimodal", "OCR", "Diffusion"], key="build_type_local")
        st.subheader(f"Download {build_model_type} Model")
        hf_model_id = st.text_input("Hugging Face Model ID", placeholder=f"e.g., {'google/gemma-2-9b-it' if build_model_type == 'Causal LM' else 'llava-hf/llava-1.5-7b-hf' if build_model_type == 'Vision/Multimodal' else 'microsoft/trocr-base-handwritten' if build_model_type == 'OCR' else 'stabilityai/stable-diffusion-xl-base-1.0'}", key="build_hf_model_id")
        local_model_name = st.text_input("Local Name for this Model", value=f"{build_model_type.split('/')[0].lower()}_{os.path.basename(hf_model_id).replace('.','') if hf_model_id else 'model'}", key="build_local_name", help="A unique name to identify this model locally.")
        st.info("Private or gated models require a valid Hugging Face token (set via HF_TOKEN env var or the Custom Key in sidebar API settings).")

        if st.button(f"Download & Save '{hf_model_id}' Locally", key="build_download_button", disabled=not hf_model_id or not local_model_name):
            local_name_check = re.sub(r'[^\w\-]+', '_', local_model_name) # Sanitize proposed name for path check
            potential_path_base = os.path.join(f"{build_model_type.split('/')[0].lower()}_models", local_name_check)

            if any(os.path.basename(p) == local_name_check for p in get_local_model_paths(build_model_type.split('/')[0].lower())):
                 st.error(f"A local model folder named '{local_name_check}' already exists. Choose a different local name.")
            else:
                model_type_map = {"Causal LM": "causal", "Vision/Multimodal": "vision", "OCR": "ocr", "Diffusion": "diffusion"}
                model_type_short = model_type_map.get(build_model_type, "unknown")
                config = LocalModelConfig(name=local_model_name, hf_id=hf_model_id, model_type=model_type_short)
                save_path = config.get_full_path()
                os.makedirs(os.path.dirname(save_path), exist_ok=True)
                st.info(f"Attempting to download '{hf_model_id}' to '{save_path}'..."); progress_bar_build = st.progress(0); status_text_build = st.empty()
                token_build = st.session_state.hf_custom_key or HF_TOKEN or None
                try:
                    if build_model_type == "Diffusion":
                         if not _diffusers_available: raise ImportError("Diffusers library required.")
                         status_text_build.text("Downloading diffusion pipeline..."); pipeline_obj = StableDiffusionPipeline.from_pretrained(hf_model_id, token=token_build); status_text_build.text("Saving diffusion model pipeline..."); pipeline_obj.save_pretrained(save_path)
                         st.session_state.local_models[save_path] = {'type': 'diffusion', 'hf_id': hf_model_id, 'model':None, 'processor':None} # Mark as downloaded
                         st.success(f"Diffusion model '{hf_model_id}' downloaded and saved to {save_path}")
                         del pipeline_obj # Free memory
                    else:
                         status_text_build.text("Downloading model components...")
                         if model_type_short == 'causal': model_class, proc_tok_class = AutoModelForCausalLM, AutoTokenizer; proc_name="tokenizer"
                         elif model_type_short == 'vision': model_class, proc_tok_class = AutoModelForVision2Seq, AutoProcessor; proc_name="processor"
                         elif model_type_short == 'ocr': model_class, proc_tok_class = AutoModelForVision2Seq, AutoProcessor; proc_name="processor"
                         else: raise ValueError(f"Unknown model type: {model_type_short}")

                         model_obj = model_class.from_pretrained(hf_model_id, token=token_build); model_obj.save_pretrained(save_path)
                         status_text_build.text(f"Model saved. Downloading {proc_name}..."); proc_tok_obj = proc_tok_class.from_pretrained(hf_model_id, token=token_build); proc_tok_obj.save_pretrained(save_path)
                         status_text_build.text(f"Components saved. Loading '{local_model_name}' into memory...")
                         device = "cuda" if torch.cuda.is_available() else "cpu"
                         # Use trust_remote_code cautiously if needed for specific models
                         reloaded_model = model_class.from_pretrained(save_path).to(device)
                         reloaded_proc_tok = proc_tok_class.from_pretrained(save_path)
                         st.session_state.local_models[save_path] = {'type': model_type_short, 'hf_id': hf_model_id, 'model': reloaded_model, proc_name: reloaded_proc_tok}
                         # Add tokenizer specifically if it's nested in processor
                         if proc_name == "processor" and hasattr(reloaded_proc_tok, 'tokenizer'):
                             st.session_state.local_models[save_path]['tokenizer'] = reloaded_proc_tok.tokenizer
                         st.success(f"{build_model_type} model '{hf_model_id}' downloaded to {save_path} and loaded ({device})."); st.session_state.selected_local_model_path = save_path
                         del model_obj, proc_tok_obj # Free memory from download cache if possible
                except (RepositoryNotFoundError, GatedRepoError) as e: st.error(f"Download failed: Repo not found or requires access/token. Error: {e}"); logger.error(f"Download failed for {hf_model_id}: {e}"); #if os.path.exists(save_path): shutil.rmtree(save_path)
                except ImportError as e: st.error(f"Download failed: Library missing. {e}"); logger.error(f"ImportError for {hf_model_id}: {e}")
                except Exception as e: st.error(f"Download error: {e}"); logger.error(f"Download failed for {hf_model_id}: {e}", exc_info=True); #if os.path.exists(save_path): shutil.rmtree(save_path)
                finally: progress_bar_build.progress(1.0); status_text_build.empty(); #st.rerun() # Rerun removed

        st.subheader("Manage Local Models")
        # Refresh list for display
        loaded_model_paths = list(st.session_state.get('local_models', {}).keys())
        if not loaded_model_paths: st.info("No models downloaded yet.")
        else:
            models_df_data = []
            for path in loaded_model_paths:
                 data = st.session_state.local_models.get(path, {}) # Safely get data
                 models_df_data.append({
                      "Local Name": os.path.basename(path), "Type": data.get('type', '?'),
                      "HF ID": data.get('hf_id', '?'), "Loaded": "Yes" if data.get('model') else "No", "Path": path })
            models_df = pd.DataFrame(models_df_data); st.dataframe(models_df, use_container_width=True, hide_index=True, column_order=["Local Name", "Type", "HF ID", "Loaded"])
            model_to_delete = st.selectbox("Select model to delete", [""] + [os.path.basename(p) for p in loaded_model_paths], key="delete_model_select")
            if model_to_delete and st.button(f"Delete Local Model '{model_to_delete}'", type="primary"):
                path_to_delete = next((p for p in loaded_model_paths if os.path.basename(p) == model_to_delete), None)
                if path_to_delete:
                    try:
                        # Explicitly delete model objects from memory first if they exist
                        if path_to_delete in st.session_state.local_models:
                             model_data_to_del = st.session_state.local_models[path_to_delete]
                             if model_data_to_del.get('model'): del model_data_to_del['model']
                             if model_data_to_del.get('tokenizer'): del model_data_to_del['tokenizer']
                             if model_data_to_del.get('processor'): del model_data_to_del['processor']
                             if _transformers_available and torch.cuda.is_available(): torch.cuda.empty_cache() # Try to clear VRAM

                        # Remove from session state
                        st.session_state.local_models.pop(path_to_delete, None)
                        if st.session_state.selected_local_model_path == path_to_delete: st.session_state.selected_local_model_path = None
                        # Delete from disk
                        if os.path.exists(path_to_delete): shutil.rmtree(path_to_delete)
                        st.success(f"Deleted model '{model_to_delete}'."); logger.info(f"Deleted local model: {path_to_delete}"); st.rerun()
                    except Exception as e: st.error(f"Failed to delete model '{model_to_delete}': {e}"); logger.error(f"Failed to delete model {path_to_delete}: {e}")


# --- Tab 5: PDF Process (HF) ---
with tabs[4]:
    st.header("PDF Page Process with HF Models 📄")
    st.markdown("Upload PDFs, view pages, and extract text/info using selected HF models (API or Local Vision/OCR).")
    pdf_use_api = st.radio("Choose Processing Method", ["Hugging Face API", "Loaded Local Model"], key="pdf_process_source", horizontal=True, help="API uses settings from sidebar. Local uses the selected local model (if suitable for vision/OCR).")
    if pdf_use_api == "Hugging Face API": st.info(f"Using API Model: {st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model} (likely image-to-text)")
    else:
        if st.session_state.selected_local_model_path: st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}")
        else: st.warning("No local model selected.")

    uploaded_pdfs_process_hf = st.file_uploader("Upload PDF files to process", type=["pdf"], accept_multiple_files=True, key="pdf_process_uploader_hf")
    if uploaded_pdfs_process_hf:
        process_all_pages_pdf = st.checkbox("Process All Pages (can be slow/expensive)", value=False, key="pdf_process_all_hf")
        pdf_prompt = st.text_area("Prompt for PDF Page Processing", "Extract the text content from this page.", key="pdf_process_prompt_hf")
        if st.button("Process Uploaded PDFs with HF", key="process_uploaded_pdfs_hf"):
            if pdf_use_api == "Loaded Local Model" and not st.session_state.selected_local_model_path: st.error("Cannot process locally: No local model selected.")
            else:
                combined_text_output_hf = f"# HF PDF Processing Results ({'API' if pdf_use_api else 'Local'})\n\n"; total_pages_processed_hf = 0; output_placeholder_hf = st.container()
                for pdf_file in uploaded_pdfs_process_hf:
                    output_placeholder_hf.markdown(f"--- \n### Processing: {pdf_file.name}")
                    try:
                        pdf_bytes = pdf_file.read(); doc = fitz.open("pdf", pdf_bytes); num_pages = len(doc)
                        pages_to_process = range(num_pages) if process_all_pages_pdf else range(min(1, num_pages))
                        output_placeholder_hf.info(f"Processing {len(pages_to_process)} of {num_pages} pages..."); doc_text = f"## File: {pdf_file.name}\n\n"
                        for i in pages_to_process:
                            page_placeholder = output_placeholder_hf.empty(); page_placeholder.info(f"Processing Page {i + 1}/{num_pages}...")
                            page = doc.load_page(i); pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                            cols_pdf = output_placeholder_hf.columns(2); cols_pdf[0].image(img, caption=f"Page {i+1}", use_container_width=True)
                            with cols_pdf[1], st.spinner("Processing page with HF model..."): hf_text = process_image_hf(img, pdf_prompt, use_api=pdf_use_api)
                            st.text_area(f"Result (Page {i+1})", hf_text, height=250, key=f"pdf_hf_out_{pdf_file.name}_{i}")
                            doc_text += f"### Page {i + 1}\n\n{hf_text}\n\n---\n\n"; total_pages_processed_hf += 1; page_placeholder.empty()
                        combined_text_output_hf += doc_text; doc.close()
                    except Exception as e: output_placeholder_hf.error(f"Error processing {pdf_file.name}: {str(e)}")
                if total_pages_processed_hf > 0:
                    st.markdown("--- \n### Combined Processing Results"); st.text_area("Full Output", combined_text_output_hf, height=400, key="combined_pdf_hf_output")
                    output_filename_pdf_hf = generate_filename("hf_processed_pdfs", "md")
                    try:
                        with open(output_filename_pdf_hf, "w", encoding="utf-8") as f: f.write(combined_text_output_hf)
                        st.success(f"Combined output saved to {output_filename_pdf_hf}")
                        st.markdown(get_download_link(output_filename_pdf_hf, "text/markdown", "Download Combined MD"), unsafe_allow_html=True)
                        st.session_state['asset_checkboxes'][output_filename_pdf_hf] = False; update_gallery() # Refresh sidebar
                    except IOError as e: st.error(f"Failed to save combined output file: {e}")

# --- Tab 6: Image Process (HF) ---
with tabs[5]:
    st.header("Image Process with HF Models 🖼️")
    st.markdown("Upload images and process them using selected HF models (API or Local).")
    img_use_api = st.radio("Choose Processing Method", ["Hugging Face API", "Loaded Local Model"], key="img_process_source_hf", horizontal=True)
    if img_use_api == "Hugging Face API": st.info(f"Using API Model: {st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model} (likely image-to-text)")
    else:
        if st.session_state.selected_local_model_path: st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}")
        else: st.warning("No local model selected.")
    img_prompt_hf = st.text_area("Prompt for Image Processing", "Describe this image in detail.", key="img_process_prompt_hf")
    uploaded_images_process_hf = st.file_uploader("Upload image files", type=["png", "jpg", "jpeg"], accept_multiple_files=True, key="image_process_uploader_hf")
    if uploaded_images_process_hf:
        if st.button("Process Uploaded Images with HF", key="process_images_hf"):
             if img_use_api == "Loaded Local Model" and not st.session_state.selected_local_model_path: st.error("Cannot process locally: No local model selected.")
             else:
                combined_img_text_hf = f"# HF Image Processing Results ({'API' if img_use_api else 'Local'})\n\n**Prompt:** {img_prompt_hf}\n\n---\n\n"; images_processed_hf = 0; output_img_placeholder_hf = st.container()
                for img_file in uploaded_images_process_hf:
                    output_img_placeholder_hf.markdown(f"### Processing: {img_file.name}")
                    try:
                        img = Image.open(img_file); cols_img_hf = output_img_placeholder_hf.columns(2); cols_img_hf[0].image(img, caption=f"Input: {img_file.name}", use_container_width=True)
                        with cols_img_hf[1], st.spinner("Processing image with HF model..."): hf_img_text = process_image_hf(img, img_prompt_hf, use_api=img_use_api)
                        st.text_area(f"Result", hf_img_text, height=300, key=f"img_hf_out_{img_file.name}")
                        combined_img_text_hf += f"## Image: {img_file.name}\n\n{hf_img_text}\n\n---\n\n"; images_processed_hf += 1
                    except UnidentifiedImageError: output_img_placeholder_hf.error(f"Invalid Image: {img_file.name}. Skipping.")
                    except Exception as e: output_img_placeholder_hf.error(f"Error processing {img_file.name}: {str(e)}")
                if images_processed_hf > 0:
                    st.markdown("--- \n### Combined Processing Results"); st.text_area("Full Output", combined_img_text_hf, height=400, key="combined_img_hf_output")
                    output_filename_img_hf = generate_filename("hf_processed_images", "md")
                    try:
                        with open(output_filename_img_hf, "w", encoding="utf-8") as f: f.write(combined_img_text_hf)
                        st.success(f"Combined output saved to {output_filename_img_hf}"); st.markdown(get_download_link(output_filename_img_hf, "text/markdown", "Download Combined MD"), unsafe_allow_html=True)
                        st.session_state['asset_checkboxes'][output_filename_img_hf] = False; update_gallery() # Refresh sidebar
                    except IOError as e: st.error(f"Failed to save combined output file: {e}")

# --- Tab 7: Text Process (HF) ---
with tabs[6]:
    st.header("Text Process with HF Models 📝")
    st.markdown("Process Markdown (.md) or Text (.txt) files using selected HF models (API or Local).")
    text_use_api = st.radio("Choose Processing Method", ["Hugging Face API", "Loaded Local Model"], key="text_process_source_hf", horizontal=True)
    if text_use_api == "Hugging Face API": st.info(f"Using API Model: {st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model}")
    else:
        if st.session_state.selected_local_model_path: st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}")
        else: st.warning("No local model selected.")
    text_files_hf = get_gallery_files(['md', 'txt'])
    if not text_files_hf: st.warning("No .md or .txt files in gallery to process.")
    else:
        selected_text_file_hf = st.selectbox("Select Text/MD File to Process", options=[""] + text_files_hf, format_func=lambda x: os.path.basename(x) if x else "Select a file...", key="text_process_select_hf")
        if selected_text_file_hf:
             st.write(f"Selected: {os.path.basename(selected_text_file_hf)}")
             try:
                 with open(selected_text_file_hf, "r", encoding="utf-8", errors='ignore') as f: content_text_hf = f.read()
                 st.text_area("File Content Preview", content_text_hf[:1000] + ("..." if len(content_text_hf) > 1000 else ""), height=200, key="text_content_preview_hf")
                 prompt_text_hf = st.text_area("Enter Prompt for this File", "Summarize the key points of this text.", key="text_individual_prompt_hf")
                 if st.button(f"Process '{os.path.basename(selected_text_file_hf)}' with HF", key=f"process_text_hf_btn"):
                     if text_use_api == "Loaded Local Model" and not st.session_state.selected_local_model_path: st.error("Cannot process locally: No local model selected.")
                     else:
                        with st.spinner("Processing text with HF model..."): result_text_processed = process_text_hf(content_text_hf, prompt_text_hf, use_api=text_use_api)
                        st.markdown("### Processing Result"); st.markdown(result_text_processed)
                        output_filename_text_hf = generate_filename(f"hf_processed_{os.path.splitext(os.path.basename(selected_text_file_hf))[0]}", "md")
                        try:
                            with open(output_filename_text_hf, "w", encoding="utf-8") as f: f.write(result_text_processed)
                            st.success(f"Result saved to {output_filename_text_hf}"); st.markdown(get_download_link(output_filename_text_hf, "text/markdown", "Download Result MD"), unsafe_allow_html=True)
                            st.session_state['asset_checkboxes'][output_filename_text_hf] = False; update_gallery() # Refresh sidebar
                        except IOError as e: st.error(f"Failed to save result file: {e}")
             except FileNotFoundError: st.error("Selected file not found.")
             except Exception as e: st.error(f"Error reading file: {e}")

# --- Tab 8: Test OCR (HF) ---
with tabs[7]:
    st.header("Test OCR with HF Models 🔍")
    st.markdown("Select an image/PDF and run OCR using HF models (API or Local - requires suitable local model).")
    ocr_use_api = st.radio("Choose OCR Method", ["Hugging Face API (Basic Captioning/OCR)", "Loaded Local OCR Model"], key="ocr_source_hf", horizontal=True, help="API uses basic image-to-text. Local requires a dedicated OCR model (e.g., TrOCR) to be loaded.")
    if ocr_use_api == "Loaded Local OCR Model":
         if st.session_state.selected_local_model_path:
              model_info = st.session_state.local_models.get(st.session_state.selected_local_model_path,{})
              model_type = model_info.get('type'); model_name = os.path.basename(st.session_state.selected_local_model_path)
              if model_type != 'ocr': st.warning(f"Selected model ({model_name}) is type '{model_type}', not 'ocr'. Results may be poor.")
              else: st.info(f"Using Local OCR Model: {model_name}")
         else: st.warning("No local model selected.")

    gallery_files_ocr_hf = get_gallery_files(['png', 'jpg', 'jpeg', 'pdf'])
    if not gallery_files_ocr_hf: st.warning("No images or PDFs in gallery.")
    else:
        selected_file_ocr_hf = st.selectbox("Select Image or PDF from Gallery for OCR", options=[""] + gallery_files_ocr_hf, format_func=lambda x: os.path.basename(x) if x else "Select a file...", key="ocr_select_file_hf")
        if selected_file_ocr_hf:
            st.write(f"Selected: {os.path.basename(selected_file_ocr_hf)}"); file_ext_ocr_hf = os.path.splitext(selected_file_ocr_hf)[1].lower(); image_to_ocr_hf = None; page_info_hf = ""
            try:
                if file_ext_ocr_hf in ['.png', '.jpg', '.jpeg']: image_to_ocr_hf = Image.open(selected_file_ocr_hf)
                elif file_ext_ocr_hf == '.pdf':
                    doc = fitz.open(selected_file_ocr_hf)
                    if len(doc) > 0: pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); image_to_ocr_hf = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); page_info_hf = " (Page 1)"
                    else: st.warning("Selected PDF is empty.")
                    doc.close()
                if image_to_ocr_hf:
                    st.image(image_to_ocr_hf, caption=f"Image for OCR{page_info_hf}", use_container_width=True)
                    if st.button("Run HF OCR on this Image 🚀", key="ocr_run_button_hf"):
                        if ocr_use_api == "Loaded Local OCR Model" and not st.session_state.selected_local_model_path: st.error("Cannot run locally: No local model selected.")
                        else:
                            output_ocr_file_hf = generate_filename(f"hf_ocr_{os.path.splitext(os.path.basename(selected_file_ocr_hf))[0]}", "txt"); st.session_state['processing']['ocr'] = True
                            with st.spinner("Performing OCR with HF model..."): ocr_result_hf = asyncio.run(process_hf_ocr(image_to_ocr_hf, output_ocr_file_hf, use_api=ocr_use_api))
                            st.session_state['processing']['ocr'] = False; st.text_area("OCR Result", ocr_result_hf, height=300, key="ocr_result_display_hf")
                            if ocr_result_hf and not ocr_result_hf.startswith("Error") and not ocr_result_hf.startswith("["):
                                entry = f"HF OCR: {selected_file_ocr_hf}{page_info_hf} -> {output_ocr_file_hf}"
                                st.session_state['history'].append(entry)
                                if len(ocr_result_hf) > 5: st.success(f"OCR output saved to {output_ocr_file_hf}"); st.markdown(get_download_link(output_ocr_file_hf, "text/plain", "Download OCR Text"), unsafe_allow_html=True); st.session_state['asset_checkboxes'][output_ocr_file_hf] = False; update_gallery() # Refresh sidebar
                                else: st.warning("OCR output seems short/empty.")
                            else: st.error(f"OCR failed. {ocr_result_hf}")
            except Exception as e: st.error(f"Error loading file for OCR: {e}")

# --- Tab 9: Test Image Gen (Diffusers) ---
with tabs[8]:
    st.header("Test Image Generation (Diffusers) 🎨")
    st.markdown("Generate images using Stable Diffusion models loaded locally via the Diffusers library.")
    if not _diffusers_available: st.error("Diffusers library is required.")
    else:
        local_diffusion_paths = get_local_model_paths("diffusion") # Check diffusion_models folder
        if not local_diffusion_paths: st.warning("No local diffusion models found. Download one using the 'Build Titan' tab."); selected_diffusion_model_path = None
        else: selected_diffusion_model_path = st.selectbox("Select Local Diffusion Model", options=[""] + local_diffusion_paths, format_func=lambda x: os.path.basename(x) if x else "Select...", key="imggen_diffusion_model_select")
        prompt_imggen_diff = st.text_area("Image Generation Prompt", "A photorealistic cat wearing sunglasses, studio lighting", key="imggen_prompt_diff")
        neg_prompt_imggen_diff = st.text_area("Negative Prompt (Optional)", "ugly, deformed, blurry, low quality", key="imggen_neg_prompt_diff")
        steps_imggen_diff = st.slider("Inference Steps", 10, 100, 25, key="imggen_steps"); guidance_imggen_diff = st.slider("Guidance Scale", 1.0, 20.0, 7.5, step=0.5, key="imggen_guidance")
        if st.button("Generate Image 🚀", key="imggen_run_button_diff", disabled=not selected_diffusion_model_path):
            if not prompt_imggen_diff: st.warning("Please enter a prompt.")
            else:
                 status_imggen = st.empty()
                 try:
                     status_imggen.info(f"Loading diffusion pipeline: {os.path.basename(selected_diffusion_model_path)}..."); device = "cuda" if _transformers_available and torch.cuda.is_available() else "cpu"; dtype = torch.float16 if device == "cuda" else torch.float32
                     pipe = StableDiffusionPipeline.from_pretrained(selected_diffusion_model_path, torch_dtype=dtype).to(device); pipe.safety_checker = None # Optional
                     status_imggen.info(f"Generating image on {device} ({dtype})..."); start_gen_time = time.time()
                     gen_output = pipe(prompt=prompt_imggen_diff, negative_prompt=neg_prompt_imggen_diff or None, num_inference_steps=steps_imggen_diff, guidance_scale=guidance_imggen_diff)
                     gen_image = gen_output.images[0]; elapsed_gen = int(time.time() - start_gen_time); status_imggen.success(f"Image generated in {elapsed_gen}s!")
                     output_imggen_file_diff = generate_filename("diffusion_gen", "png"); gen_image.save(output_imggen_file_diff)
                     st.image(gen_image, caption=f"Generated: {output_imggen_file_diff}", use_container_width=True)
                     st.markdown(get_download_link(output_imggen_file_diff, "image/png", "Download Generated Image"), unsafe_allow_html=True)
                     st.session_state['asset_checkboxes'][output_imggen_file_diff] = False; update_gallery() # Refresh sidebar
                     st.session_state['history'].append(f"Diffusion Gen: '{prompt_imggen_diff[:30]}...' -> {output_imggen_file_diff}")
                 except ImportError: st.error("Diffusers or Torch library not found.")
                 except Exception as e: st.error(f"Image generation failed: {e}"); logger.error(f"Diffusion generation failed for {selected_diffusion_model_path}: {e}", exc_info=True)
                 finally: 
                     if 'pipe' in locals(): 
                         del pipe; torch.cuda.empty_cache() if device == "cuda" else None # Clear VRAM

# --- Tab 10: Character Editor ---
with tabs[9]:
    st.header("Character Editor 🧑‍🎨"); st.subheader("Create Your Character")
    load_characters(); existing_char_names = [c['name'] for c in st.session_state.get('characters', [])]
    form_key = f"character_form_{st.session_state.get('char_form_reset_key', 0)}"
    with st.form(key=form_key):
        st.markdown("**Create New Character**")
        if st.form_submit_button("Randomize Content 🎲"): st.session_state['char_form_reset_key'] += 1; st.rerun()
        rand_name, rand_gender, rand_intro, rand_greeting = randomize_character_content()
        name_char = st.text_input("Name (3-25 chars...)", value=rand_name, max_chars=25, key="char_name_input")
        gender_char = st.radio("Gender", ["Male", "Female"], index=["Male", "Female"].index(rand_gender), key="char_gender_radio")
        intro_char = st.text_area("Intro (Public description)", value=rand_intro, max_chars=300, height=100, key="char_intro_area")
        greeting_char = st.text_area("Greeting (First message)", value=rand_greeting, max_chars=300, height=100, key="char_greeting_area")
        tags_char = st.text_input("Tags (comma-separated)", "OC, friendly", key="char_tags_input")
        submitted = st.form_submit_button("Create Character ✨")
        if submitted:
            error = False; # Validation checks...
            if not (3 <= len(name_char) <= 25): st.error("Name must be 3-25 characters."); error = True
            if not re.match(r'^[a-zA-Z0-9 _-]+$', name_char): st.error("Name contains invalid characters."); error = True
            if name_char in existing_char_names: st.error(f"Name '{name_char}' already exists!"); error = True
            if not intro_char or not greeting_char: st.error("Intro/Greeting cannot be empty."); error = True
            if not error:
                tag_list = [tag.strip() for tag in tags_char.split(',') if tag.strip()]
                character_data = {"name": name_char, "gender": gender_char, "intro": intro_char, "greeting": greeting_char, "created_at": datetime.now(pytz.timezone("US/Central")).strftime('%Y-%m-%d %H:%M:%S %Z'), "tags": tag_list}
                if save_character(character_data): st.success(f"Character '{name_char}' created!"); st.session_state['char_form_reset_key'] += 1; st.rerun()

# --- Tab 11: Character Gallery ---
with tabs[10]:
    st.header("Character Gallery 🖼️"); load_characters(); characters_list = st.session_state.get('characters', [])
    if not characters_list: st.warning("No characters created yet.")
    else:
        st.subheader(f"Your Characters ({len(characters_list)})"); search_term = st.text_input("Search Characters by Name", key="char_gallery_search")
        if search_term: characters_list = [c for c in characters_list if search_term.lower() in c['name'].lower()]
        cols_char_gallery = st.columns(3); chars_to_delete = []
        for idx, char in enumerate(characters_list):
            with cols_char_gallery[idx % 3], st.container(border=True):
                st.markdown(f"**{char['name']}**"); st.caption(f"Gender: {char.get('gender', 'N/A')}")
                st.markdown("**Intro:**"); st.markdown(f"> {char.get('intro', '')}")
                st.markdown("**Greeting:**"); st.markdown(f"> {char.get('greeting', '')}")
                st.caption(f"Tags: {', '.join(char.get('tags', ['N/A']))}"); st.caption(f"Created: {char.get('created_at', 'N/A')}")
                delete_key_char = f"delete_char_{char['name']}_{idx}";
                if st.button(f"Delete", key=delete_key_char, type="primary", help=f"Delete {char['name']}"): chars_to_delete.append(char['name']) # Shorten button label
        if chars_to_delete:
             current_characters = st.session_state.get('characters', []); updated_characters = [c for c in current_characters if c['name'] not in chars_to_delete]
             st.session_state['characters'] = updated_characters
             try:
                 with open("characters.json", "w", encoding='utf-8') as f: json.dump(updated_characters, f, indent=2)
                 logger.info(f"Deleted characters: {', '.join(chars_to_delete)}"); st.success(f"Deleted: {', '.join(chars_to_delete)}"); st.rerun()
             except IOError as e: logger.error(f"Failed to save characters.json after deletion: {e}"); st.error("Failed to update character file.")

# --- Footer and Persistent Sidebar Elements ------------
st.sidebar.markdown("---")
# Update Sidebar Gallery (Call this at the end to reflect all changes)
update_gallery()

# Action Logs in Sidebar
st.sidebar.subheader("Action Logs 📜")
log_expander = st.sidebar.expander("View Logs", expanded=False)
with log_expander:
    # Display logs in reverse order (newest first)
    log_text = "\n".join([f"{record.levelname}: {record.message}" for record in reversed(log_records)])
    st.code(log_text, language='log')

# History in Sidebar
st.sidebar.subheader("Session History 📜")
history_expander = st.sidebar.expander("View History", expanded=False)
with history_expander:
     for entry in reversed(st.session_state.get("history", [])):
         if entry: history_expander.write(f"- {entry}")

st.sidebar.markdown("---")
st.sidebar.info("Using Hugging Face models for AI tasks.")
st.sidebar.caption("App Modified by AI Assistant")