Blame for main.py

import os
import sys
import uuid
import asyncio
import aiofiles
import concurrent.futures
import logging
import json
import re
import time
import random
import hmac
import shutil
from datetime import datetime
from typing import List, Tuple, Optional, Dict, Any
from contextlib import asynccontextmanager
from rapidfuzz import fuzz
from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks, Depends, HTTPException, Request
from fastapi.responses import JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.wsgi import WSGIMiddleware
from zipfile import ZipFile

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

from langchain_core.prompts import PromptTemplate
from langchain_core.documents import Document

# 1) Imports for OpenAI (OpenAIEmbeddings, ChatOpenAI)
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

# 1b) Imports for Gemini (Google Generative AI)
try:
    from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
except ImportError:
    ChatGoogleGenerativeAI = None
    GoogleGenerativeAIEmbeddings = None

import requests as _requests

try:
    from openai import OpenAI as OpenAIClient
except ImportError:
    OpenAIClient = None

from pdf_text_processing import (
    extract_text_from_pdf,
    extract_and_split,
    get_pdf_title,
    entity_match
)
from vectorstore import get_vectorstore, compute_file_hash, CHROMA_DB_BASE
from system_chart_extraction import system_chart_extraction
from system_chart_service import SystemChartService
from hypotheses_service import HypothesesService
from hypotheses_generation import PATTERN as CHART_PATTERN
from prompts import create_hypotheses_generation_prompt, create_design_goals_analysis_prompt
from conversation_storage import save_conversation, get_conversations, get_all_conversations, update_conversation
from plotextract.app import app as plotextract_app

# Metadata management functions
def load_metadata(database: str) -> Dict[str, Dict[str, Any]]:
    """
    Load metadata from the METADATA file in the database directory.
    Returns a dictionary with file_hash -> metadata mapping.
    """
    category_upload_dir = os.path.join(UPLOAD_DIR, database)
    metadata_file = os.path.join(category_upload_dir, "METADATA")
    
    if not os.path.exists(metadata_file):
        logger.debug(f"📋 No METADATA file found for {database}")
        return {}
    
    try:
        with open(metadata_file, 'r', encoding='utf-8') as f:
            metadata = json.load(f)
            logger.debug(f"📋 Loaded metadata for {database}: {len(metadata)} entries")
            logger.debug(f"📋 Hash keys: {[h[:16] + '...' for h in metadata.keys()]}")
            return metadata
    except Exception as e:
        logger.error(f"Error loading metadata for {database}: {str(e)}")
        return {}

async def embed_stream_into_vs(
    texts, metas, vectorstore, embedding_model,
    batch_size: int = 512,  
    max_concurrency: int = 12,
    progress_key: str = None
):
    """
    Embeds texts in parallel batches and writes each batch immediately to the vectorstore.
    Retries on transient errors with exponential backoff.
    Provides real-time progress tracking for large document uploads.
    """
    assert len(texts) == len(metas)
    n = len(texts)
    if n == 0:
        return
    
    batches = [(i, min(i + batch_size, n)) for i in range(0, n, batch_size)]
    sem = asyncio.Semaphore(max_concurrency)
    
    # Thread-safe progress tracking
    progress_lock = asyncio.Lock()
    processed_count = 0

    async def run_batch(batch_idx, bi, bj):
        nonlocal processed_count
        async with sem:
            delay = 1.0
            for attempt in range(6):
                try:
                    # Generate embeddings for this batch
                    embeds = await embedding_model.aembed_documents(texts[bi:bj])
                    
                    # Add to vectorstore
                    vectorstore.add_texts(
                        texts=texts[bi:bj],
                        metadatas=metas[bi:bj],
                        embeddings=embeds
                    )
                    
                    # Update progress atomically
                    async with progress_lock:
                        processed_count += (bj - bi)
                        
                        if progress_key and progress_key in EMBEDDING_PROGRESS:
                            # Calculate progress (10% to 90% for embedding phase)
                            progress_percent = min(90, int(10 + (processed_count / n * 80)))
                            
                            EMBEDDING_PROGRESS[progress_key].update({
                                "status": "embedding",
                                "processed_chunks": processed_count,
                                "progress": progress_percent,
                                "message": f"Embedded {processed_count}/{n} chunks ({len(batches)} batches)"
                            })
                            
                            logger.debug(f"Batch {batch_idx + 1}/{len(batches)} completed: {processed_count}/{n} chunks ({progress_percent}%)")
                    
                    return
                    
                except Exception as e:
                    # Handle retryable errors with exponential backoff
                    msg = str(e).lower()
                    if any(keyword in msg for keyword in ["rate", "429", "timeout", "temporarily", "overloaded", "quota"]):
                        if attempt < 5:  # Don't sleep on last attempt
                            jitter = random.uniform(0, 0.5)
                            sleep_time = delay + jitter
                            logger.warning(f"Retryable error in batch {batch_idx + 1}, attempt {attempt + 1}/6: {str(e)}. Sleeping {sleep_time:.2f}s")
                            await asyncio.sleep(sleep_time)
                            delay *= 1.5  # More conservative backoff
                            continue
                    
                    # Log the error with context
                    logger.error(f"Failed to embed batch {batch_idx + 1}/{len(batches)} (chunks {bi}-{bj}) after {attempt + 1} attempts: {str(e)}")
                    raise

    # Run all batches in parallel
    try:
        await asyncio.gather(*[
            run_batch(idx, bi, bj) 
            for idx, (bi, bj) in enumerate(batches)
        ])
        
        # Final progress update before persisting
        if progress_key and progress_key in EMBEDDING_PROGRESS:
            EMBEDDING_PROGRESS[progress_key].update({
                "status": "saving",
                "progress": 95,
                "message": f"Saving {n} embeddings to vector database..."
            })
        
        # Persist the vectorstore
        vectorstore.persist()
        logger.info(f"✅ Successfully embedded and saved {n} chunks in {len(batches)} parallel batches")
        
    except Exception as e:
        # Update progress on failure
        if progress_key and progress_key in EMBEDDING_PROGRESS:
            EMBEDDING_PROGRESS[progress_key].update({
                "status": "error",
                "progress": EMBEDDING_PROGRESS[progress_key].get("progress", 0),
                "message": f"Embedding failed: {str(e)}"
            })
        logger.error(f"❌ Embedding failed after processing {processed_count}/{n} chunks: {str(e)}")
        raise


def upsert_metadata(database: str, new_entries: Dict[str, Dict[str, Any]]) -> None:
    category_upload_dir = os.path.join(UPLOAD_DIR, database)
    os.makedirs(category_upload_dir, exist_ok=True)
    metadata_file = os.path.join(category_upload_dir, "METADATA")
    try:
        existing = {}
        if os.path.exists(metadata_file):
            with open(metadata_file, "r", encoding="utf-8") as f:
                try:
                    existing = json.load(f)
                except Exception:
                    logger.warning("⚠️ METADATA unreadable; rebuilding.")
                    existing = {}
        existing.update(new_entries)
        with open(metadata_file, "w", encoding="utf-8") as f:
            json.dump(existing, f, indent=2, ensure_ascii=False)
        logger.info(f"📋 METADATA updated; total entries: {len(existing)}")
    except Exception as e:
        logger.error(f"Error updating METADATA: {e}")


def save_metadata(database: str, metadata: Dict[str, Dict[str, Any]]):
    """
    Save metadata to the METADATA file in the database directory.
    PROTECTION: Will not overwrite existing METADATA files to prevent changes.
    """
    category_upload_dir = os.path.join(UPLOAD_DIR, database)
    os.makedirs(category_upload_dir, exist_ok=True)
    metadata_file = os.path.join(category_upload_dir, "METADATA")
    
    # PROTECTION: Check if METADATA file already exists
    if os.path.exists(metadata_file):
        logger.warning(f"🔒 METADATA file already exists for {database}. Skipping save to prevent modification.")
        return
    
    try:
        with open(metadata_file, 'w', encoding='utf-8') as f:
            json.dump(metadata, f, indent=2, ensure_ascii=False)
        logger.info(f"📋 Created new METADATA file for {database}")
    except Exception as e:
        logger.error(f"Error saving metadata for {database}: {str(e)}")

def extract_file_metadata(file_path: str, original_filename: str, file_type: str, file_hash: str) -> Dict[str, Any]:
    """
    Extract metadata from a file including title, size, and other relevant information.
    """
    metadata = {
        "original_filename": original_filename,
        "saved_filename": os.path.basename(file_path),
        "file_type": file_type,
        "file_hash": file_hash,
        "upload_timestamp": datetime.now().isoformat(),
        "file_size": os.path.getsize(file_path) if os.path.exists(file_path) else 0
    }
    
    # Extract title based on file type
    if file_type == "pdf":
        # For PDFs, extract title from the PDF content
        try:
            from pdf_text_processing import get_pdf_title
            title = get_pdf_title(file_path)
            metadata["title"] = title if title else original_filename
        except Exception as e:
            logger.warning(f"Could not extract title from PDF {original_filename}: {str(e)}")
            metadata["title"] = original_filename
    elif file_type == "text":
        # For text files, try to extract title from the file content
        try:
            from pdf_text_processing import get_file_title
            title = get_file_title(file_path)
            metadata["title"] = title if title else original_filename
        except Exception as e:
            logger.warning(f"Could not extract title from text file {original_filename}: {str(e)}")
            metadata["title"] = original_filename
    else:
        # For other files (images), use original filename as title
        metadata["title"] = original_filename
    
    return metadata

def get_file_metadata(database: str, file_hash: str) -> Dict[str, Any]:
    """
    Get metadata for a specific file by its hash.
    """
    metadata = load_metadata(database)
    file_metadata = metadata.get(file_hash, {})
    
    # Debug logging
    if file_metadata:
        logger.debug(f"✅ Found metadata for hash {file_hash[:16]}... in database {database}")
    else:
        logger.debug(f"❌ No metadata found for hash {file_hash[:16]}... in database {database}")
        logger.debug(f"📋 Available metadata hashes: {[h[:16] + '...' for h in metadata.keys()]}")
    
    return file_metadata

def get_file_title_from_metadata(database: str, file_hash: str) -> str:
    """
    Get the title of a file from its metadata.
    """
    metadata = get_file_metadata(database, file_hash)
    title = metadata.get("title", "Unknown")
    
    # Debug logging
    if title == "Unknown":
        logger.debug(f"🔍 No title found in metadata for hash {file_hash[:16]}... in database {database}")
    else:
        logger.debug(f"✅ Found title in metadata for hash {file_hash[:16]}...: {title}")
    
    return title

def resource_path(relative_path: str) -> str:
    """
    Get the absolute path to a resource, works for dev and for PyInstaller.
    """
    if hasattr(sys, '_MEIPASS'):
        return os.path.join(sys._MEIPASS, relative_path)
    return os.path.join(os.path.dirname(__file__), relative_path)

def get_base_directory() -> str:
    """
    HF Spaces-safe data root:
    - Prefer explicit override via APP_DATA_DIR or PERSISTENT_STORAGE_PATH
    - Else use /data/<app> only if it exists and is writable
    - Else use $HOME/.cache/<app>
    - Else use /tmp/<app>
    """
    override = (os.environ.get("APP_DATA_DIR") or os.environ.get("PERSISTENT_STORAGE_PATH") or "").strip()
    if override:
        os.makedirs(override, exist_ok=True)
        return override

    if os.path.isdir("/data") and os.access("/data", os.W_OK):
        base = "/data/pdf-rag"
        os.makedirs(base, exist_ok=True)
        return base

    home = os.environ.get("HOME")
    if home:
        base = os.path.join(home, ".cache", "pdf-rag")
        os.makedirs(base, exist_ok=True)
        return base

    base = "/tmp/pdf-rag"
    os.makedirs(base, exist_ok=True)
    return base

# Define paths relative to the base directory for persistent storage
APP_DATA_ROOT = get_base_directory()
UPLOAD_DIR = os.path.join(APP_DATA_ROOT, "uploaded_files")
os.makedirs(UPLOAD_DIR, exist_ok=True)
logger.info(f"UPLOAD_DIR set to: {os.path.abspath(UPLOAD_DIR)}")

@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    FastAPI lifespan:
    1) Load existing databases from previous sessions
    """
    load_existing_databases()
    yield

app = FastAPI(lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)

app.mount("/plotextract", WSGIMiddleware(plotextract_app))

# A registry of hash -> filename for each "database" category
FILE_REGISTRY = {}

# Progress tracking for embedding tasks
EMBEDDING_PROGRESS = {}

# Hypothesis generation context storage for continuous conversation
HYPOTHESIS_GENERATION_CONTEXT = {}

# Conversation storage for system chart extraction is now handled by conversation_storage module

def load_existing_databases():
    """
    Scan the structured uploaded files directory and existing databases to populate FILE_REGISTRY
    with information about previously uploaded files. Now uses category subdirectories.
    Note: Only loads text and image files since PDFs are now converted to text during upload.
    Also handles METADATA files for preserving file information.
    """
    global FILE_REGISTRY
    
    logger.info("🔍 Scanning for existing databases...")
    
    # Get all database directories from root databases
    database_dirs = []
    chroma_db_base = CHROMA_DB_BASE
    
    if os.path.exists(chroma_db_base):
        for item in os.listdir(chroma_db_base):
            db_path = os.path.join(chroma_db_base, item)
            if os.path.isdir(db_path) and not item.startswith('.'):
                database_dirs.append(item)
    
    logger.info(f"Found database directories: {database_dirs}")
    
    # For each database directory, check if there's a corresponding uploaded files subdirectory
    for db_name in database_dirs:
        if db_name not in FILE_REGISTRY:
            FILE_REGISTRY[db_name] = {}
        
        # Check for corresponding uploaded files subdirectory
        category_upload_dir = os.path.join(UPLOAD_DIR, db_name)
        
        if os.path.exists(category_upload_dir) and os.path.isdir(category_upload_dir):
            logger.info(f"📁 Found uploaded files directory for {db_name}")
            
            # Load existing metadata
            existing_metadata = load_metadata(db_name)
            
            # Scan all files in this category directory
            files_found = 0
            metadata_needs_migration = False
            migrated_metadata = {}
            
            for file in os.listdir(category_upload_dir):
                if not file.startswith('.') and file != "METADATA":
                    file_path = os.path.join(category_upload_dir, file)
                    if os.path.isfile(file_path):
                        try:
                            # Calculate file hash and add to registry
                            with open(file_path, 'rb') as f:
                                content = f.read()
                            file_hash = compute_file_hash(content)
                            FILE_REGISTRY[db_name][file_hash] = file
                            files_found += 1
                            
                            # Check if this file exists in metadata with the current hash
                            if file_hash in existing_metadata:
                                # Perfect match - use existing metadata
                                migrated_metadata[file_hash] = existing_metadata[file_hash]
                            else:
                                # Hash mismatch - try to find by filename pattern for .txt files (migrated from PDF)
                                found_metadata = None
                                if file.endswith('.txt'):
                                    # Try to find metadata entry that matches this file's original filename pattern
                                    for old_hash, meta in existing_metadata.items():
                                        if meta.get('saved_filename') == file:
                                            logger.info(f"🔄 Migrating metadata for {file}: {old_hash[:16]}... -> {file_hash[:16]}...")
                                            found_metadata = meta.copy()
                                            found_metadata['file_hash'] = file_hash  # Update to new hash
                                            migrated_metadata[file_hash] = found_metadata
                                            metadata_needs_migration = True
                                            break
                                    
                                    # If still not found, try matching by original filename
                                    if not found_metadata:
                                        for old_hash, meta in existing_metadata.items():
                                            # Extract the original filename without extension and compare
                                            original_name = os.path.splitext(meta.get('original_filename', ''))[0]
                                            file_name_part = file.split('_', 1)[1] if '_' in file else file  # Remove UUID prefix
                                            file_name_part = os.path.splitext(file_name_part)[0]  # Remove .txt extension
                                            
                                            if original_name and original_name == file_name_part:
                                                logger.info(f"🔄 Migrating metadata by filename match for {file}: {original_name}")
                                                found_metadata = meta.copy()
                                                found_metadata['file_hash'] = file_hash  # Update to new hash
                                                found_metadata['saved_filename'] = file  # Update saved filename
                                                migrated_metadata[file_hash] = found_metadata
                                                metadata_needs_migration = True
                                                break
                                
                                # If no metadata found, create new metadata
                                if not found_metadata:
                                    logger.info(f"📋 Creating metadata for existing file: {file}")
                                    # Determine file type
                                    file_ext = os.path.splitext(file.lower())[1]
                                    if file_ext in ['.txt', '.text']:
                                        file_type = "text"  # Could be from PDF or original text
                                    elif file_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']:
                                        file_type = "image"
                                    else:
                                        file_type = "unknown"
                                    
                                    # Create metadata for the file
                                    file_metadata = extract_file_metadata(file_path, file, file_type, file_hash)
                                    migrated_metadata[file_hash] = file_metadata
                                    metadata_needs_migration = True
                            
                        except Exception as e:
                            logger.warning(f"Error processing file {file}: {e}")
            
            # Update existing_metadata with migrated data
            existing_metadata = migrated_metadata
            
            # Log migration results
            if metadata_needs_migration:
                logger.info(f"🔄 Metadata migration completed for {db_name} database")
            
            # Save updated metadata (only if METADATA file doesn't exist yet)
            if existing_metadata:
                metadata_file = os.path.join(category_upload_dir, "METADATA")
                if not os.path.exists(metadata_file):
                    save_metadata(db_name, existing_metadata)
                else:
                    logger.info(f"🔒 METADATA file already exists for {db_name}. Skipping metadata updates.")
            
            logger.info(f"📁 Added {files_found} files to {db_name} database")
        else:
            logger.info(f"⚠️  No uploaded files directory found for {db_name} (expected: {category_upload_dir})")
    
    # Also check for any uploaded files subdirectories that don't have corresponding databases yet
    if os.path.exists(UPLOAD_DIR):
        for item in os.listdir(UPLOAD_DIR):
            item_path = os.path.join(UPLOAD_DIR, item)
            if os.path.isdir(item_path) and not item.startswith('.'):
                if item not in FILE_REGISTRY:
                    # Found a files directory without a corresponding database
                    logger.info(f"📂 Found uploaded files directory '{item}' without corresponding database")
                    FILE_REGISTRY[item] = {}
                    
                    # Load existing metadata
                    existing_metadata = load_metadata(item)
                    
                    # Add files from this directory
                    files_found = 0
                    metadata_needs_migration = False
                    migrated_metadata = {}
                    
                    for file in os.listdir(item_path):
                        if not file.startswith('.') and file != "METADATA":
                            file_path = os.path.join(item_path, file)
                            if os.path.isfile(file_path):
                                try:
                                    with open(file_path, 'rb') as f:
                                        content = f.read()
                                    file_hash = compute_file_hash(content)
                                    FILE_REGISTRY[item][file_hash] = file
                                    files_found += 1
                                    
                                    # Check if this file exists in metadata with the current hash
                                    if file_hash in existing_metadata:
                                        # Perfect match - use existing metadata
                                        migrated_metadata[file_hash] = existing_metadata[file_hash]
                                    else:
                                        # Hash mismatch - try to find by filename pattern for .txt files (migrated from PDF)
                                        found_metadata = None
                                        if file.endswith('.txt'):
                                            # Try to find metadata entry that matches this file's original filename pattern
                                            for old_hash, meta in existing_metadata.items():
                                                if meta.get('saved_filename') == file:
                                                    logger.info(f"🔄 Migrating metadata for {file}: {old_hash[:16]}... -> {file_hash[:16]}...")
                                                    found_metadata = meta.copy()
                                                    found_metadata['file_hash'] = file_hash  # Update to new hash
                                                    migrated_metadata[file_hash] = found_metadata
                                                    metadata_needs_migration = True
                                                    break
                                            
                                            # If still not found, try matching by original filename
                                            if not found_metadata:
                                                for old_hash, meta in existing_metadata.items():
                                                    # Extract the original filename without extension and compare
                                                    original_name = os.path.splitext(meta.get('original_filename', ''))[0]
                                                    file_name_part = file.split('_', 1)[1] if '_' in file else file  # Remove UUID prefix
                                                    file_name_part = os.path.splitext(file_name_part)[0]  # Remove .txt extension
                                                    
                                                    if original_name and original_name == file_name_part:
                                                        logger.info(f"🔄 Migrating metadata by filename match for {file}: {original_name}")
                                                        found_metadata = meta.copy()
                                                        found_metadata['file_hash'] = file_hash  # Update to new hash
                                                        found_metadata['saved_filename'] = file  # Update saved filename
                                                        migrated_metadata[file_hash] = found_metadata
                                                        metadata_needs_migration = True
                                                        break
                                        
                                        # If no metadata found, create new metadata
                                        if not found_metadata:
                                            logger.info(f"📋 Creating metadata for existing file: {file}")
                                            # Determine file type
                                            file_ext = os.path.splitext(file.lower())[1]
                                            if file_ext in ['.txt', '.text']:
                                                file_type = "text"
                                            elif file_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']:
                                                file_type = "image"
                                            else:
                                                file_type = "unknown"
                                            
                                            # Create metadata for the file
                                            file_metadata = extract_file_metadata(file_path, file, file_type, file_hash)
                                            migrated_metadata[file_hash] = file_metadata
                                            metadata_needs_migration = True
                                        
                                except Exception as e:
                                    logger.warning(f"Error processing file {file}: {e}")
                    
                    # Update existing_metadata with migrated data
                    existing_metadata = migrated_metadata
                    
                    # Log migration results
                    if metadata_needs_migration:
                        logger.info(f"🔄 Metadata migration completed for {item} database")
                    
                    # Save updated metadata (only if METADATA file doesn't exist yet)
                    if existing_metadata:
                        metadata_file = os.path.join(item_path, "METADATA")
                        if not os.path.exists(metadata_file):
                            save_metadata(item, existing_metadata)
                        else:
                            logger.info(f"🔒 METADATA file already exists for {item}. Skipping metadata updates.")
                    
                    logger.info(f"📁 Added {files_found} files to {item} (files-only database)")
    
    # Clean up empty databases
    empty_dbs = [db for db, files in FILE_REGISTRY.items() if not files]
    for db in empty_dbs:
        del FILE_REGISTRY[db]
        logger.info(f"🗑️  Removed empty database: {db}")
    
    logger.info(f"✅ Loaded {len(FILE_REGISTRY)} databases with files:")
    for db_name, files in FILE_REGISTRY.items():
        logger.info(f"  📚 {db_name}: {len(files)} files")

# Utility functions for file type detection
def is_image_file(filename: str) -> bool:
    """Check if a file is an image based on its extension."""
    image_extensions = ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']
    file_ext = os.path.splitext(filename.lower())[1]
    return file_ext in image_extensions

def is_text_or_pdf_file(filename: str) -> bool:
    """Check if a file is a PDF or text file based on its extension."""
    text_pdf_extensions = ['.pdf', '.txt', '.text']
    file_ext = os.path.splitext(filename.lower())[1]
    return file_ext in text_pdf_extensions

def get_file_type(filename: str) -> str:
    """Return the file type category: 'image', 'text_pdf', or 'unknown'."""
    if is_image_file(filename):
        return 'image'
    elif is_text_or_pdf_file(filename):
        return 'text_pdf'
    else:
        return 'unknown'

# Conversation management functions are now handled by conversation_storage module

# --------------------  Per-browser workspaces  -------------------- #
# Each browser gets a random workspace id in a cookie. Every server-side key
# (database, progress/extraction id, conversation, hypothesis run) is stored as
# "<workspace>-<name>", so users only see their own data; the client only ever
# sees the unprefixed names. API keys and models are held per workspace too.
WORKSPACE_COOKIE = "madtool_workspace"
WORKSPACE_MODELS: Dict[str, Dict[str, Any]] = {}  # workspace -> {"embedding", "qa", "chart"}


def is_workspace_id(value: Optional[str]) -> bool:
    return bool(value) and re.fullmatch(r"[0-9a-f]{32}", value) is not None


def workspace(request: Request) -> str:
    ws = request.cookies.get(WORKSPACE_COOKIE)
    if not is_workspace_id(ws):
        raise HTTPException(status_code=401, detail="Your session has no workspace. Please reload the page.")
    return ws


def scoped(ws: str, name: str) -> str:
    return f"{ws}-{name}"


def unscoped(ws: str, key: str) -> Optional[str]:
    """The client-facing name of a key in workspace `ws`, or None if the key isn't theirs."""
    prefix = f"{ws}-"
    return key[len(prefix):] if key.startswith(prefix) else None


@app.get("/")
def get_index(request: Request):
    response = FileResponse(resource_path("index.html"))
    if not is_workspace_id(request.cookies.get(WORKSPACE_COOKIE)):
        # Deployed, the site is always HTTPS (HF's proxy terminates TLS and doesn't say so),
        # and SameSite=None lets the cookie work inside the huggingface.co Space iframe.
        # Local http dev keeps a plain Lax cookie.
        secure = request.url.hostname not in ("localhost", "127.0.0.1")
        response.set_cookie(
            WORKSPACE_COOKIE, uuid.uuid4().hex, max_age=365 * 24 * 3600, httponly=True,
            secure=secure, samesite="none" if secure else "lax",
        )
    return response


# --------------------  Model / API related info  -------------------- #
OPENAI_MODELS = [
    "gpt-5-nano",
    "gpt-5-mini",
    "gpt-5.2"
]

GEMINI_MODELS = [
    "gemini-3-flash-preview",
    "gemini-3-pro-preview",
    "gemini-2.5-flash-lite"
]

AVAILABLE_MODELS = OPENAI_MODELS + GEMINI_MODELS

DEFAULT_MODEL = "gpt-5.2"
# Shared output-length cap for QA, system chart extraction, and hypotheses.
# LangChain's max_tokens maps to provider-specific names:
# OpenAI max_output_tokens / max_completion_tokens, Gemini maxOutputTokens.
DEFAULT_MAX_OUTPUT_TOKENS = 8192

ADMIN_DELETE_CONFIRMATION_TEXT = "DELETE APP DATA"

def get_model_provider(model: str) -> Optional[str]:
    if model in OPENAI_MODELS:
        return "openai"
    if model in GEMINI_MODELS:
        return "gemini"
    return None


def detect_provider_from_key(api_key: str) -> Optional[str]:
    """Infer provider from an API key prefix.

    AIza... → Gemini, sk-... → OpenAI.
    Anthropic (sk-ant-...) keys are detected but unsupported — they must be
    checked *before* the generic sk- prefix so they are not misread as OpenAI;
    callers reject the "anthropic" value with a clear message.
    """
    key = (api_key or "").strip()
    if not key:
        return None
    if key.startswith("sk-ant-"):
        return "anthropic"
    if key.startswith("AIza"):
        return "gemini"
    if key.startswith("sk-"):
        return "openai"
    return None


UNSUPPORTED_ANTHROPIC_KEY_MESSAGE = (
    "Anthropic (sk-ant-...) keys are not supported: Anthropic does not provide an "
    "embedding model, which this tool requires. Please use an OpenAI (sk-...) or "
    "Gemini (AIza...) key."
)


def provider_display_name(provider: str) -> str:
    return {"openai": "OpenAI", "gemini": "Gemini"}.get(provider, provider or "Unknown")


# Substrings of OpenAI model ids that are not text-chat models (same list
# PlotExtract uses, plus completions-only "instruct" models). "search" also
# covers the tool-only deep-research models.
OPENAI_NON_CHAT_MARKERS = (
    "audio",
    "dall-e",
    "embedding",
    "gpt-image",
    "instruct",
    "moderation",
    "realtime",
    "search",
    "transcribe",
    "tts",
    "whisper",
)


def is_openai_chat_model(model_id: str) -> bool:
    mid = model_id.lower()
    is_family = "gpt" in mid or re.match(r"^o\d+(?:-|$)", mid)  # GPT and o-series reasoning models
    return bool(is_family) and not any(marker in mid for marker in OPENAI_NON_CHAT_MARKERS)


def list_openai_chat_models(api_key: str) -> List[Dict[str, str]]:
    if OpenAIClient is None:
        raise RuntimeError("openai SDK is not installed.")
    client = OpenAIClient(api_key=api_key)
    page = client.models.list()
    out: List[Dict[str, Any]] = []
    seen = set()
    for item in getattr(page, "data", []) or []:
        model_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None)
        if not model_id or model_id in seen:
            continue
        if not is_openai_chat_model(model_id):
            continue
        seen.add(model_id)
        created = getattr(item, "created", None) or (item.get("created") if isinstance(item, dict) else None)
        out.append({"id": model_id, "display_name": model_id, "created": created or 0})
    # Newest first by release timestamp
    out.sort(key=lambda m: (m["created"], m["id"]), reverse=True)
    for m in out:
        m.pop("created")
    return out


def _fetch_gemini_model_list(api_key: str) -> List[Dict[str, Any]]:
    url = "https://generativelanguage.googleapis.com/v1beta/models"
    response = _requests.get(url, params={"key": api_key}, timeout=15)
    if response.status_code != 200:
        try:
            err = response.json().get("error", {}).get("message") or response.text
        except Exception:
            err = response.text
        raise RuntimeError(f"Gemini API returned {response.status_code}: {err}")
    payload = response.json() or {}
    return payload.get("models", []) or []


# Substrings of Gemini model ids that support generateContent but are not
# text-chat models (speech, image generation, live audio, tool-only agents).
GEMINI_NON_CHAT_MARKERS = (
    "aqa",
    "audio",
    "computer-use",
    "embedding",
    "image",
    "live",
    "robotics",
    "tts",
)


def gemini_sort_key(model_id: str) -> Tuple:
    """Newest-first ordering for Gemini ids, which carry no release date: the
    version number tracks release order, then pro > flash > flash-lite, then
    stable before preview/experimental. "-latest" aliases float to the top.
    """
    mid = model_id.lower()
    version = re.search(r"-(\d+(?:\.\d+)?)(?:-|$)", mid)
    if "lite" in mid:
        tier = 1
    elif "flash" in mid:
        tier = 2
    elif "pro" in mid:
        tier = 3
    else:
        tier = 0
    return (
        "latest" in mid,
        mid.startswith("gemini"),
        float(version.group(1)) if version else 0.0,
        tier,
        not any(tag in mid for tag in ("preview", "exp")),
        mid,
    )


def list_gemini_models(api_key: str) -> List[Dict[str, str]]:
    """List Gemini chat models via the public REST API to avoid pinning conflicts
    between google-generativeai and langchain-google-genai.
    """
    out: List[Dict[str, str]] = []
    seen = set()
    for item in _fetch_gemini_model_list(api_key):
        methods = item.get("supportedGenerationMethods") or []
        if "generateContent" not in methods:
            continue
        raw_id = item.get("name", "")
        model_id = raw_id.split("/", 1)[-1] if raw_id else ""
        if not model_id or model_id in seen:
            continue
        if any(marker in model_id.lower() for marker in GEMINI_NON_CHAT_MARKERS):
            continue
        seen.add(model_id)
        display = item.get("displayName") or model_id
        out.append({"id": model_id, "display_name": str(display)})
    out.sort(key=lambda m: gemini_sort_key(m["id"]), reverse=True)
    return out


def pick_gemini_embedding_model(api_key: str) -> str:
    """Find a Gemini embedding model this key actually has access to.

    The v1beta endpoint exposes different embedding models depending on the key's
    tier and rollout — `embedding-001` and `text-embedding-004` are sometimes 404.
    We list models, filter to those supporting `embedContent`, and pick the
    newest one we recognise.
    """
    try:
        embedding_models: List[str] = []
        for item in _fetch_gemini_model_list(api_key):
            methods = item.get("supportedGenerationMethods") or []
            if "embedContent" not in methods:
                continue
            name = item.get("name") or ""
            if name:
                embedding_models.append(name)
        for preferred_prefix in (
            "models/gemini-embedding",
            "models/text-embedding",
            "models/embedding",
        ):
            for name in embedding_models:
                if name.startswith(preferred_prefix):
                    return name
        if embedding_models:
            return embedding_models[0]
    except Exception as exc:
        logger.warning(f"Failed to pick Gemini embedding model dynamically, falling back: {exc}")
    return "models/gemini-embedding-001"


def fetch_models_for_provider(provider: str, api_key: str) -> List[Dict[str, str]]:
    if provider == "openai":
        return list_openai_chat_models(api_key)
    if provider == "gemini":
        return list_gemini_models(api_key)
    return []


def choose_default_model(models: List[Dict[str, str]], provider: str) -> str:
    # The provider lists are already sorted newest-first, so the default is
    # simply the top entry.
    return models[0]["id"] if models else ""

class AsyncEmbeddingAdapter:
    def __init__(self, embedder):
        self._embedder = embedder

    async def aembed_documents(self, texts):
        if hasattr(self._embedder, "aembed_documents"):
            return await self._embedder.aembed_documents(texts)
        return await asyncio.to_thread(self._embedder.embed_documents, texts)

    def __getattr__(self, name):
        return getattr(self._embedder, name)

def ensure_async_embeddings(embedder):
    if hasattr(embedder, "aembed_documents"):
        return embedder
    return AsyncEmbeddingAdapter(embedder)


def llm_content_to_text(content: Any) -> str:
    """
    Normalize LLM content payloads (string/list/dict) into plain text.
    Some providers (e.g., Responses API) may return structured blocks.
    """
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: List[str] = []
        for item in content:
            text = llm_content_to_text(item)
            if text:
                parts.append(text)
        return "\n".join(parts)
    if isinstance(content, dict):
        # Common structured message shapes from provider SDKs.
        for key in ("text", "content", "output_text"):
            if key in content:
                return llm_content_to_text(content.get(key))
        return json.dumps(content, ensure_ascii=False)
    return str(content)

@app.post("/set_api_key")
async def set_api_key(
    api_key_input: Optional[str] = Form(None),
    model: str = Form(DEFAULT_MODEL),
    ws: str = Depends(workspace),
):
    """
    Sets up this workspace's LLM and embeddings for proprietary models.
    """
    if not api_key_input:
        return JSONResponse(
            {"message": "Error: API key is required for proprietary models."},
            status_code=400
        )
    api_key = api_key_input.strip()
    model_provider = detect_provider_from_key(api_key)
    if model_provider == "anthropic":
        return JSONResponse(
            content={"message": UNSUPPORTED_ANTHROPIC_KEY_MESSAGE},
            status_code=400
        )
    if not model_provider:
        return JSONResponse(
            content={"message": "Unsupported API key format. Use an OpenAI (sk-...) or Gemini (AIza...) key."},
            status_code=400
        )
    if not model:
        return JSONResponse(
            content={"message": "Please select a model."},
            status_code=400
        )
    DEFAULT_MODEL_QA = model
    DEFAULT_MODEL_SYSTEM_CHART_EXTRACTION = model

    try:
        if not api_key:
            return JSONResponse(
                content={"message": "Error: API key is required for proprietary models."},
                status_code=400
            )

        if model_provider == "openai":
            embedding_model = OpenAIEmbeddings(openai_api_key=api_key)
            embedding_model = ensure_async_embeddings(embedding_model)
            openai_common_kwargs: Dict[str, Any] = {
                "openai_api_key": api_key,
                "model": DEFAULT_MODEL_QA,
                "use_responses_api": True,
                "max_tokens": DEFAULT_MAX_OUTPUT_TOKENS,
            }
            openai_chart_kwargs: Dict[str, Any] = {
                "openai_api_key": api_key,
                "model": DEFAULT_MODEL_SYSTEM_CHART_EXTRACTION,
                "use_responses_api": True,
                "max_tokens": DEFAULT_MAX_OUTPUT_TOKENS,
            }
            llm_QA = ChatOpenAI(
                **openai_common_kwargs
            )
            llm_system_chart_extraction = ChatOpenAI(
                **openai_chart_kwargs
            )
            provider_label = "OpenAI"
        elif model_provider == "gemini":
            if ChatGoogleGenerativeAI is None or GoogleGenerativeAIEmbeddings is None:
                return JSONResponse(
                    content={"message": "Gemini models require langchain-google-genai. Please install it to continue."},
                    status_code=500
                )
            gemini_embedding_model_name = pick_gemini_embedding_model(api_key)
            logger.info(f"Using Gemini embedding model: {gemini_embedding_model_name}")
            embedding_model = GoogleGenerativeAIEmbeddings(
                google_api_key=api_key,
                model=gemini_embedding_model_name
            )
            embedding_model = ensure_async_embeddings(embedding_model)
            gemini_common_kwargs: Dict[str, Any] = {
                "google_api_key": api_key,
                "model": DEFAULT_MODEL_QA,
                "max_tokens": DEFAULT_MAX_OUTPUT_TOKENS,
            }
            gemini_chart_kwargs: Dict[str, Any] = {
                "google_api_key": api_key,
                "model": DEFAULT_MODEL_SYSTEM_CHART_EXTRACTION,
                "max_tokens": DEFAULT_MAX_OUTPUT_TOKENS,
            }
            llm_QA = ChatGoogleGenerativeAI(
                **gemini_common_kwargs
            )
            llm_system_chart_extraction = ChatGoogleGenerativeAI(
                **gemini_chart_kwargs
            )
            provider_label = "Gemini"
        else:
            return JSONResponse(
                content={"message": f"Unsupported provider: {model_provider}."},
                status_code=400
            )

        WORKSPACE_MODELS[ws] = {"embedding": embedding_model, "qa": llm_QA, "chart": llm_system_chart_extraction}
        return JSONResponse(
            content={"message": f"{provider_label} model '{model}' set successfully!"},
            status_code=200
        )
    except Exception as e:
        return JSONResponse(
            content={"message": f"Error setting up model: {str(e)}"},
            status_code=500
        )

@app.get("/available_models")
async def get_available_models():
    return JSONResponse(content={
        "models": AVAILABLE_MODELS
    })


@app.post("/list_models")
async def list_models_for_key(api_key_input: Optional[str] = Form(None)):
    """Return the models available for the provider that the supplied API key belongs to.

    Provider is inferred from the key prefix; the dropdown only ever shows models
    matching the detected provider.
    """
    key = (api_key_input or "").strip()
    if not key:
        return JSONResponse({"message": "API key is required."}, status_code=400)
    provider = detect_provider_from_key(key)
    if provider == "anthropic":
        return JSONResponse(
            {"message": UNSUPPORTED_ANTHROPIC_KEY_MESSAGE},
            status_code=400,
        )
    if not provider:
        return JSONResponse(
            {"message": "Unsupported API key format. Use an OpenAI (sk-...) or Gemini (AIza...) key."},
            status_code=400,
        )
    try:
        models = fetch_models_for_provider(provider, key)
    except Exception as exc:
        return JSONResponse(
            {"message": f"Failed to fetch models from {provider_display_name(provider)}.", "details": str(exc)},
            status_code=400,
        )
    if not models:
        return JSONResponse(
            {"message": f"No models available for this {provider_display_name(provider)} key."},
            status_code=400,
        )
    return JSONResponse({
        "provider": provider,
        "provider_label": provider_display_name(provider),
        "models": models,
        "default_model": choose_default_model(models, provider),
    })

# --------------------  Upload  -------------------- #
async def save_file(file: UploadFile, database: str) -> Tuple[str, str]:
    """
    Saves text or image files to disk in a subdirectory for the database.
    For PDFs: extracts text and saves only the .txt file (not the original PDF).
    For text files: saves the text file as-is.
    For images: saves the image file as-is.
    Also extracts and saves metadata including paper titles to a METADATA file.
    Returns (file_path, duplicate_filename).
    If duplicate is detected, returns (None, original_filename).
    """
    # Check if file type is supported
    file_ext = os.path.splitext(file.filename.lower())[1]
    supported_extensions = ['.pdf', '.txt', '.text', '.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']
    if file_ext not in supported_extensions:
        logger.warning(f"⚠️ Unsupported file type: {file.filename}")
        return None, f"Unsupported file type: {file.filename}"
    
    content = await file.read()
    
    # For non-PDF files, check duplicates immediately
    # For PDFs, duplicate check happens after text extraction (due to hash consistency)
    if file_ext != '.pdf':
        file_hash = compute_file_hash(content)
        if file_hash in FILE_REGISTRY[database]:
            logger.warning(f"⚠️ Duplicate detected: {file.filename}")
            return None, file.filename

    # Create category-specific directory
    category_upload_dir = os.path.join(UPLOAD_DIR, database)
    os.makedirs(category_upload_dir, exist_ok=True)

    # PROTECTION: Check if METADATA file already exists
    metadata_file = os.path.join(category_upload_dir, "METADATA")
    if os.path.exists(metadata_file):
        error_msg = f"🔒 Database '{database}' is locked. Cannot add files to databases with existing METADATA files."
        logger.warning(error_msg)
        return None, error_msg

    file_id = str(uuid.uuid4())[:8]
    image_extensions = ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']
    
    if file_ext == '.pdf':
        # For PDFs: Extract text AND preserve original PDF file
        logger.info(f"Processing PDF: {file.filename} (will save both text and original PDF)")
        
        # Create temporary file to extract text and metadata from
        import tempfile
        with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as temp_pdf:
            temp_pdf.write(content)
            temp_pdf_path = temp_pdf.name
        
        try:
            # Extract text from the temporary PDF
            docs = extract_text_from_pdf(temp_pdf_path)
            if not docs:
                logger.warning(f"⚠️ No text could be extracted from PDF: {file.filename}")
                os.unlink(temp_pdf_path)  # Clean up temp file
                return None, f"No text could be extracted from: {file.filename}"
            
            # Save the extracted text file
            text_filename = f"{file_id}_{os.path.splitext(file.filename)[0]}.txt"
            text_path = os.path.join(category_upload_dir, text_filename)
            
            async with aiofiles.open(text_path, "w", encoding="utf-8") as tfile:
                await tfile.write(docs[0].page_content)
            
            # ALSO save the original PDF file
            pdf_filename = f"{file_id}_{file.filename}"
            pdf_path = os.path.join(category_upload_dir, pdf_filename)
            
            async with aiofiles.open(pdf_path, "wb") as pfile:
                await pfile.write(content)
            
            # CRITICAL FIX: Recompute hash from the saved .txt file content, not original PDF
            # This ensures hash consistency between upload and loading phases
            with open(text_path, "rb") as txt_file:
                txt_content = txt_file.read()
            file_hash = compute_file_hash(txt_content)  # Use TXT hash instead of PDF hash
            
            # Check for duplicates with the new hash
            if file_hash in FILE_REGISTRY[database]:
                logger.warning(f"⚠️ Duplicate detected after text extraction: {file.filename}")
                os.unlink(temp_pdf_path)  # Clean up temp file
                os.unlink(pdf_path)  # Clean up saved PDF
                return None, file.filename
            
            # Extract and save metadata (including title from original PDF and PDF path)
            file_metadata = extract_file_metadata(temp_pdf_path, file.filename, "pdf", file_hash)
            file_metadata["pdf_path"] = pdf_filename  # Store PDF filename
            upsert_metadata(database, {file_hash: file_metadata})
            
            # Register the text file in FILE_REGISTRY (for text-based operations)
            FILE_REGISTRY[database][file_hash] = text_filename
            logger.info(f"📄 Extracted and saved text: {text_filename} in {database}")
            logger.info(f"📄 Also saved original PDF: {pdf_filename}")
            logger.info(f"📋 Saved metadata with title: {file_metadata.get('title', 'Unknown')}")
            logger.info(f"🔑 Using TXT file hash: {file_hash[:16]}... for consistency")
            
            os.unlink(temp_pdf_path)  # Clean up temp file
            return text_path, None
            
        except Exception as e:
            logger.error(f"❌ Error processing PDF {file.filename}: {str(e)}")
            os.unlink(temp_pdf_path)  # Clean up temp file
            return None, f"Error processing PDF: {str(e)}"
    
    elif file_ext in image_extensions:
        # For images: Save the image file as-is
        filename = f"{file_id}_{file.filename}"
        file_path = os.path.join(category_upload_dir, filename)
        logger.info(f"Saving image to: {os.path.abspath(file_path)}")
        
        async with aiofiles.open(file_path, "wb") as f:
            await f.write(content)
        
        # Extract and save metadata for image (file_hash already computed above)
        file_metadata = extract_file_metadata(file_path, file.filename, "image", file_hash)
        upsert_metadata(database, {file_hash: file_metadata})
            
        FILE_REGISTRY[database][file_hash] = filename
        logger.info(f"📸 Saved image: {filename} in {database}")
        logger.info(f"📋 Saved metadata for image: {file.filename}")
        return file_path, None
    
    else:
        # For text files: Save as text with UTF-8 encoding
        filename = f"{file_id}_{file.filename}"
        file_path = os.path.join(category_upload_dir, filename)
        logger.info(f"Saving text file to: {os.path.abspath(file_path)}")
        
        async with aiofiles.open(file_path, "w", encoding="utf-8") as f:
            await f.write(content.decode('utf-8'))
        
        # Extract and save metadata for text file (file_hash already computed above)
        file_metadata = extract_file_metadata(file_path, file.filename, "text", file_hash)
        upsert_metadata(database, {file_hash: file_metadata})
            
        FILE_REGISTRY[database][file_hash] = filename
        logger.info(f"📝 Saved text file: {filename} in {database}")
        logger.info(f"📋 Saved metadata with title: {file_metadata.get('title', 'Unknown')}")
        return file_path, None

def save_file_from_bytes(content: bytes, original_filename: str, database: str) -> Tuple[str, str]:
    """
    Save a single file provided as raw bytes with duplicate checks and metadata handling.
    Returns (file_path, duplicate_or_error_message). If duplicate or error, file_path is None.
    Mirrors logic in save_file() but for in-memory content (e.g., ZIP entries).
    """
    # Check if file type is supported
    file_ext = os.path.splitext(original_filename.lower())[1]
    supported_extensions = ['.pdf', '.txt', '.text', '.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']
    if file_ext not in supported_extensions:
        logger.warning(f"⚠️ Unsupported file type in ZIP: {original_filename}")
        return None, f"Unsupported file type: {original_filename}"

    # For non-PDF files, check duplicates immediately
    if file_ext != '.pdf':
        file_hash = compute_file_hash(content)
        if file_hash in FILE_REGISTRY[database]:
            logger.warning(f"⚠️ Duplicate detected in ZIP: {original_filename}")
            return None, original_filename

    # Create category-specific directory
    category_upload_dir = os.path.join(UPLOAD_DIR, database)
    os.makedirs(category_upload_dir, exist_ok=True)

    file_id = str(uuid.uuid4())[:8]
    image_extensions = ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']

    if file_ext == '.pdf':
        # For PDFs: Extract text AND preserve original PDF file
        logger.info(f"Processing PDF from ZIP: {original_filename} (will save both text and original PDF)")

        import tempfile
        # Create temporary file to extract text and metadata from
        with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as temp_pdf:
            temp_pdf.write(content)
            temp_pdf_path = temp_pdf.name

        try:
            # Extract text from the temporary PDF
            docs = extract_text_from_pdf(temp_pdf_path)
            if not docs:
                logger.warning(f"⚠️ No text could be extracted from PDF (ZIP): {original_filename}")
                os.unlink(temp_pdf_path)
                return None, f"No text could be extracted from: {original_filename}"

            # Save the extracted text file
            text_filename = f"{file_id}_{os.path.splitext(original_filename)[0]}.txt"
            text_path = os.path.join(category_upload_dir, text_filename)
            with open(text_path, "w", encoding="utf-8") as tfile:
                tfile.write(docs[0].page_content)

            # ALSO save the original PDF file
            pdf_filename = f"{file_id}_{original_filename}"
            pdf_path = os.path.join(category_upload_dir, pdf_filename)
            with open(pdf_path, "wb") as pfile:
                pfile.write(content)

            # Recompute hash from the saved .txt file content
            with open(text_path, "rb") as txt_file:
                txt_content = txt_file.read()
            file_hash = compute_file_hash(txt_content)

            # Check for duplicates with the new hash
            if file_hash in FILE_REGISTRY[database]:
                logger.warning(f"⚠️ Duplicate detected after text extraction (ZIP): {original_filename}")
                os.unlink(temp_pdf_path)
                os.unlink(pdf_path)
                try:
                    os.unlink(text_path)
                except Exception:
                    pass
                return None, original_filename

            # Extract and save metadata (including title from original PDF and PDF path)
            file_metadata = extract_file_metadata(temp_pdf_path, original_filename, "pdf", file_hash)
            file_metadata["pdf_path"] = pdf_filename
            upsert_metadata(database, {file_hash: file_metadata})

            # Register the text file in FILE_REGISTRY
            FILE_REGISTRY[database][file_hash] = text_filename
            logger.info(f"📄 Extracted and saved text (ZIP): {text_filename} in {database}")
            logger.info(f"📄 Also saved original PDF (ZIP): {pdf_filename}")

            os.unlink(temp_pdf_path)
            return text_path, None

        except Exception as e:
            logger.error(f"❌ Error processing PDF from ZIP {original_filename}: {str(e)}")
            try:
                os.unlink(temp_pdf_path)
            except Exception:
                pass
            return None, f"Error processing PDF: {str(e)}"

    elif file_ext in image_extensions:
        # For images: Save the image file as-is
        filename = f"{file_id}_{original_filename}"
        file_path = os.path.join(category_upload_dir, filename)
        logger.info(f"Saving image (ZIP) to: {os.path.abspath(file_path)}")
        with open(file_path, "wb") as f:
            f.write(content)

        # Extract and save metadata for image
        file_hash = compute_file_hash(content)
        file_metadata = extract_file_metadata(file_path, original_filename, "image", file_hash)
        upsert_metadata(database, {file_hash: file_metadata})

        FILE_REGISTRY[database][file_hash] = filename
        logger.info(f"📸 Saved image (ZIP): {filename} in {database}")
        return file_path, None

    else:
        # For text files: Save as text with UTF-8 encoding
        filename = f"{file_id}_{original_filename}"
        file_path = os.path.join(category_upload_dir, filename)
        logger.info(f"Saving text file (ZIP) to: {os.path.abspath(file_path)}")
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(content.decode('utf-8'))

        # Extract and save metadata for text file
        file_hash = compute_file_hash(content)
        file_metadata = extract_file_metadata(file_path, original_filename, "text", file_hash)
        upsert_metadata(database, {file_hash: file_metadata})

        FILE_REGISTRY[database][file_hash] = filename
        logger.info(f"📝 Saved text file (ZIP): {filename} in {database}")
        return file_path, None

def obtain_text_from_files(database: str):
    """
    Returns a dict of title -> full file text for all text files in the given database.
    Now that PDFs are converted to .txt files during upload, we only need to read the text files.
    Uses metadata to get proper titles including original PDF titles.
    """
    file_dict = {}

    if database not in FILE_REGISTRY or not FILE_REGISTRY[database]:
        logger.warning(f"Database '{database}' not found or empty; returning no papers.")
        return {}
    
    # Create path to category-specific directory
    category_upload_dir = os.path.join(UPLOAD_DIR, database)
    
    for file_hash, filename in FILE_REGISTRY[database].items():
        file_path = os.path.join(category_upload_dir, filename)
        
        # Skip image files - we only want text content
        if is_image_file(filename):
            continue
            
        # Since PDFs are now converted to .txt files during upload,
        # we only need to read text files directly
        if filename.lower().endswith(('.txt', '.text')):
            try:
                # Get title from metadata (preserves original PDF titles)
                file_title = get_file_title_from_metadata(database, file_hash)
                if not file_title or file_title == "Unknown":
                    # Fallback to extracting title from file content
                    logger.info(f"🔍 Metadata lookup failed for {filename} (hash: {file_hash[:16]}...), using fallback title extraction")
                    from pdf_text_processing import get_file_title
                    temp_title = get_file_title(file_path)
                    file_title = temp_title if temp_title else filename
                    logger.info(f"📝 Fallback title for {filename}: {file_title}")
                else:
                    logger.info(f"✅ Using title from metadata for {filename}: {file_title}")

                # Read the text file content
                with open(file_path, "r", encoding="utf-8") as f:
                    file_text = f.read()

                file_dict[file_title] = file_text
            except Exception as e:
                logger.error(f"❌ Error reading file {filename}: {str(e)}")
                continue
    
    return file_dict



def obtain_text_from_pdf(database: str):
    """
    Legacy function name for backward compatibility.
    Now handles both PDF and text files.
    """
    return obtain_text_from_files(database)



def extract_all_files(file_paths: List[str]) -> List[Document]:
    """
    Runs text extraction + chunk-splitting on each file (text files only, since PDFs are now converted to .txt during upload).
    """
    all_docs = []
    with concurrent.futures.ThreadPoolExecutor(
        max_workers=min(32, os.cpu_count())
    ) as executor:
        futures = [executor.submit(extract_and_split, path) for path in file_paths]
        for future in concurrent.futures.as_completed(futures):
            result = future.result()
            if result:
                all_docs.extend(result)
    return all_docs

def save_to_vectorstore(embeddings, texts, metadatas, vectorstore):
    """
    Save a batch of (embeddings, texts, metadatas) to the vectorstore in chunks.
    """
    try:
        batch_size = 4000  # tune as needed (e.g., 500-2000)
        n = len(texts)
        for i in range(0, n, batch_size):
            j = min(i + batch_size, n)
            vectorstore.add_texts(
                texts=texts[i:j],
                metadatas=metadatas[i:j],
                embeddings=embeddings[i:j],
            )
        vectorstore.persist()
        logger.info(f"✅ Vectorstore saved: {n} chunks (in batches of {batch_size})")
    except Exception as e:
        logger.error(f"❌ Vectorstore save error: {e}")

def extract_texts_and_metas(docs: List[Document]):
    texts = [doc.page_content for doc in docs]
    metas = [doc.metadata for doc in docs]
    return texts, metas

# -------------------- Progress tracking endpoints -------------------- #
@app.get("/embedding_progress/{database}")
async def get_embedding_progress(database: str, ws: str = Depends(workspace)):
    """
    Returns the current progress of embedding tasks for a specific database.
    """
    key = scoped(ws, database)
    if key not in EMBEDDING_PROGRESS:
        print(f"❌ No task found for {database}")
        return JSONResponse(content={"status": "no_task", "progress": 0, "message": "No embedding task in progress"})
    
    progress_data = EMBEDDING_PROGRESS[key].copy()
    # Remove zip_data from response to avoid sending large binary data
    if "zip_data" in progress_data:
        del progress_data["zip_data"]
    return JSONResponse(content=progress_data)

@app.get("/download_extraction/{extraction_id}")
async def download_extraction(extraction_id: str, ws: str = Depends(workspace)):
    """
    Download the system chart extraction ZIP file directly from memory.
    """
    progress_data = EMBEDDING_PROGRESS.get(scoped(ws, extraction_id))
    if progress_data is None:
        return JSONResponse(content={"error": "Extraction not found"}, status_code=404)
    
    if progress_data.get("status") != "completed":
        return JSONResponse(content={"error": "Extraction not completed yet"}, status_code=400)
    
    zip_data = progress_data.get("zip_data")
    if not zip_data:
        return JSONResponse(content={"error": "Extraction data not available"}, status_code=404)
    
    filename = progress_data.get("filename", "system_chart_extraction.zip")
    
    # Return the ZIP file as a direct download
    from fastapi.responses import Response
    return Response(
        content=zip_data,
        media_type="application/zip",
        headers={
            "Content-Disposition": f"attachment; filename={filename}",
            "Content-Length": str(len(zip_data)),
            "Content-Type": "application/zip",
            "Cache-Control": "no-cache, no-store, must-revalidate",
            "Pragma": "no-cache",
            "Expires": "0"
        }
    )

async def process_files_async_all(file_paths: List[str], database: str, embedding_model):
    """
    Process all files with parallel embedding for improved performance.
    Uses the new embed_stream_into_vs function for concurrent batch processing.
    Note: PDFs are now converted to .txt files during upload, so we only process text files.
    """
    global EMBEDDING_PROGRESS
    if embedding_model is None:
        logger.error("❌ Embedding model is None. Did you forget /set_api_key?")
        return

    total_files = len(file_paths)
    logger.info(f"🚀 Starting parallel processing of {total_files} files for database '{database}'")
    
    # Initialize progress tracking
    EMBEDDING_PROGRESS[database] = {
        "status": "extracting",
        "progress": 0,
        "total_files": total_files,
        "processed_files": 0,
        "message": f"Extracting text from {total_files} files..."
    }

    try:
        # Extract text from all files
        vectorstore = get_vectorstore(database, embedding_model)
        all_docs = extract_all_files(file_paths)
        
        if not all_docs:
            EMBEDDING_PROGRESS[database] = {
                "status": "completed", 
                "progress": 100, 
                "message": "No documents found to embed"
            }
            logger.warning(f"⚠️ No documents extracted from {total_files} files")
            return

        # Extract texts and metadata
        texts, metas = extract_texts_and_metas(all_docs)
        total_chunks = len(texts)
        
        logger.info(f"📄 Extracted {total_chunks} chunks from {total_files} files")
        
        # Update progress after extraction
        EMBEDDING_PROGRESS[database] = {
            "status": "embedding",
            "progress": 10,
            "total_chunks": total_chunks,
            "processed_chunks": 0,
            "message": f"Starting parallel embedding of {total_chunks} chunks..."
        }

        # Use parallel embedding with progress tracking
        await embed_stream_into_vs(
            texts, metas, vectorstore, embedding_model,
            batch_size=512,
            max_concurrency=12,
            progress_key=database
        )

        # Final completion status
        EMBEDDING_PROGRESS[database] = {
            "status": "completed",
            "progress": 100,
            "message": f"Successfully embedded {total_chunks} chunks from {total_files} files using parallel processing"
        }
        
        logger.info(f"✅ Completed parallel processing: {total_chunks} chunks from {total_files} files")
        
    except asyncio.CancelledError:
        EMBEDDING_PROGRESS[database] = {
            "status": "error", 
            "progress": EMBEDDING_PROGRESS[database].get("progress", 0), 
            "message": "Embedding process was cancelled"
        }
        logger.warning(f"⚠️ Embedding cancelled for database '{database}'")
        raise
    except Exception as e:
        current_progress = EMBEDDING_PROGRESS[database].get("progress", 0)
        EMBEDDING_PROGRESS[database] = {
            "status": "error", 
            "progress": current_progress, 
            "message": f"Error during parallel embedding: {str(e)}"
        }
        logger.error(f"❌ Embedding failed for database '{database}': {str(e)}")
        raise

# Removed per-file embedding mode; only all-at-once embedding is supported.

async def process_system_chart_extraction_async(category: str, llm, extraction_id: str, enable_visualization: bool = False, run_system_chart_normalization: bool = False):
    """
    Background task for system chart extraction with progress tracking.
    Uses the SystemChartService for business logic while handling API-specific concerns.
    """
    global EMBEDDING_PROGRESS
    
    print(f"🚀 System Chart Background task started for extraction_id: {extraction_id}")
    print(f"📊 Visualization enabled: {enable_visualization}")
    
    # Initialize progress tracking
    EMBEDDING_PROGRESS[extraction_id] = {
        "status": "starting",
        "progress": 0,
        "message": "Initializing system chart extraction..."
    }
    
    try:
        # Create system chart service instance
        system_chart_service = SystemChartService()
        
        # Get PDF text dictionary
        pdf_dict = obtain_text_from_pdf(category)
        if not pdf_dict:
            EMBEDDING_PROGRESS[extraction_id] = {
                "status": "error",
                "progress": 0,
                "message": "No files found in the specified category."
            }
            return
        
        # Define progress callback for API-specific progress tracking
        def progress_callback(progress_data):
            print(f"📊 Progress update: {progress_data}")
            EMBEDDING_PROGRESS[extraction_id] = progress_data
            print(f"🔄 EMBEDDING_PROGRESS[{extraction_id}] updated to: {EMBEDDING_PROGRESS[extraction_id]}")
            print(f"📋 Current keys in EMBEDDING_PROGRESS: {list(EMBEDDING_PROGRESS.keys())}")
        
        # Perform system chart extraction using the service
        run_norm_effective = bool(enable_visualization or run_system_chart_normalization)
        result = await system_chart_service.extract_system_charts(
            category=category,
            llm=llm,
            pdf_dict=pdf_dict,
            enable_visualization=enable_visualization,
            run_system_chart_normalization=run_norm_effective,
            progress_callback=progress_callback,
            extraction_id=extraction_id  # Pass extraction_id to save conversations
        )
        
        # Handle the result
        if result["success"]:
            EMBEDDING_PROGRESS[extraction_id] = {
                "status": "completed",
                "progress": 100,
                "message": result["message"],
                "zip_data": result["zip_data"],
                "filename": result["filename"]
            }
        else:
            EMBEDDING_PROGRESS[extraction_id] = {
                "status": "error",
                "progress": 0,
                "message": result["message"]
            }
        
    except Exception as e:
        EMBEDDING_PROGRESS[extraction_id] = {
            "status": "error",
            "progress": 0,
            "message": f"Error during system chart extraction: {str(e)}"
        }
        logger.error(f"System chart extraction error: {str(e)}")

async def process_multi_database_system_chart_extraction_async(ws: str, categories: List[str], llm, extraction_id: str, enable_visualization: bool = False, run_system_chart_normalization: bool = False):
    """
    Background task for system chart extraction across multiple databases with progress tracking.
    Combines text from all databases and runs system chart extraction on the combined dataset.
    """
    global EMBEDDING_PROGRESS
    
    print(f"🚀 Multi-Database System Chart Background task started for extraction_id: {extraction_id}")
    print(f"📊 Processing {len(categories)} databases: {categories}")
    print(f"📊 Visualization enabled: {enable_visualization}")
    
    # Initialize progress tracking
    EMBEDDING_PROGRESS[extraction_id] = {
        "status": "gathering_data",
        "progress": 5,
        "message": f"Gathering text from {len(categories)} databases...",
        "databases": categories,
        "current_database": None,
        "databases_processed": 0,
        "total_databases": len(categories)
    }
    
    try:
        # Create system chart service instance
        system_chart_service = SystemChartService()
        
        # Gather text from all databases
        combined_pdf_dict = {}
        databases_processed = 0
        
        for i, category in enumerate(categories):
            EMBEDDING_PROGRESS[extraction_id].update({
                "progress": 5 + (i * 20 // len(categories)),  # 5-25% for data gathering
                "message": f"Gathering text from database: {category}",
                "current_database": category,
                "databases_processed": i
            })
            
            print(f"📁 Processing database {i+1}/{len(categories)}: {category}")
            
            # Get PDF text dictionary for this database
            pdf_dict = obtain_text_from_pdf(scoped(ws, category))
            if pdf_dict:
                # Add database prefix to avoid title conflicts
                for title, content in pdf_dict.items():
                    prefixed_title = f"[{category}] {title}"
                    combined_pdf_dict[prefixed_title] = content
                print(f"✅ Added {len(pdf_dict)} documents from {category}")
            else:
                print(f"⚠️  No files found in database: {category}")
            
            databases_processed += 1
        
        # Check if we have any documents to process
        if not combined_pdf_dict:
            EMBEDDING_PROGRESS[extraction_id] = {
                "status": "error",
                "progress": 0,
                "message": f"No files found in any of the selected databases: {', '.join(categories)}"
            }
            return
        
        print(f"📚 Combined {len(combined_pdf_dict)} documents from {len(categories)} databases")
        
        # Update progress for starting extraction
        EMBEDDING_PROGRESS[extraction_id].update({
            "status": "extracting",
            "progress": 30,
            "message": f"Starting system chart extraction on {len(combined_pdf_dict)} documents...",
            "current_database": "combined",
            "databases_processed": len(categories),
            "total_documents": len(combined_pdf_dict)
        })
        
        # Define progress callback for API-specific progress tracking
        def progress_callback(progress_data):
            # Adjust progress to account for the 30% already used for data gathering
            adjusted_progress = 30 + (progress_data.get("progress", 0) * 0.7)  # 30-100%
            
            # Merge with our multi-database tracking data
            updated_data = {
                "status": progress_data.get("status", "processing"),
                "progress": int(adjusted_progress),
                "message": progress_data.get("message", "Processing..."),
                "databases": categories,
                "databases_processed": len(categories),
                "total_databases": len(categories),
                "total_documents": len(combined_pdf_dict)
            }
            
            # Add specific progress details if available
            for key in ["current_paper", "total_papers", "current_question", "total_questions", 
                       "completed_questions", "current_paper_type", "questions_in_paper"]:
                if key in progress_data:
                    updated_data[key] = progress_data[key]
            
            print(f"📊 Multi-DB Progress update: {updated_data}")
            EMBEDDING_PROGRESS[extraction_id] = updated_data
        
        run_norm_effective = bool(enable_visualization or run_system_chart_normalization)
        # Perform system chart extraction using the service on combined data
        # Use a combined category name for the service
        combined_category_name = "_".join(categories)
        
        result = await system_chart_service.extract_system_charts(
            category=combined_category_name,
            llm=llm,
            pdf_dict=combined_pdf_dict,
            enable_visualization=enable_visualization,
            run_system_chart_normalization=run_norm_effective,
            progress_callback=progress_callback,
            extraction_id=extraction_id  # Pass extraction_id to save conversations
        )
        
        # Handle the result
        if result["success"]:
            EMBEDDING_PROGRESS[extraction_id] = {
                "status": "completed",
                "progress": 100,
                "message": f"System chart extraction completed for {len(categories)} databases ({len(combined_pdf_dict)} documents)",
                "zip_data": result["zip_data"],
                "filename": result["filename"],
                "databases": categories,
                "databases_processed": len(categories),
                "total_databases": len(categories),
                "total_documents": len(combined_pdf_dict)
            }
        else:
            EMBEDDING_PROGRESS[extraction_id] = {
                "status": "error",
                "progress": 0,
                "message": f"Error during extraction: {result['message']}",
                "databases": categories,
                "total_documents": len(combined_pdf_dict)
            }
        
    except Exception as e:
        EMBEDDING_PROGRESS[extraction_id] = {
            "status": "error",
            "progress": 0,
            "message": f"Error during multi-database system chart extraction: {str(e)}",
            "databases": categories
        }
        logger.error(f"Multi-database system chart extraction error: {str(e)}")

async def process_hypotheses_generation_async(system_charts_content_list: List[str], llm, extraction_id: str, material_goal_1: str, material_goal_2: str, iterations_count: int = 1):
    """
    Background task for synergistic hypotheses generation with progress tracking.
    Uses the HypothesesService for business logic while handling API-specific concerns.
    """
    global EMBEDDING_PROGRESS, HYPOTHESIS_GENERATION_CONTEXT
    
    print(f"🚀 Hypotheses Generation Background task started for extraction_id: {extraction_id}")
    
    # Store context for continuous conversation
    HYPOTHESIS_GENERATION_CONTEXT[extraction_id] = {
        'system_charts_content_list': system_charts_content_list,
        'material_goal_1': material_goal_1,
        'material_goal_2': material_goal_2,
        'iterations_count': iterations_count
    }
    
    # Initialize progress tracking
    EMBEDDING_PROGRESS[extraction_id] = {
        "status": "starting",
        "progress": 0,
        "message": "Initializing hypothesis generation..."
    }
    
    try:
        # Create hypotheses service instance
        hypotheses_service = HypothesesService()
        
        # Validate the system charts content list first
        validation = HypothesesService.validate_system_charts_content(system_charts_content_list)
        if not validation["valid"]:
            EMBEDDING_PROGRESS[extraction_id] = {
                "status": "error",
                "progress": 0,
                "message": validation["message"]
            }
            return
        
        print(f"📊 Validated {validation['valid_charts']} system charts from {len(system_charts_content_list)} files for hypothesis generation")
        
        # Define progress callback for API-specific progress tracking
        def progress_callback(progress_data):
            EMBEDDING_PROGRESS[extraction_id] = progress_data
        
        # Perform hypotheses generation using the service
        result = await hypotheses_service.generate_hypotheses(
            system_charts_content_list=system_charts_content_list,  # Pass the list
            llm=llm,
            material_goal_1=material_goal_1,
            material_goal_2=material_goal_2,
            iterations_count=iterations_count,
            progress_callback=progress_callback
        )
        
        # Handle the result
        if result["success"]:
            HYPOTHESIS_GENERATION_CONTEXT[extraction_id]['hypotheses'] = result["hypotheses"]
            EMBEDDING_PROGRESS[extraction_id] = {
                "status": "completed",
                "progress": 100,
                "message": result["message"],
                "zip_data": result["zip_data"],
                "filename": result["filename"],
                "iterations_count": iterations_count,
                "hypotheses": result["hypotheses"]
            }
        else:
            EMBEDDING_PROGRESS[extraction_id] = {
                "status": "error",
                "progress": 0,
                "message": result["message"]
            }
        
    except Exception as e:
        EMBEDDING_PROGRESS[extraction_id] = {
            "status": "error",
            "progress": 0,
            "message": f"Error during hypothesis generation: {str(e)}"
        }
        logger.error(f"Hypothesis generation error: {str(e)}")

# ------------  NEW: background QA with chunk-level progress  ------------ #

async def process_qa_async(
    question: str,
    categories: List[str],
    k: int,
    embedding_model,
    llm_QA,
    progress_id: str,
):
    """
    End-to-end QA answer generation that updates EMBEDDING_PROGRESS
    after each chunk is summarised.
    """
    try:
        # ----  gather relevant chunks -------------------------------------------------
        all_docs = []
        for cat in categories:
            vs = get_vectorstore(cat, embedding_model)
            q_emb = await embedding_model.aembed_documents([question])
            k_cat = k  # Use k chunks from each database (not divided across databases)
            res = vs._collection.query(query_embeddings=q_emb, n_results=k_cat)
            for doc_txt, meta in zip(res["documents"][0], res["metadatas"][0]):
                meta["database"] = cat
                all_docs.append(Document(page_content=doc_txt, metadata=meta))

        total_chunks = len(all_docs)
        if not total_chunks:
            EMBEDDING_PROGRESS[progress_id] = {
                "status": "error",
                "progress": 0,
                "message": "No relevant chunks found."
            }
            return

        EMBEDDING_PROGRESS[progress_id] = {
            "status": "processing",
            "progress": 0,
            "total_chunks": total_chunks,
            "processed_chunks": 0,
            "answer": None
        }

        # ----  chunk-wise map ---------------------------------------------------------
        async def summarise(idx, doc):
            from langchain_core.messages import HumanMessage

            response = await llm_QA.ainvoke(
                [HumanMessage(content=f"Summarise this in relation to '{question}':\n{doc.page_content}")]
            )
            summary = llm_content_to_text(getattr(response, "content", response))
            # update progress *after* the chunk is done
            done = EMBEDDING_PROGRESS[progress_id]["processed_chunks"] + 1
            EMBEDDING_PROGRESS[progress_id]["processed_chunks"] = done
            EMBEDDING_PROGRESS[progress_id]["progress"] = int(done / total_chunks * 90)  # keep 10 % for reduce
            return summary

        summaries = await asyncio.gather(
            *[summarise(i, d) for i, d in enumerate(all_docs)]
        )

        # ----  reduce -----------------------------------------------------------------
        from langchain_core.messages import HumanMessage

        final_response = await llm_QA.ainvoke(
            [HumanMessage(content="Combine the following partial answers into one:\n" + "\n".join(summaries))]
        )
        final = llm_content_to_text(getattr(final_response, "content", final_response))
        EMBEDDING_PROGRESS[progress_id] = {
            "status": "completed",
            "progress": 100,
            "answer": final
        }

    except Exception as e:
        EMBEDDING_PROGRESS[progress_id] = {
            "status": "error",
            "progress": 0,
            "message": f"QA failed: {e}"
        }

@app.post("/upload")
async def upload_files(
    background_tasks: BackgroundTasks,
    category: str = Form(...),
    files: List[UploadFile] = File(...),
    ws: str = Depends(workspace),
):
    """
    Upload PDF, text, and image files to the specified category ('database').
    PDFs are converted to .txt files (original PDF is not saved to save storage space).
    Text files are saved as-is. Image files are saved as-is.
    Background embedding tasks are started for text files only.
    Supported file types: .pdf, .txt, .text, .png, .jpg, .jpeg, .gif, .bmp, .tiff, .webp
    """
    global FILE_REGISTRY
    embedding_model = WORKSPACE_MODELS.get(ws, {}).get("embedding")
    # Basic check to ensure we have an embedding_model:
    if embedding_model is None:
        return JSONResponse(
            content={
                "message": "No embedding model set. Call /set_api_key first."
            },
            status_code=400
        )

    display_category, category = category, scoped(ws, category)
    if category not in FILE_REGISTRY:
        FILE_REGISTRY[category] = {}

    # PROTECTION: Check if database is locked before processing files
    category_upload_dir = os.path.join(UPLOAD_DIR, category)
    metadata_file = os.path.join(category_upload_dir, "METADATA")
    if os.path.exists(metadata_file):
        return JSONResponse(
            content={
                "message": f"🔒 Database '{display_category}' is locked. Cannot add files to databases with existing METADATA files. Once a database is created and has a METADATA file, it becomes immutable to prevent data corruption."
            },
            status_code=400
        )

    # Auto-detect single ZIP upload and route through ZIP logic
    if len(files) == 1 and files[0].filename.lower().endswith('.zip'):
        # Reuse the ZIP handler logic inline to keep response identical to /upload_zip
        import io
        import zipfile
        zip_bytes = await files[0].read()
        zf = zipfile.ZipFile(io.BytesIO(zip_bytes))

        saved_file_paths, duplicates = [], []
        uploaded_files_info = []

        for info in zf.infolist():
            if info.is_dir():
                continue
            inner_name = info.filename
            base_name = os.path.basename(inner_name)
            if not base_name:
                continue
            try:
                data = zf.read(info)
            except Exception as e:
                logger.warning(f"⚠️ Skipping '{inner_name}' due to read error: {e}")
                continue

            path, dup = save_file_from_bytes(data, base_name, category)
            if dup:
                duplicates.append(dup)
                continue
            if path:
                saved_file_paths.append(path)

                original_ext = os.path.splitext(base_name.lower())[1]
                if original_ext == '.pdf' or original_ext in ['.txt', '.text']:
                    file_type = "text_pdf"
                elif original_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']:
                    file_type = "image"
                else:
                    file_type = "unknown"

                uploaded_files_info.append({
                    "path": path,
                    "filename": os.path.basename(path),
                    "original_filename": base_name,
                    "type": file_type
                })

        # Separate files by type
        text_pdf_files = [info["path"] for info in uploaded_files_info if info["type"] == "text_pdf"]
        image_files = [info["path"] for info in uploaded_files_info if info["type"] == "image"]

        # Kick off background tasks for embedding (only for PDF/text files)
        if text_pdf_files and embedding_model is not None:
            start_time = time.perf_counter()
            task = asyncio.create_task(process_files_async_all(text_pdf_files, category, embedding_model))
            task.add_done_callback(lambda _, s = start_time: print(f"Embedding time in {(time.perf_counter() - s) * 1000:.2f} ms"))

        return JSONResponse(content={
            "message": "✅ ZIP upload complete.",
            "uploaded": [info["filename"] for info in uploaded_files_info],
            "duplicates": duplicates,
            "file_types": {
                "text_pdf_count": len(text_pdf_files),
                "image_count": len(image_files),
                "has_images": len(image_files) > 0,
                "has_text_pdf": len(text_pdf_files) > 0
            },
            "embedding_started": len(text_pdf_files) > 0 and embedding_model is not None
        })

    tasks = [(save_file(file, category), file.filename) for file in files]
    results = await asyncio.gather(*[task[0] for task in tasks])
    saved_file_paths, duplicates = [], []
    uploaded_files_info = []

    for (path, dup), (_, original_filename) in zip(results, tasks):
        if dup:
            duplicates.append(dup)
        elif path:
            saved_file_paths.append(path)
            
            # Determine file type based on original filename, not the saved path
            # (since PDFs are now converted to .txt files)
            original_ext = os.path.splitext(original_filename.lower())[1]
            if original_ext == '.pdf' or original_ext in ['.txt', '.text']:
                file_type = "text_pdf"
            elif original_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']:
                file_type = "image"
            else:
                file_type = "unknown"
            
            uploaded_files_info.append({
                "path": path,
                "filename": os.path.basename(path),
                "original_filename": original_filename,
                "type": file_type
            })

    # Separate files by type
    text_pdf_files = [info["path"] for info in uploaded_files_info if info["type"] == "text_pdf"]
    image_files = [info["path"] for info in uploaded_files_info if info["type"] == "image"]

    # Kick off background tasks for embedding (only for PDF/text files)
    if text_pdf_files and embedding_model is not None:
        start_time = time.perf_counter()
        task = asyncio.create_task(process_files_async_all(text_pdf_files, category, embedding_model))
        task.add_done_callback(lambda _, s = start_time: print(f"Embedding time in {(time.perf_counter() - s) * 1000:.2f} ms"))

    # Prepare response with file type information
    return JSONResponse(content={
        "message": "Upload complete.",
        "uploaded": [info["filename"] for info in uploaded_files_info],
        "duplicates": duplicates,
        "file_types": {
            "text_pdf_count": len(text_pdf_files),
            "image_count": len(image_files),
            "has_images": len(image_files) > 0,
            "has_text_pdf": len(text_pdf_files) > 0
        },
        "embedding_started": len(text_pdf_files) > 0 and embedding_model is not None,
        "note": "PDFs are converted to text files during upload to save storage space. Only the extracted text is saved."
    })

@app.post("/upload_zip")
async def upload_zip(
    background_tasks: BackgroundTasks,
    category: str = Form(...),
    zip_file: UploadFile = File(...),
    ws: str = Depends(workspace),
):
    """
    Upload a single ZIP file containing PDFs, text files, and/or images.
    Files are extracted server-side and saved using the same rules as /upload.
    Then, background embedding is kicked off for text/PDF-derived files.
    """
    global FILE_REGISTRY
    embedding_model = WORKSPACE_MODELS.get(ws, {}).get("embedding")
    if embedding_model is None:
        return JSONResponse(
            content={"message": "No embedding model set. Please call /set_api_key first."},
            status_code=400
        )

    display_category, category = category, scoped(ws, category)
    if category not in FILE_REGISTRY:
        FILE_REGISTRY[category] = {}

    # PROTECTION: Check if database is locked before processing files
    category_upload_dir = os.path.join(UPLOAD_DIR, category)
    metadata_file = os.path.join(category_upload_dir, "METADATA")
    if os.path.exists(metadata_file):
        return JSONResponse(
            content={
                "message": f"🔒 Database '{display_category}' is locked. Cannot add files to databases with existing METADATA files. Once a database is created and has a METADATA file, it becomes immutable to prevent data corruption."
            },
            status_code=400
        )

    import io
    import zipfile

    # Read ZIP into memory (could be optimized to temp file for very large zips)
    zip_bytes = await zip_file.read()
    zf = zipfile.ZipFile(io.BytesIO(zip_bytes))

    saved_file_paths, duplicates = [], []
    uploaded_files_info = []

    for info in zf.infolist():
        if info.is_dir():
            continue
        inner_name = info.filename
        base_name = os.path.basename(inner_name)
        if not base_name:  # skip entries without a filename
            continue
        try:
            data = zf.read(info)
        except Exception as e:
            logger.warning(f"⚠️ Skipping '{inner_name}' due to read error: {e}")
            continue

        path, dup = save_file_from_bytes(data, base_name, category)
        if dup:
            duplicates.append(dup)
            continue
        if path:
            saved_file_paths.append(path)

            original_ext = os.path.splitext(base_name.lower())[1]
            if original_ext == '.pdf' or original_ext in ['.txt', '.text']:
                file_type = "text_pdf"
            elif original_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']:
                file_type = "image"
            else:
                file_type = "unknown"

            uploaded_files_info.append({
                "path": path,
                "filename": os.path.basename(path),
                "original_filename": base_name,
                "type": file_type
            })

    # Separate files by type
    text_pdf_files = [info["path"] for info in uploaded_files_info if info["type"] == "text_pdf"]
    image_files = [info["path"] for info in uploaded_files_info if info["type"] == "image"]

    # Kick off background tasks for embedding (only for PDF/text files)
    if text_pdf_files and embedding_model is not None:
        start_time = time.perf_counter()
        task = asyncio.create_task(process_files_async_all(text_pdf_files, category, embedding_model))
        task.add_done_callback(lambda _, s = start_time: print(f"Embedding time in {(time.perf_counter() - s) * 1000:.2f} ms"))

    return JSONResponse(content={
        "message": "✅ ZIP upload complete.",
        "uploaded": [info["filename"] for info in uploaded_files_info],
        "duplicates": duplicates,
        "file_types": {
            "text_pdf_count": len(text_pdf_files),
            "image_count": len(image_files),
            "has_images": len(image_files) > 0,
            "has_text_pdf": len(text_pdf_files) > 0
        },
        "embedding_started": len(text_pdf_files) > 0 and embedding_model is not None
    })

@app.get("/available_databases")
async def get_available_databases(ws: str = Depends(workspace)):
    """
    Returns this workspace's databases/categories with their file information.
    """
    global FILE_REGISTRY
    
    databases = []
    for category, files in FILE_REGISTRY.items():
        name = unscoped(ws, category)
        if name is not None and files:  # Only this workspace's categories that have files
            # Count file types in the category
            text_pdf_count = 0
            image_count = 0
            pdf_count = 0
            
            # Load metadata to check original file types (especially for converted PDFs)
            metadata_dict = load_metadata(category)
            
            for file_hash, filename in files.items():
                file_type = get_file_type(filename)
                if file_type == "text_pdf":
                    text_pdf_count += 1
                elif file_type == "image":
                    image_count += 1
                
                # Check for PDFs using metadata (for converted PDFs) or filename (for direct PDFs)
                if filename.lower().endswith('.pdf'):
                    pdf_count += 1
                elif file_hash in metadata_dict:
                    original_filename = metadata_dict[file_hash].get('original_filename', '')
                    if original_filename.lower().endswith('.pdf'):
                        pdf_count += 1
            
            # Determine available modes
            available_modes = []
            if text_pdf_count > 0:
                available_modes.extend(["QA", "System Chart Extraction"])
            
            databases.append({
                "name": name,
                "file_count": len(files),
                "text_pdf_count": text_pdf_count,
                "image_count": image_count,
                "pdf_count": pdf_count,
                "available_modes": available_modes,
                "file_types": {
                    "has_images": image_count > 0,
                    "has_text_pdf": text_pdf_count > 0,
                    "has_pdfs": pdf_count > 0
                }
            })
    
    return JSONResponse(content={"databases": databases})

# --------------------  Conversation Management Endpoints -------------------- #

@app.get("/all_conversations")
async def get_every_conversation(ws: str = Depends(workspace)):
    """
    Get conversations across all of this workspace's extractions (and therefore all
    its databases), so the frontend dropdown is not limited to the most recent run.
    """
    all_convs = get_all_conversations()

    conversation_list = []
    for key, conversations in all_convs.items():
        ext_id = unscoped(ws, key)
        if ext_id is None:
            continue
        for conv_key, conv_data in conversations.items():
            conversation_list.append({
                "extraction_id": ext_id,
                "key": conv_key,
                "title": conv_data["paper_title"],
                "paper_type": conv_data["paper_type"],
                "created_at": conv_data["created_at"],
                "last_updated": conv_data["last_updated"],
                "message_count": len(conv_data["conversation_history"])
            })

    conversation_list.sort(key=lambda x: x["created_at"], reverse=True)

    return JSONResponse(content={
        "conversations": conversation_list,
        "total_count": len(conversation_list)
    })

@app.get("/conversations/{extraction_id}")
async def get_extraction_conversations(extraction_id: str, ws: str = Depends(workspace)):
    """
    Get all conversations for a given extraction ID.
    Returns a list of conversation metadata for frontend dropdown.
    """
    conversations = get_conversations(scoped(ws, extraction_id))
    
    if not conversations:
        return JSONResponse(content={
            "conversations": [],
            "message": "No conversations found for this extraction ID"
        })
    
    # Format conversations for frontend dropdown
    conversation_list = []
    for conv_key, conv_data in conversations.items():
        conversation_list.append({
            "key": conv_key,
            "title": conv_data["paper_title"],
            "paper_type": conv_data["paper_type"],
            "created_at": conv_data["created_at"],
            "last_updated": conv_data["last_updated"],
            "message_count": len(conv_data["conversation_history"])
        })
    
    # Sort by creation time (newest first)
    conversation_list.sort(key=lambda x: x["created_at"], reverse=True)
    
    return JSONResponse(content={
        "conversations": conversation_list,
        "total_count": len(conversation_list)
    })

@app.post("/continue_conversation/{extraction_id}/{conversation_key}")
async def continue_conversation(
    extraction_id: str,
    conversation_key: str,
    question: str = Form(...),
    background_tasks: BackgroundTasks = None,
    ws: str = Depends(workspace),
):
    """
    Continue a conversation from system chart extraction with a follow-up question.
    """
    llm_system_chart_extraction = WORKSPACE_MODELS.get(ws, {}).get("chart")
    if llm_system_chart_extraction is None:
        return JSONResponse(
            content={"error": "No model configured. Please set up a model first."},
            status_code=400
        )
    
    # Get the conversation
    extraction_id = scoped(ws, extraction_id)
    conversations = get_conversations(extraction_id)
    if conversation_key not in conversations:
        return JSONResponse(
            content={"error": "Conversation not found"},
            status_code=404
        )
    
    conversation = conversations[conversation_key]
    
    try:
        # Import the required message classes
        from langchain_core.messages import HumanMessage, AIMessage
        
        # Get the current conversation history
        current_history = conversation["conversation_history"].copy()
        
        # Add the new question with scientific formatting instructions
        enhanced_question = f"""
{question}

Please provide a detailed, scientific answer. Follow these guidelines:
1. Provide quantitative details including exact values, percentages, and units when available
2. When referencing specific information, mention which paper or section it comes from
3. For compositional data, organize information clearly
4. Include relevant scientific context and mechanisms
"""
        
        new_message = HumanMessage(content=enhanced_question)
        current_history.append(new_message)
        
        # Get response from the model
        response = await llm_system_chart_extraction.ainvoke(current_history)
        answer_text = llm_content_to_text(response.content)
        ai_response = AIMessage(content=answer_text)

        # Update the conversation
        update_conversation(extraction_id, conversation_key, new_message, ai_response)

        return JSONResponse(content={
            "question": question,
            "answer": answer_text,
            "paper_title": conversation["paper_title"],
            "paper_type": conversation["paper_type"],
            "conversation_updated": True
        })
        
    except Exception as e:
        logger.error(f"Error continuing conversation: {str(e)}")
        return JSONResponse(
            content={"error": f"Error processing question: {str(e)}"},
            status_code=500
        )

@app.get("/conversation_history/{extraction_id}/{conversation_key}")
async def get_conversation_history(extraction_id: str, conversation_key: str, ws: str = Depends(workspace)):
    """
    Get the full conversation history for a specific conversation.
    """
    conversations = get_conversations(scoped(ws, extraction_id))
    if conversation_key not in conversations:
        return JSONResponse(
            content={"error": "Conversation not found"},
            status_code=404
        )
    
    conversation = conversations[conversation_key]
    
    # Format the conversation history for frontend display
    formatted_history = []
    for i, message in enumerate(conversation["conversation_history"]):
        # Check if it's a human or AI message
        if hasattr(message, 'type'):
            message_type = message.type
            content = message.content
        else:
            # Fallback for different message formats
            message_type = "human" if i % 2 == 0 else "ai"
            content = str(message)
        
        formatted_history.append({
            "type": message_type,
            "content": content,
            "index": i
        })
    
    return JSONResponse(content={
        "paper_title": conversation["paper_title"],
        "paper_type": conversation["paper_type"],
        "created_at": conversation["created_at"],
        "last_updated": conversation["last_updated"],
        "history": formatted_history
    })

# --------------------  Ask Endpoint -------------------- #
@app.post("/ask")
async def ask_question(
    background_tasks: BackgroundTasks,
    categories: str = Form(...),  # Changed from category to categories (comma-separated)
    # Default to "" instead of required: newer FastAPI treats an empty form value as
    # missing, and the hidden question field submits "" in System Chart Extraction mode.
    question: str = Form(""),
    k: int = Form(10),
    goal: str = Form("QA"),
    enable_visualization: bool = Form(False),
    run_system_chart_normalization: bool = Form(False),
    ws: str = Depends(workspace),
):
    """
    Ask a question with different 'goal' modes across multiple databases:
      - QA (Chunk-based Q&A across multiple databases)
      - System Chart Extraction (pulls pure text from PDFs, special method)
    """
    models = WORKSPACE_MODELS.get(ws, {})
    embedding_model, llm_QA, llm_system_chart_extraction = models.get("embedding"), models.get("qa"), models.get("chart")
    if embedding_model is None or llm_QA is None:
        return JSONResponse(
            content={"answer": "No model/embedding set. Please call /set_api_key first."},
            status_code=400
        )

    # Parse categories (comma-separated string)
    category_list = [cat.strip() for cat in categories.split(",") if cat.strip()]
    if not category_list:
        return JSONResponse(content={
            "answer": "No databases selected. Please select at least one database."
        })

    # Validate that all categories exist and have files
    valid_categories = []
    for category in category_list:
        if FILE_REGISTRY.get(scoped(ws, category)):
            valid_categories.append(category)
    
    if not valid_categories:
        return JSONResponse(content={
            "answer": f"No files found in the selected databases: {', '.join(category_list)}"
        })

    # System Chart Extraction can now work with multiple databases
    if goal == "System Chart Extraction":
        # Create a combined extraction ID that represents all databases
        database_names = "_".join(valid_categories)
        extraction_id = f"system_chart_multi_{uuid.uuid4().hex[:8]}"
        print(f"🎯 Created multi-database extraction_id: {extraction_id}")
        print(f"📊 Will process databases: {valid_categories}")
        
        EMBEDDING_PROGRESS[scoped(ws, extraction_id)] = {
            "status": "starting",
            "progress": 0,
            "message": f"Initializing system chart extraction for {len(valid_categories)} databases...",
            "databases": valid_categories,
            "total_databases": len(valid_categories)
        }
        
        # Start background task for multi-database system chart extraction
        background_tasks.add_task(
            process_multi_database_system_chart_extraction_async, 
            ws,
            valid_categories, 
            llm_system_chart_extraction, 
            scoped(ws, extraction_id),
            enable_visualization,
            run_system_chart_normalization
        )
        
        return JSONResponse(content={
            "answer": f"System chart extraction started for {len(valid_categories)} databases: {', '.join(valid_categories)}",
            "extraction_id": extraction_id,
            "check_progress_url": f"/embedding_progress/{extraction_id}",
            "databases": valid_categories,
            "total_databases": len(valid_categories)
        })


    # For QA goal, we can query multiple databases
    if goal == "QA":
        if not question.strip():
            return JSONResponse(
                content={"answer": "Please enter a question before submitting."},
                status_code=400
            )
        if llm_QA is None:
            return JSONResponse(
                content={"answer": "QA model not configured."},
                status_code=400
            )
        
        qa_id = f"qa_{uuid.uuid4().hex[:8]}"
        EMBEDDING_PROGRESS[scoped(ws, qa_id)] = {"status": "starting", "progress": 0}

        # fire-and-forget
        background_tasks.add_task(
            process_qa_async,
            question,
            [scoped(ws, c) for c in valid_categories],
            k,
            embedding_model,
            llm_QA,
            scoped(ws, qa_id)
        )

        return JSONResponse({
            "answer": "Your question is being processed.",
            "extraction_id": qa_id,               # same field name as other tasks
            "check_progress_url": f"/embedding_progress/{qa_id}",
            "databases_searched": valid_categories
        })

    # If we get here, it's an unknown goal
    return JSONResponse(content={
        "answer": f"Unknown goal: {goal}"
    })

# --------------------  Hypotheses Generation Endpoint -------------------- #
@app.post("/generate_hypotheses")
async def generate_hypotheses(
    background_tasks: BackgroundTasks,
    system_charts_files: List[UploadFile] = File(...),
    material_goal_1: str = Form("Property #1 for xxx material"),
    material_goal_2: str = Form("Property #2 for xxx material"),
    iterations_count: int = Form(1),
    ws: str = Depends(workspace),
):
    """
    Generate synergistic scientific hypotheses from multiple uploaded system charts files.
    Accepts multiple system_charts.txt files and generates hypotheses targeting specific material design goals.
    """
    llm_QA = WORKSPACE_MODELS.get(ws, {}).get("qa")
    
    # Check if models are configured
    if llm_QA is None:
        return JSONResponse(
            content={"answer": "No model/embedding set. Please call /set_api_key first."},
            status_code=400
        )
    
    # Use the QA LLM for hypothesis generation
    selected_llm = llm_QA
    
    # Validate that we have at least one file
    if not system_charts_files or len(system_charts_files) == 0:
        return JSONResponse(
            content={"answer": "Please upload at least one .txt file containing system charts with citations."},
            status_code=400
        )
    
    # Validate file types
    invalid_files = [f.filename for f in system_charts_files if not f.filename.endswith('.txt')]
    if invalid_files:
        return JSONResponse(
            content={"answer": f"All files must be .txt files. Invalid files: {', '.join(invalid_files)}"},
            status_code=400
        )
    
    # Validate iterations count
    if iterations_count < 1 or iterations_count > 100:
        return JSONResponse(
            content={"answer": "Number of iterations must be between 1 and 100."},
            status_code=400
        )
    
    try:
        # Read all uploaded file contents
        system_charts_content_list = []
        for file in system_charts_files:
            file_content = await file.read()
            content = file_content.decode('utf-8')
            system_charts_content_list.append(content)
        
        # Validate that all contents are suitable for synergistic hypothesis generation
        validation = HypothesesService.validate_system_charts_content(system_charts_content_list)
        if not validation["valid"]:
            return JSONResponse(
                content={"answer": f"Content validation failed: {validation['message']}"},
                status_code=400
            )
        
        # Initialize progress tracking for hypotheses generation
        extraction_id = f"hypotheses_{uuid.uuid4().hex[:8]}"
        run_key = scoped(ws, extraction_id)
        HYPOTHESIS_GENERATION_CONTEXT[run_key] = {
            'system_charts_content_list': system_charts_content_list,
            'material_goal_1': material_goal_1,
            'material_goal_2': material_goal_2,
            'iterations_count': iterations_count
        }
        print(f"🎯 Created hypotheses extraction_id: {extraction_id}")
        
        EMBEDDING_PROGRESS[run_key] = {
            "status": "starting",
            "progress": 0,
            "message": "Initializing hypothesis generation..."
        }
        print(f"📝 Initial hypotheses progress set for {extraction_id}: {EMBEDDING_PROGRESS[run_key]}")
        
        # Start background task for hypotheses generation
        print(f"🚀 Adding hypotheses background task for {extraction_id}")
        background_tasks.add_task(
            process_hypotheses_generation_async,
            system_charts_content_list,  # Pass the list of contents
            selected_llm,
            run_key,
            material_goal_1,
            material_goal_2,
            iterations_count
        )
        print(f"✅ Hypotheses background task added for {extraction_id}")
        
        file_names = [f.filename for f in system_charts_files]
        return JSONResponse(content={
            "answer": f"Synergistic hypothesis generation started for {len(system_charts_files)} file{'s' if len(system_charts_files) > 1 else ''}: {', '.join(file_names)} targeting '{material_goal_1}' and '{material_goal_2}' with {iterations_count} iteration{'s' if iterations_count > 1 else ''}.",
            "extraction_id": extraction_id,
            "check_progress_url": f"/embedding_progress/{extraction_id}",
            "data_entries": validation.get('data_entries', 0),
            "method": "synergistic_analysis",
            "material_goals": [material_goal_1, material_goal_2],
            "charts_found": validation.get('chart_count', len(system_charts_files)),
            "files_processed": len(system_charts_files)
        })
        
    except UnicodeDecodeError:
        return JSONResponse(
            content={"answer": "Error: Could not read the uploaded file. Please ensure it's a valid UTF-8 text file."},
            status_code=400
        )
    except Exception as e:
        return JSONResponse(
            content={"answer": f"Error processing uploaded file: {str(e)}"},
            status_code=500
        )

@app.post("/analyze_design_goals")
async def analyze_design_goals(description: str = Form(...), ws: str = Depends(workspace)):
    """
    Analyze a free-form materials-design description and suggest the two target
    properties for hypothesis generation. Runs a single LLM call and returns the
    two properties so the frontend can pre-fill the editable goal fields.
    """
    llm_QA = WORKSPACE_MODELS.get(ws, {}).get("qa")
    if llm_QA is None:
        return JSONResponse(
            content={"error": "No model set. Please call /set_api_key first."},
            status_code=400
        )

    description = (description or "").strip()
    if not description:
        return JSONResponse(
            content={"error": "Please describe your materials design goals first."},
            status_code=400
        )

    try:
        from langchain_core.messages import HumanMessage

        prompt = create_design_goals_analysis_prompt(description)
        response = await llm_QA.ainvoke([HumanMessage(content=prompt)])
        raw = llm_content_to_text(response.content).strip()
        print(f"🎯 Design-goals analysis raw output: {raw[:300]}")

        # Parse the strict "Property 1: ... / Property 2: ..." format robustly.
        goal_1, goal_2 = _parse_design_goals(raw)

        if not goal_1 or not goal_2:
            return JSONResponse(
                content={
                    "error": "Could not confidently extract two design properties. Please refine your description or fill in the fields manually.",
                    "raw": raw
                },
                status_code=422
            )

        return JSONResponse(content={
            "success": True,
            "goal_1": goal_1,
            "goal_2": goal_2
        })

    except Exception as e:
        print(f"❌ Error analyzing design goals: {str(e)}")
        return JSONResponse(
            content={"error": f"Error analyzing design goals: {str(e)}"},
            status_code=500
        )

def _parse_design_goals(raw: str):
    """Extract the two properties from the model output.

    Handles the expected 'Property 1: ...' / 'Property 2: ...' format and a few
    lenient fallbacks (numbered lists, bullet points).
    """
    goal_1 = ""
    goal_2 = ""

    def _clean(value: str) -> str:
        return value.strip().strip('"').strip("'").strip("*").strip("-").strip().rstrip(".")

    m1 = re.search(r"Property\s*1\s*[:\-]\s*(.+)", raw, re.IGNORECASE)
    m2 = re.search(r"Property\s*2\s*[:\-]\s*(.+)", raw, re.IGNORECASE)
    if m1:
        goal_1 = _clean(m1.group(1).splitlines()[0])
    if m2:
        goal_2 = _clean(m2.group(1).splitlines()[0])

    # Fallback: take the first two non-empty lines if labels were not found.
    if not (goal_1 and goal_2):
        lines = [l.strip() for l in raw.splitlines() if l.strip()]
        cleaned = []
        for line in lines:
            line = re.sub(r"^\s*(?:property\s*\d+|[0-9]+[.)]|[-*•])\s*[:\-]?\s*", "", line, flags=re.IGNORECASE)
            line = _clean(line)
            if line:
                cleaned.append(line)
        if not goal_1 and len(cleaned) >= 1:
            goal_1 = cleaned[0]
        if not goal_2 and len(cleaned) >= 2:
            goal_2 = cleaned[1]

    return goal_1, goal_2

# --------------------  Hypothesis Continuous Conversation Endpoints -------------------- #
@app.post("/start_hypothesis_conversation")
async def start_hypothesis_conversation(
    extraction_id: str = Form(...),
    hypothesis_number: int = Form(...),
    pasted_papers: str = Form("{}"),
    ws: str = Depends(workspace),
):
    """
    Open a conversation about one generated hypothesis. Resolves the charts it cites
    and their source papers from every database in this workspace. `pasted_papers` is a JSON
    object {chart_number: paper_text} supplied by the user for papers not found.
    Replaces any previous conversation for this run.
    """
    run = HYPOTHESIS_GENERATION_CONTEXT.get(scoped(ws, extraction_id))
    if not run or 'hypotheses' not in run:
        return JSONResponse(content={"error": "Hypothesis run not found. Please generate hypotheses first."}, status_code=404)
    if WORKSPACE_MODELS.get(ws, {}).get("qa") is None:
        return JSONResponse(content={"error": "LLM not configured. Please set up the model first."}, status_code=400)

    try:
        pasted = {int(k): v for k, v in json.loads(pasted_papers).items() if str(v).strip()}
        context, sources = build_hypothesis_context(ws, run, hypothesis_number, pasted)
    except ValueError as e:
        return JSONResponse(content={"error": str(e)}, status_code=400)
    except Exception as e:
        return JSONResponse(content={"error": f"Error starting hypothesis conversation: {str(e)}"}, status_code=500)

    run.update({'conversation_context': context, 'conversation_history': []})
    print(f"💬 Hypothesis #{hypothesis_number} conversation started for {extraction_id}: {len(context)} chars, sources={sources}")
    return JSONResponse(content={"success": True, "hypothesis_number": hypothesis_number, "sources": sources})


@app.post("/hypothesis_conversation_message")
async def hypothesis_conversation_message(
    extraction_id: str = Form(...),
    message: str = Form(...),
    ws: str = Depends(workspace),
):
    """
    Send a message in the hypothesis conversation and get an AI response with full context.
    """
    run = HYPOTHESIS_GENERATION_CONTEXT.get(scoped(ws, extraction_id))
    llm_QA = WORKSPACE_MODELS.get(ws, {}).get("qa")
    if not run or 'conversation_context' not in run:
        return JSONResponse(content={"error": "Hypothesis conversation not found. Please click Discuss on a hypothesis first."}, status_code=404)
    if llm_QA is None:
        return JSONResponse(content={"error": "LLM not configured. Please set up the model first."}, status_code=400)

    try:
        history = run['conversation_history']
        prompt = run['conversation_context'] + "\nCONVERSATION HISTORY:\n"
        for entry in history[-10:]:
            prompt += f"Human: {entry['human']}\nAI: {entry['ai']}\n\n"
        prompt += f"""CURRENT QUESTION:
Human: {message}

Please provide a detailed, scientific answer based on the research context above. Follow these guidelines:
1. Include specific information from the source papers and system charts
2. Cite sources with the same number used in the hypothesis, e.g. "Paper [2]" or "Chart [2]"
3. For compositional data, organize information clearly
4. Provide quantitative details including exact values, percentages, and units when available
5. Explain the scientific significance of the data

AI: """
        print(f"🤖 Sending hypothesis prompt to LLM ({len(prompt)} chars)")
        from langchain_core.messages import HumanMessage
        response = await llm_QA.ainvoke([HumanMessage(content=prompt)])
        ai_response = llm_content_to_text(response.content)

        history.append({'human': message, 'ai': ai_response})
        return JSONResponse(content={"success": True, "response": ai_response, "conversation_length": len(history)})

    except Exception as e:
        print(f"❌ Error processing hypothesis conversation message: {str(e)}")
        return JSONResponse(content={"error": f"Error processing message: {str(e)}"}, status_code=500)

def get_admin_password() -> Optional[str]:
    password = os.environ.get("ADMIN_PASSWORD")
    if password:
        return password.strip()
    return None

def get_admin_wipe_targets() -> List[str]:
    return [
        os.path.abspath(UPLOAD_DIR),
        os.path.abspath(CHROMA_DB_BASE),
    ]

def is_safe_target(path: str, data_root: str) -> bool:
    normalized_path = os.path.abspath(path)
    normalized_root = os.path.abspath(data_root)

    if normalized_path in {"/", ""}:
        return False

    try:
        return os.path.commonpath([normalized_root, normalized_path]) == normalized_root
    except ValueError:
        return False

def delete_directory_contents(path: str) -> Dict[str, int]:
    deleted_files = 0
    deleted_dirs = 0
    if not os.path.isdir(path):
        return {"deleted_files": 0, "deleted_dirs": 0}
    for name in os.listdir(path):
        target = os.path.join(path, name)
        if os.path.isfile(target) or os.path.islink(target):
            os.unlink(target)
            deleted_files += 1
        elif os.path.isdir(target):
            shutil.rmtree(target)
            deleted_dirs += 1
    return {"deleted_files": deleted_files, "deleted_dirs": deleted_dirs}

@app.post("/admin/verify_password")
async def admin_verify_password(password: str = Form(...)):
    admin_password = get_admin_password()
    if not admin_password:
        return JSONResponse(
            content={"ok": False, "message": "ADMIN_PASSWORD is not configured on server."},
            status_code=500
        )

    if not hmac.compare_digest(password or "", admin_password):
        return JSONResponse(
            content={"ok": False, "message": "Invalid admin password."},
            status_code=403
        )

    return JSONResponse(content={"ok": True, "message": "Password verified."}, status_code=200)

@app.post("/admin/wipe_data")
async def admin_wipe_data(
    password: str = Form(...),
    confirmation_text: str = Form(...)
):
    admin_password = get_admin_password()
    if not admin_password:
        return JSONResponse(
            content={"ok": False, "message": "ADMIN_PASSWORD is not configured on server."},
            status_code=500
        )

    if not hmac.compare_digest(password or "", admin_password):
        return JSONResponse(
            content={"ok": False, "message": "Invalid admin password."},
            status_code=403
        )

    if (confirmation_text or "").strip() != ADMIN_DELETE_CONFIRMATION_TEXT:
        return JSONResponse(
            content={
                "ok": False,
                "message": f"Invalid confirmation text. Type exactly: {ADMIN_DELETE_CONFIRMATION_TEXT}"
            },
            status_code=400
        )

    targets = get_admin_wipe_targets()
    data_root = os.path.abspath(APP_DATA_ROOT)

    for target in targets:
        if not is_safe_target(target, data_root):
            return JSONResponse(
                content={
                    "ok": False,
                    "message": f"Refusing to wipe unsafe target: {target}",
                    "data_root": data_root,
                    "targets": targets
                },
                status_code=400
            )

    try:
        results = []
        total_files = 0
        total_dirs = 0

        for target in targets:
            if not os.path.exists(target):
                results.append({"target": target, "status": "missing"})
                continue

            result = delete_directory_contents(target)
            total_files += result["deleted_files"]
            total_dirs += result["deleted_dirs"]
            results.append({"target": target, "status": "wiped", **result})

        FILE_REGISTRY.clear()
        EMBEDDING_PROGRESS.clear()
        HYPOTHESIS_GENERATION_CONTEXT.clear()

        logger.warning(
            "Admin data wipe executed for targets=%s (files=%s, dirs=%s)",
            targets,
            total_files,
            total_dirs
        )
        return JSONResponse(
            content={
                "ok": True,
                "message": "Wipe completed.",
                "data_root": data_root,
                "targets": results,
                "deleted_files": total_files,
                "deleted_dirs": total_dirs
            },
            status_code=200
        )
    except Exception as e:
        logger.exception("Admin data wipe failed: %s", str(e))
        return JSONResponse(
            content={"ok": False, "message": f"Failed to wipe data: {str(e)}"},
            status_code=500
        )

# --------------------  Hypothesis conversation helpers -------------------- #
HYPOTHESIS_PAPER_CHAR_CAP = 60_000     # per source paper
HYPOTHESIS_CONTEXT_CHAR_CAP = 250_000  # all source papers combined


def load_all_papers(ws: str) -> List[tuple]:
    """The workspace's papers as (database, title, text), databases in alphabetical order."""
    return [(unscoped(ws, key), title, text) for key in sorted(FILE_REGISTRY) if unscoped(ws, key) is not None
            for title, text in obtain_text_from_files(key).items()]


def match_paper(title: str, papers: List[tuple]):
    """Find a paper by title: exact or substring first, then fuzzy (score >= 80). Returns (database, text)."""
    title = (title or "").strip()
    best = (0, None, None)
    for db, stored, text in papers:
        if title and (title == stored or title in stored or stored in title):
            return db, text
        score = fuzz.token_set_ratio(title, stored) if title else 0
        if score > best[0]:
            best = (score, db, text)
    return (best[1], best[2]) if best[0] >= 80 else (None, None)


def build_hypothesis_context(ws: str, run: dict, hypothesis_number: int, pasted: Dict[int, str]):
    """
    Build the fixed conversation context for one hypothesis: the hypothesis text plus,
    for every chart it cites as [n], that chart and its source paper (full text, capped).
    Returns (context_text, sources) where sources describes how each citation resolved.
    """
    hypothesis = next((h for h in run['hypotheses'] if h['number'] == hypothesis_number), None)
    if hypothesis is None:
        raise ValueError(f"Hypothesis #{hypothesis_number} not found in this run.")

    charts = [(m.group('paper').strip(), m.group('chart').strip())
              for blob in run['system_charts_content_list'] for m in CHART_PATTERN.finditer(blob)]
    cited = sorted({int(n) for n in re.findall(r'\[(\d+)\]', hypothesis['text']) if 0 < int(n) <= len(charts)})
    papers = load_all_papers(ws) if any(n not in pasted for n in cited) else []

    context = f"=== HYPOTHESIS #{hypothesis_number} UNDER DISCUSSION ===\n{hypothesis['text']}\n\n"
    sources, budget = [], HYPOTHESIS_CONTEXT_CHAR_CAP
    for n in cited:
        title, chart = charts[n - 1]
        db, text = ("pasted", pasted[n]) if n in pasted else match_paper(title, papers)
        truncated = False
        if text:
            allowed = min(HYPOTHESIS_PAPER_CHAR_CAP, budget)
            truncated = len(text) > allowed
            text = text[:allowed]
            budget -= len(text)
        sources.append({'number': n, 'title': title, 'database': db, 'truncated': truncated})
        context += (f"=== SOURCE [{n}]: {title} ===\n"
                    f"CHART [{n}]:\n{chart}\n\n"
                    f"PAPER [{n}]:\n{text or 'NOT AVAILABLE'}\n\n")
    context += "=== END OF CONTEXT ===\n\n"
    return context, sources


if __name__ == "__main__":
    import uvicorn
    
    # Log important startup information for diagnostics
    logger.info("=" * 50)
    logger.info("Starting PDF-RAG application")
    logger.info(f"Running as executable: {hasattr(sys, '_MEIPASS')}")
    logger.info(f"Current working directory: {os.getcwd()}")
    if hasattr(sys, '_MEIPASS'):
        logger.info(f"PyInstaller _MEIPASS: {sys._MEIPASS}")
        logger.info(f"Executable path: {sys.executable}")
    logger.info(f"Base directory for storage: {get_base_directory()}")
    logger.info(f"UPLOAD_DIR: {os.path.abspath(UPLOAD_DIR)}")
    logger.info(f"CHROMA_DB_BASE: {os.path.abspath(CHROMA_DB_BASE)}")
    logger.info("=" * 50)
       
    print("Starting server...")

    uvicorn.run(app, host="127.0.0.1", port=8000)