processing_pipeline.py

# --- File: ./processing_pipeline.py (UPDATED FOR SQLITE-VEC - SYNTAX FIXED) ---
import sqlite3
import hashlib
import traceback
import os
import datetime
import re
import email
import json
import fnmatch
from email.header import decode_header
from email.utils import parsedate_to_datetime
from pathlib import Path
from itertools import combinations
from typing import Union, List

import fitz
import spacy
from bs4 import BeautifulSoup
import ollama
import numpy as np
import sqlite_vec

# --- Configuration ---
# FIXED: Importing absolute paths directly from config to prevent worker displacement
from project.config import EMBEDDING_MODEL, resolve_document_path, DATABASE_FILE, DOCUMENTS_DIR

SPACY_MODEL = "en_core_web_lg"
CHUNK_SIZE = 400
CHUNK_OVERLAP = 50

NLP_MODEL = None

def load_spacy_model():
    """Loads the spaCy model into the global variable if not already loaded."""
    global NLP_MODEL
    if NLP_MODEL is None:
        print(f"Loading spaCy model '{SPACY_MODEL}' in process {os.getpid()}...")
        try:
            disabled_pipes = ["lemmatizer"]
            NLP_MODEL = spacy.load(SPACY_MODEL, disable=disabled_pipes)
            print(f"spaCy model loaded successfully in process {os.getpid()}.")
        except OSError:
            print(f"FATAL: spaCy model '{SPACY_MODEL}' not found.")
            print(f"Please run: python -m spacy download {SPACY_MODEL}")
            raise
    return NLP_MODEL

def get_db_conn():
    """Gets a fresh, process-safe database connection."""
    conn = sqlite3.connect(DATABASE_FILE, timeout=30)
    
    # --- Load sqlite-vec extension for native vector search ---
    conn.enable_load_extension(True)
    sqlite_vec.load(conn)
    conn.enable_load_extension(False)
    
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON;")
    conn.execute("PRAGMA journal_mode = WAL;")
    return conn

def _paginate_text(text, words_per_page=300):
    """Splits a string of text into pages of roughly N words."""
    words = text.split()
    page_content_map = {}
    for i in range(0, len(words), words_per_page):
        page_number = (i // words_per_page) + 1
        page_content_map[page_number] = " ".join(words[i:i + words_per_page])
    if not page_content_map:
        page_content_map[1] = ""
    return page_content_map

def extract_text_for_copying(doc_path: Path, file_type: str, start_page=None, end_page=None, doc_id=None) -> str:
    if file_type in ['TXT', 'HTML', 'SRT', 'EML']:
        conn = None
        try:
            conn = get_db_conn()
            
            # If doc_id wasn't provided, fall back to the old path-guessing logic
            if not doc_id:
                if doc_path.is_absolute() and not str(doc_path).startswith(str(DOCUMENTS_DIR)):
                     pass 
                else:
                     rel_path_str = str(doc_path.relative_to(DOCUMENTS_DIR).as_posix())
                     doc_id_row = conn.execute("SELECT id FROM documents WHERE relative_path = ?", (rel_path_str,)).fetchone()
                     if doc_id_row:
                         doc_id = doc_id_row['id']
            
            if not doc_id:
                return f"Error: Document not found in database for path {doc_path}"
            
            # For these types, the database treats them as a single continuous block of text (Page 1)
            if file_type in ['HTML', 'SRT', 'EML']:
                start_page = None
                end_page = None

            query = "SELECT page_content FROM content_index WHERE doc_id = ?"
            params = [doc_id]
            
            if start_page is not None and end_page is not None and start_page <= end_page:
                query += " AND page_number BETWEEN ? AND ?"
                params.extend([start_page, end_page])
            elif start_page is not None:
                query += " AND page_number = ?"
                params.append(start_page)
            
            query += " ORDER BY page_number ASC"
            
            pages = [row['page_content'] for row in conn.execute(query, params).fetchall()]
            return "\n\n".join(pages)

        except Exception as e:
            print(f"Error extracting text from DB for {file_type} file {doc_path}: {e}")
            return f"Error extracting text from database: {e}"
        finally:
            if conn:
                conn.close()

    elif file_type == 'PDF':
        full_text = []
        try:
            doc = fitz.open(doc_path)
            first_page_index = max(0, start_page - 1) if start_page is not None else 0
            last_page_index = doc.page_count - 1
            if end_page is not None:
                last_page_index = min(doc.page_count - 1, end_page - 1)
            elif start_page is not None:
                last_page_index = first_page_index
            if first_page_index > last_page_index:
                return ""
            for page_num in range(first_page_index, last_page_index + 1):
                page = doc.load_page(page_num)
                text = page.get_text("text", sort=True)
                if text:
                    full_text.append(text.strip())
            doc.close()
        except Exception as e:
            print(f"Error extracting text from PDF {doc_path}: {e}")
            return f"Error extracting text: {e}"
        return "\n\n".join(full_text)
    else:
        return f"Cannot extract text from unsupported file type: {file_type}"

def _extract_text_from_pipermail(html_content: str) -> str:
    soup = BeautifulSoup(html_content, 'lxml')
    title, subject, author, body = soup.find('h1'), soup.find('b'), soup.find('i'), soup.find('pre')
    parts = []
    if title: parts.append(f"List: {title.get_text(strip=True)}")
    if subject: parts.append(f"Subject: {subject.get_text(strip=True)}")
    if author: parts.append(f"Author: {author.get_text(strip=True)}")
    if body: parts.append(f"--- Message Body ---\n{body.get_text()}")
    return "\n\n".join(parts)

def _extract_text_with_block_separation(html_content: str) -> str:
    soup = BeautifulSoup(html_content, 'lxml')
    for script_or_style in soup(["script", "style"]):
        script_or_style.decompose()
    text_blocks = []
    title_tag = soup.find('title')
    if title_tag and title_tag.get_text(strip=True):
        text_blocks.append(f"Title: {title_tag.get_text(strip=True)}")
    for tag in soup.find_all(['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'li', 'blockquote', 'pre', 'tr', 'div']):
        text = tag.get_text(separator=' ', strip=True)
        if text:
            text_blocks.append(text)
    return "\n\n".join(text_blocks)

def _extract_text_from_srt(srt_content: str) -> str:
    lines = []
    for block in srt_content.strip().split('\n\n'):
        parts = block.strip().split('\n')
        if len(parts) >= 3:
            lines.append(" ".join(parts[2:]))
    return "\n".join(lines)

def _parse_srt_for_db(srt_content: str) -> list[dict]:
    cues = []
    cue_pattern = re.compile(r'(\d+)\s*(\d{2}:\d{2}:\d{2},\d{3}\s*-->\s*\d{2}:\d{2}:\d{2},\d{3})\s*(.*?)\n\n', re.DOTALL)
    for match in cue_pattern.finditer(srt_content + '\n\n'):
        sequence, timestamp, dialogue_block = match.groups()
        clean_dialogue = re.sub(r'<[^>]+>', '', dialogue_block).replace('\n', ' ').strip()
        cues.append({'sequence': int(sequence), 'timestamp': timestamp.strip(), 'dialogue': clean_dialogue})
    return cues

def _get_srt_duration(srt_content: str) -> Union[int, None]:
    timestamps = re.findall(r'\d{2}:\d{2}:\d{2},\d{3}\s*-->\s*(\d{2}:\d{2}:\d{2},\d{3})', srt_content)
    if not timestamps: return None
    try:
        h, m, s, ms = map(int, re.split('[:,]', timestamps[-1]))
        return h * 3600 + m * 60 + s + ms / 1000
    except (ValueError, IndexError):
        return None

def _extract_text_from_pdf_doc(doc: fitz.Document) -> dict:
    page_content_map = {}
    try:
        for page_num, page in enumerate(doc, 1):
            text = page.get_text("text", sort=True)
            if text: page_content_map[page_num] = text.strip()
    except Exception as e:
        print(f"Error extracting text from PDF: {e}")
    return page_content_map

def _chunk_text(text: str, size: int, overlap: int) -> List[str]:
    if not text: return []
    words = text.split()
    chunks = []
    for i in range(0, len(words), size - overlap):
        chunks.append(" ".join(words[i:i + size]))
    return chunks

def _generate_embeddings_for_page(page_text: str) -> list[tuple[str, bytes]]:
    if not page_text.strip(): return []
    chunks = _chunk_text(page_text, CHUNK_SIZE, CHUNK_OVERLAP)
    embeddings_to_store = []
    for chunk in chunks:
        try:
            response = ollama.embeddings(model=EMBEDDING_MODEL, prompt=chunk)
            embedding_vector = response['embedding']
            embedding_blob = np.array(embedding_vector, dtype=np.float32).tobytes()
            embeddings_to_store.append((chunk, embedding_blob))
        except Exception as e:
            print(f"WORKER WARNING: Could not generate embedding for a chunk. Error: {e}")
            continue
    return embeddings_to_store

def _decode_header_text(header_value):
    """Decodes email headers to handle different charsets."""
    if not header_value:
        return ""
    decoded_parts = []
    for part, charset in decode_header(header_value):
        if isinstance(part, bytes):
            decoded_parts.append(part.decode(charset or 'utf-8', 'ignore'))
        else:
            decoded_parts.append(part)
    return "".join(decoded_parts)

def _parse_eml_content(eml_bytes: bytes) -> dict:
    """Parses an .eml file's content and returns structured metadata and body text."""
    msg = email.message_from_bytes(eml_bytes)
    
    metadata = {
        'from_address': _decode_header_text(msg.get('From')),
        'to_addresses': _decode_header_text(msg.get('To')),
        'cc_addresses': _decode_header_text(msg.get('Cc')),
        'subject': _decode_header_text(msg.get('Subject')),
        'sent_at': None
    }
    
    date_str = msg.get('Date')
    if date_str:
        try:
            metadata['sent_at'] = parsedate_to_datetime(date_str)
        except (ValueError, TypeError):
            pass

    body_plain, body_html = "", ""
    if msg.is_multipart():
        for part in msg.walk():
            content_type = part.get_content_type()
            if content_type == 'text/plain' and not body_plain:
                try:
                    body_plain = part.get_payload(decode=True).decode(part.get_content_charset() or 'utf-8', 'ignore')
                except (LookupError, AttributeError):
                    pass
            elif content_type == 'text/html' and not body_html:
                 try:
                    body_html = part.get_payload(decode=True).decode(part.get_content_charset() or 'utf-8', 'ignore')
                 except (LookupError, AttributeError):
                    pass
    else:
        try:
            body_plain = msg.get_payload(decode=True).decode(msg.get_content_charset() or 'utf-8', 'ignore')
        except (LookupError, AttributeError):
            pass

    final_body = body_plain.strip()
    if not final_body and body_html:
        soup = BeautifulSoup(body_html, 'lxml')
        final_body = _extract_text_with_block_separation(str(soup))
        
    return {'metadata': metadata, 'body': final_body}

def _extract_data_from_pages(page_content_map: dict) -> dict:
    nlp = load_spacy_model()
    data_to_store = {"entities": set(), "appearances": set(), "relationships": [], "content": [], "super_chunks": []}
    
    for page_num, page_text in page_content_map.items():
        data_to_store["content"].append((page_num, page_text))
        if not page_text or len(page_text) > nlp.max_length: continue
        
        doc_nlp = nlp(page_text)
        
        for ent in doc_nlp.ents:
            entity_tuple = (ent.text.strip(), ent.label_)
            if entity_tuple[0]:
                data_to_store["entities"].add(entity_tuple)
                data_to_store["appearances"].add((entity_tuple, page_num))
        
        for sent in doc_nlp.sents:
            unique_ents = list(dict.fromkeys(sent.ents))
            if len(unique_ents) < 2: continue
            for ent1, ent2 in combinations(unique_ents, 2):
                start, end = min(ent1.end_char, ent2.end_char), max(ent1.start_char, ent2.start_char)
                if end > start and (end - start) < 75:
                    phrase = ' '.join(page_text[start:end].strip().split())
                    if phrase:
                        subj = (ent1.text.strip(), ent1.label_) if ent1.start_char < ent2.start_char else (ent2.text.strip(), ent2.label_)
                        obj = (ent2.text.strip(), ent2.label_) if ent1.start_char < ent2.start_char else (ent1.text.strip(), ent1.label_)
                        if subj[0] and obj[0]:
                             data_to_store["relationships"].append((subj, obj, phrase, page_num))

        # New "Super Embedding" Chunk Logic
        for ent in doc_nlp.ents:
            if not ent.text.strip(): continue
            verb = ent.root.head
            if verb.pos_ in ("VERB", "AUX"):
                action_phrase = " ".join(token.text for token in verb.subtree).strip().replace('\n', ' ').replace('  ', ' ')
                if action_phrase:
                    super_chunk_text = f"{ent.text} ({ent.label_}): {action_phrase}"
                    entity_tuple = (ent.text.strip(), ent.label_)
                    data_to_store["super_chunks"].append({
                        "entity": entity_tuple,
                        "page_number": page_num,
                        "chunk_text": super_chunk_text
                    })

    return data_to_store

def process_document(doc_id):
    """Worker function using the original, stable, self-contained transaction model."""
    conn = None
    try:
        conn = get_db_conn()
        cursor = conn.cursor()

        conn.execute("UPDATE documents SET status = ?, status_message = ? WHERE id = ?", ('Indexing', 'Worker process started...', doc_id))
        conn.commit()

        doc_info = conn.execute("SELECT relative_path, file_type FROM documents WHERE id = ?", (doc_id,)).fetchone()
        if not doc_info:
            raise ValueError(f"No document found with ID: {doc_id}")

        print(f"--- Worker {os.getpid()} processing Doc ID: {doc_id} (Type: {doc_info['file_type']}) ---")
        
        full_path = resolve_document_path(doc_info['relative_path'])
        
        page_content_map, page_count, duration_seconds = {}, 0, None
        extracted_data = {"embeddings": [], "super_chunks": []}
        csl_json_text = None

        eml_meta_to_insert = None 

        if doc_info['file_type'] == 'SRT':
            content = full_path.read_text(encoding='utf-8', errors='ignore')
            parsed_cues = _parse_srt_for_db(content)
            duration_seconds = _get_srt_duration(content)
            page_count = len(parsed_cues)
            
            nlp = load_spacy_model()
            extracted_data.update({"entities": set(), "appearances": set(), "relationships": [], "content": [], "cues": parsed_cues})
            
            SRT_CHUNK_SIZE_CUES, SRT_CHUNK_OVERLAP_CUES = 20, 5
            for i in range(0, len(parsed_cues), SRT_CHUNK_SIZE_CUES - SRT_CHUNK_OVERLAP_CUES):
                chunk_cues = parsed_cues[i:i + SRT_CHUNK_SIZE_CUES]
                if not chunk_cues: continue
                chunk_dialogue = " ".join(c['dialogue'] for c in chunk_cues)
                try:
                    response = ollama.embeddings(model=EMBEDDING_MODEL, prompt=chunk_dialogue)
                    embedding_blob = np.array(response['embedding'], dtype=np.float32).tobytes()
                    extracted_data["embeddings"].append((doc_id, chunk_cues[0]['sequence'], chunk_dialogue, embedding_blob))
                except Exception as e:
                    print(f"WORKER WARNING: Could not generate embedding for an SRT chunk. Error: {e}")
            
            full_dialogue_text = " ".join(c['dialogue'] for c in parsed_cues)
            extracted_data["content"].append((1, full_dialogue_text))

            if full_dialogue_text and len(full_dialogue_text) <= nlp.max_length:
                cue_char_boundaries = []
                current_pos = 0
                for cue in parsed_cues:
                    dialogue_len = len(cue['dialogue'])
                    cue_char_boundaries.append((current_pos, current_pos + dialogue_len))
                    current_pos += dialogue_len + 1
                doc_nlp_full = nlp(full_dialogue_text)
                for ent in doc_nlp_full.ents:
                    ent_tuple = (ent.text.strip(), ent.label_)
                    if ent_tuple[0]: extracted_data["entities"].add(ent_tuple)
                    for i, (cue_start, cue_end) in enumerate(cue_char_boundaries):
                        if ent.start_char >= cue_start and ent.start_char < cue_end:
                            extracted_data["appearances"].add((ent_tuple, parsed_cues[i]['sequence']))
                            break
                for sent in doc_nlp_full.sents:
                    unique_ents = list(dict.fromkeys(sent.ents))
                    if len(unique_ents) < 2: continue
                    for ent1, ent2 in combinations(unique_ents, 2):
                        start, end = min(ent1.end_char, ent2.end_char), max(ent1.start_char, ent2.start_char)
                        if end > start and (end - start) < 75:
                            phrase = ' '.join(full_dialogue_text[start:end].strip().split())
                            if phrase:
                                subj = (ent1.text.strip(), ent1.label_) if ent1.start_char < ent2.start_char else (ent2.text.strip(), ent2.label_)
                                obj = (ent2.text.strip(), ent2.label_) if ent1.start_char < ent2.start_char else (ent1.text.strip(), ent1.label_)
                                if subj[0] and obj[0]:
                                    rel_start_char = min(ent1.start_char, ent2.start_char)
                                    for i, (cue_start, cue_end) in enumerate(cue_char_boundaries):
                                        if rel_start_char >= cue_start and rel_start_char < cue_end:
                                            extracted_data["relationships"].append((subj, obj, phrase, parsed_cues[i]['sequence']))
                                            break
        elif doc_info['file_type'] == 'EML':
            eml_bytes = full_path.read_bytes()
            parsed_eml = _parse_eml_content(eml_bytes)
            page_content_map = {1: parsed_eml['body']}
            page_count = 1
            
            for page_num, page_text in page_content_map.items():
                for chunk_text, embedding_blob in _generate_embeddings_for_page(page_text):
                    extracted_data["embeddings"].append((doc_id, page_num, chunk_text, embedding_blob))
            extracted_data.update(_extract_data_from_pages(page_content_map))
            
            eml_meta = parsed_eml['metadata']

            csl_data = {
                "id": f"doc-{doc_id}",
                "type": "personal_communication",
                "title": eml_meta.get('subject'),
                "medium": "Email",
                "author": [{"literal": eml_meta.get('from_address')}],
                "recipient": [{"literal": eml_meta.get('to_addresses')}],
                "issued": None
            }
            if eml_meta.get('sent_at'):
                dt = eml_meta['sent_at']
                csl_data['issued'] = {'date-parts': [[dt.year, dt.month, dt.day]]}
            
            csl_json_text = json.dumps(csl_data, indent=2)

            eml_meta_to_insert = (doc_id, eml_meta['from_address'], eml_meta['to_addresses'], eml_meta['cc_addresses'], eml_meta['subject'], eml_meta['sent_at'])

        else:
            if doc_info['file_type'] == 'PDF':
                with fitz.open(full_path) as pdf_doc:
                    page_count = pdf_doc.page_count
                    page_content_map = _extract_text_from_pdf_doc(pdf_doc)
            elif doc_info['file_type'] == 'TXT':
                content = full_path.read_text(encoding='utf-8', errors='ignore')
                page_content_map = _paginate_text(content.strip())
                page_count = len(page_content_map)
            elif doc_info['file_type'] == 'HTML':
                mode_row = conn.execute("SELECT value FROM app_settings WHERE key = 'html_parsing_mode'").fetchone()
                content = full_path.read_text(encoding='utf-8', errors='ignore')
                extracted_text = _extract_text_from_pipermail(content) if mode_row and mode_row[0] == 'pipermail' else _extract_text_with_block_separation(content)
                page_content_map = {1: extracted_text.strip()} if extracted_text.strip() else {1: ""}
                page_count = 1 if extracted_text.strip() else 0
            
            for page_num, page_text in page_content_map.items():
                for chunk_text, embedding_blob in _generate_embeddings_for_page(page_text):
                    extracted_data["embeddings"].append((doc_id, page_num, chunk_text, embedding_blob))

            extracted_data.update(_extract_data_from_pages(page_content_map))

        # --- DATABASE WRITE PHASE (ALL IN ONE TRANSACTION) ---
        cursor.execute("BEGIN TRANSACTION;")
        cursor.execute("DELETE FROM srt_cues WHERE doc_id = ?", (doc_id,))
        cursor.execute("DELETE FROM content_index WHERE doc_id = ?", (doc_id,))
        cursor.execute("DELETE FROM entity_appearances WHERE doc_id = ?", (doc_id,))
        cursor.execute("DELETE FROM entity_relationships WHERE doc_id = ?", (doc_id,))
        
        # This will trigger the cascading deletes in vec_embedding_chunks due to our SQLite triggers
        cursor.execute("DELETE FROM embedding_chunks WHERE doc_id = ?", (doc_id,))
        cursor.execute("DELETE FROM super_embedding_chunks WHERE doc_id = ?", (doc_id,))
        
        cursor.execute("DELETE FROM email_metadata WHERE doc_id = ?", (doc_id,))
        cursor.execute("DELETE FROM document_metadata WHERE doc_id = ?", (doc_id,))

        cursor.execute("UPDATE documents SET page_count = ?, duration_seconds = ? WHERE id = ?", (page_count, duration_seconds, doc_id))

        if eml_meta_to_insert:
            cursor.execute("""
                INSERT INTO email_metadata (doc_id, from_address, to_addresses, cc_addresses, subject, sent_at)
                VALUES (?, ?, ?, ?, ?, ?)
            """, eml_meta_to_insert)

        if extracted_data.get("cues"):
            cues_to_insert = [(doc_id, c['sequence'], c['timestamp'], c['dialogue']) for c in extracted_data["cues"]]
            cursor.executemany("INSERT INTO srt_cues (doc_id, sequence, timestamp, dialogue) VALUES (?, ?, ?, ?)", cues_to_insert)

        if extracted_data.get("content"):
            cursor.executemany("INSERT INTO content_index (doc_id, page_number, page_content) VALUES (?, ?, ?)", [(doc_id, pn, pt) for pn, pt in extracted_data["content"]])
        
        # --- NEW: Write to sqlite-vec virtual tables ---
        if extracted_data.get("embeddings"):
            for d_id, page_num, chunk_text, embedding_blob in extracted_data["embeddings"]:
                cursor.execute(
                    "INSERT INTO embedding_chunks (doc_id, page_number, chunk_text) VALUES (?, ?, ?)", 
                    (d_id, page_num, chunk_text)
                )
                chunk_id = cursor.lastrowid
                cursor.execute(
                    "INSERT INTO vec_embedding_chunks (chunk_id, embedding) VALUES (?, ?)",
                    (chunk_id, embedding_blob)
                )
        
        if csl_json_text:
            cursor.execute("""
                INSERT INTO document_metadata (doc_id, csl_json, last_updated) VALUES (?, ?, CURRENT_TIMESTAMP)
            """, (doc_id, csl_json_text))

        if extracted_data.get("entities"):
            entities_list = list(extracted_data["entities"])
            cursor.executemany("INSERT OR IGNORE INTO entities (text, label) VALUES (?, ?)", entities_list)
            
            entity_id_map = {}
            for text, label in entities_list:
                res = conn.execute("SELECT id FROM entities WHERE text = ? AND label = ?", (text, label)).fetchone()
                if res: entity_id_map[(text, label)] = res['id']

            appearances_to_insert = [(doc_id, entity_id_map[ent_tuple], page_num) for ent_tuple, page_num in extracted_data["appearances"] if ent_tuple in entity_id_map]
            if appearances_to_insert:
                cursor.executemany("INSERT OR IGNORE INTO entity_appearances (doc_id, entity_id, page_number) VALUES (?, ?, ?)", appearances_to_insert)
            
            relationships_to_insert = [
                (entity_id_map[subj], entity_id_map[obj], phrase, doc_id, page_num) 
                for subj, obj, phrase, page_num in extracted_data.get("relationships", []) 
                if subj in entity_id_map and obj in entity_id_map
            ]
            
            if relationships_to_insert:
                cursor.executemany(
                    "INSERT INTO entity_relationships (subject_entity_id, object_entity_id, relationship_phrase, doc_id, page_number) VALUES (?, ?, ?, ?, ?)", 
                    relationships_to_insert
                )
            
            if extracted_data.get("super_chunks"):
                for chunk_data in extracted_data["super_chunks"]:
                    entity_tuple = chunk_data["entity"]
                    if entity_tuple in entity_id_map:
                        entity_id = entity_id_map[entity_tuple]
                        chunk_text = chunk_data["chunk_text"]
                        try:
                            response = ollama.embeddings(model=EMBEDDING_MODEL, prompt=chunk_text)
                            embedding_blob = np.array(response['embedding'], dtype=np.float32).tobytes()
                            
                            # Insert metadata and get ID
                            cursor.execute(
                                "INSERT INTO super_embedding_chunks (doc_id, page_number, entity_id, chunk_text) VALUES (?, ?, ?, ?)",
                                (doc_id, chunk_data["page_number"], entity_id, chunk_text)
                            )
                            chunk_id = cursor.lastrowid
                            # Insert vector into vec0
                            cursor.execute(
                                "INSERT INTO vec_super_embedding_chunks (chunk_id, embedding) VALUES (?, ?)",
                                (chunk_id, embedding_blob)
                            )
                        except Exception as e:
                            print(f"WORKER WARNING: Could not generate super embedding for chunk '{chunk_text[:50]}...'. Error: {e}")
                            continue

        cursor.execute("UPDATE documents SET status = ?, status_message = ?, processed_at = CURRENT_TIMESTAMP WHERE id = ?", ('Indexed', 'Processing complete.', doc_id))
        conn.commit()
        print(f"--- Worker {os.getpid()} finished Doc ID: {doc_id} ---")
        return "SUCCESS"

    except Exception as e:
        print(f"!!! WORKER {os.getpid()} ERROR on Doc ID: {doc_id} !!!")
        print(traceback.format_exc())
        if conn:
            try:
                conn.rollback()
                error_message = f"Worker Error: {type(e).__name__}: {e}"
                conn.execute("UPDATE documents SET status = ?, status_message = ? WHERE id = ?", ('Error', error_message[:1000], doc_id))
                conn.commit()
            except Exception as rollback_e:
                print(f"CRITICAL: Failed to rollback or update error status for doc {doc_id}. Error: {rollback_e}")
        raise
    finally:
        if conn:
            conn.close()

def _gather_files_recursively(base_dir: Path, target_dir: Path, supported_patterns: list, virtual_prefix: str = "") -> list:
    """
    Helper function to gather files, allowing for virtual prefix mapping if
    scanning an external directory mapped by an .rlink file.
    """
    files = []
    try:
        # Use os.walk instead of rglob to gracefully bypass Windows PermissionErrors
        for root, dirs, filenames in os.walk(target_dir):
            for filename in filenames:
                if any(fnmatch.fnmatch(filename.lower(), p.lower()) for p in supported_patterns):
                    file_path = Path(root) / filename
                    
                    try:
                        if virtual_prefix:
                            rel_path_str = f"{virtual_prefix}/{file_path.relative_to(target_dir).as_posix()}"
                        else:
                            rel_path_str = file_path.relative_to(base_dir).as_posix()
                            
                        files.append((file_path, rel_path_str))
                    except ValueError:
                        continue # Skip if relative_to fails
    except Exception as e:
        print(f"[WARN] Error scanning directory {target_dir}: {e}")
        
    return files

def discover_and_register_documents():
    """Scans the source directory (and any .rlink aliases), registers new files, and moves missing ones to the Recycle Bin."""
    print(f"--- Scanning for documents in {DOCUMENTS_DIR} ---")
    conn = get_db_conn() 
    try:
        # 1. Fetch current DB state
        db_docs = conn.execute("SELECT id, relative_path, file_hash, status FROM documents").fetchall()
        db_files = {row['relative_path']: row['file_hash'] for row in db_docs}
        db_statuses = {row['relative_path']: row['status'] for row in db_docs}
        db_path_to_id = {row['relative_path']: row['id'] for row in db_docs}
        
        # --- FIX: Case-insensitive lookup map for cross-OS imports ---
        db_files_lower = {k.lower(): k for k in db_files.keys()}
        
        registered_count = 0
        restored_count = 0
        supported_patterns = ["*.pdf", "*.txt", "*.html", "*.srt", "*.eml"]
        
        all_files_with_virtual_paths = _gather_files_recursively(DOCUMENTS_DIR, DOCUMENTS_DIR, supported_patterns)
        
        for rlink_file in DOCUMENTS_DIR.glob('*.rlink'):
            if rlink_file.is_file():
                try:
                    target_path_str = rlink_file.read_text(encoding='utf-8').strip()
                    target_dir = Path(target_path_str).resolve() # Ensure path is normalized
                    if target_dir.exists() and target_dir.is_dir():
                        print(f"  [INFO] Following alias '{rlink_file.name}' to: {target_dir}")
                        all_files_with_virtual_paths.extend(
                            _gather_files_recursively(DOCUMENTS_DIR, target_dir, supported_patterns, virtual_prefix=rlink_file.name)
                        )
                    else:
                        print(f"  [WARN] Alias target does not exist or is not a directory: {target_dir}")
                except Exception as e:
                    print(f"  [WARN] Failed to read alias file {rlink_file}: {e}")
            
        found_paths_exact = set() 

        for file_path, rel_path_str in all_files_with_virtual_paths:
            
            # --- FIX: Align OS path casing with Database casing ---
            lower_rel = rel_path_str.lower()
            if lower_rel in db_files_lower:
                rel_path_str = db_files_lower[lower_rel]
                
            found_paths_exact.add(rel_path_str) 
            
            stats = file_path.stat()
            current_size = stats.st_size
            current_mtime = datetime.datetime.fromtimestamp(stats.st_mtime)
            file_type = file_path.suffix[1:].upper()

            hasher = hashlib.md5()
            with open(file_path, 'rb') as f:
                buf = f.read(65536)
                while len(buf) > 0:
                    hasher.update(buf)
                    buf = f.read(65536)
            current_hash_str = hasher.hexdigest()

            if rel_path_str not in db_files or db_files.get(rel_path_str) != current_hash_str:
                print(f"Registering new/modified file: {rel_path_str} (Type: {file_type})")
                conn.execute(
                    """
                    INSERT INTO documents (relative_path, file_hash, file_type, status, status_message, file_size_bytes, file_modified_at, processed_at) 
                    VALUES (?, ?, ?, 'New', 'Ready for processing', ?, ?, CURRENT_TIMESTAMP)
                    ON CONFLICT(relative_path) 
                    DO UPDATE SET file_hash=excluded.file_hash, 
                                  file_type=excluded.file_type,
                                  status='New', 
                                  status_message='File modified, ready for re-processing', 
                                  file_size_bytes=excluded.file_size_bytes,
                                  file_modified_at=excluded.file_modified_at,
                                  processed_at=excluded.processed_at;
                    """, 
                    (rel_path_str, current_hash_str, file_type, current_size, current_mtime)
                )
                registered_count += 1
            else:
                if db_statuses.get(rel_path_str) == 'Missing':
                    print(f"Restoring previously missing file: {rel_path_str}")
                    conn.execute("UPDATE documents SET status = 'Indexed', status_message = 'Restored from Recycle Bin' WHERE relative_path = ?", (rel_path_str,))
                    restored_count += 1

        # 2. IDENTIFY MISSING FILES (Soft Delete to Recycle Bin)
        missing_paths = set(db_files.keys()) - found_paths_exact
        missing_count = 0

        if missing_paths:
            for path in missing_paths:
                if db_statuses.get(path) != 'Missing':
                    doc_id = db_path_to_id[path]
                    print(f"  Moving document to Recycle Bin: {path}")
                    conn.execute("UPDATE documents SET status = 'Missing', status_message = 'File removed from directory or alias disconnected. View in Settings > Recycle Bin.' WHERE id = ?", (doc_id,))
                    missing_count += 1

        conn.commit()
        print(f"--- Discovery complete. Registered {registered_count}. Restored {restored_count}. Trashed {missing_count}. ---")
        return "SUCCESS"
        
    except Exception as e:
        print(f"!!! ERROR during document discovery: {e} !!!")
        print(traceback.format_exc())
        try: conn.rollback()
        except sqlite3.Error: pass
    finally:
        if conn: conn.close()

def update_browse_cache():
    """Recomputes the aggregated entity data, IGNORING files in the recycle bin."""
    print("--- Starting update of Aggregated View Cache ---")
    conn = get_db_conn() 
    try:
        conn.execute("BEGIN")
        print("Executing aggregation query...")
        query = """
            SELECT 
                e.id as entity_id, 
                e.text as entity_text, 
                e.label as entity_label, 
                COUNT(DISTINCT ea.doc_id) as document_count,
                COUNT(ea.doc_id) as appearance_count
            FROM entities e 
            JOIN entity_appearances ea ON e.id = ea.entity_id 
            JOIN documents d ON ea.doc_id = d.id
            WHERE d.status != 'Missing'
            GROUP BY e.id, e.text, e.label
        """
        aggregated_data = conn.execute(query).fetchall()
        print(f"Found {len(aggregated_data)} unique entities with document counts.")
        print("Updating browse_cache table...")
        
        conn.execute("DELETE FROM browse_cache")
        if aggregated_data:
            print(f"Inserting {len(aggregated_data)} rows into the new cache...")
            insert_query = """
                INSERT INTO browse_cache (entity_id, entity_text, entity_label, document_count, appearance_count) 
                VALUES (?, ?, ?, ?, ?)
            """
            conn.executemany(insert_query, [
                (row['entity_id'], row['entity_text'], row['entity_label'], row['document_count'], row['appearance_count']) 
                for row in aggregated_data
            ])
        conn.commit()
        print("--- Aggregated View Cache update finished successfully. ---")
    except Exception as e:
        print(f"!!! ERROR updating browse cache: {e} !!!")
        print(traceback.format_exc())
        conn.rollback()
    finally:
        if conn: conn.close()