contract-processor.py

"""
Author: Martin Bacigal, 01/2025 @ https://procureai.tech
License: MIT License
"""
import tkinter as tk
from tkinter import filedialog, messagebox, ttk, scrolledtext
from threading import Thread
from dataclasses import dataclass
from typing import Optional, Dict, List, Any, Union
from pathlib import Path
import os
import asyncio
from concurrent.futures import ThreadPoolExecutor
import logging
from functools import partial, lru_cache
import json
from enum import Enum
import mimetypes
import queue
from datetime import datetime
import psutil
import gc

import pandas as pd
from pypdf import PdfReader
import docx
import mammoth
from pyxlsb import open_workbook as open_xlsb
import pythoncom
from win32com.client import Dispatch
from openai import AzureOpenAI, OpenAI
from tenacity import retry, stop_after_attempt, wait_exponential
from dataclasses import dataclass, field
from functools import lru_cache
from openpyxl import load_workbook
from openpyxl.utils.dataframe import dataframe_to_rows
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

# Import our custom modules
from database_manager import DatabaseManager, DatabaseConfig, Document, Company
from settings_manager import SettingsManager, SettingsDialog, AppSettings
from memory_manager import (
    MemoryConfig, MemoryMonitor, MemoryEfficientDocumentProcessor,
    get_system_memory_info, estimate_file_memory_usage
)

# Enhanced Logging Setup
class LogHandler(logging.Handler):
    def __init__(self, text_widget):
        super().__init__()
        self.text_widget = text_widget
        self.queue = queue.Queue()
        self.text_widget.tag_configure('ERROR', foreground='red')
        self.text_widget.tag_configure('WARNING', foreground='orange')
        self.text_widget.tag_configure('INFO', foreground='green')
        self.update_widget()

    def emit(self, record):
        self.queue.put(record)

    def update_widget(self):
        while True:
            try:
                record = self.queue.get_nowait()
                self.text_widget.insert('end',
                    f"{datetime.now().strftime('%H:%M:%S')} - {record.getMessage()}\n",
                    record.levelname)
                self.text_widget.see('end')
            except queue.Empty:
                break
        self.text_widget.after(100, self.update_widget)

@dataclass(frozen=True)
class OpenAIConfig:
    api_key: str = field(default_factory=lambda: os.getenv("OPENAI_API_KEY", "OPENAI_API_KEY"))

    @lru_cache(maxsize=1)
    def create_client(self):
        if not self.api_key:
            logging.error("OpenAI API key is not set.")
            raise ValueError("OpenAI API key is not set.")

        try:
            client = OpenAI(api_key=self.api_key)
            logging.info("OpenAI client configured successfully.")
            return client
        except Exception as e:
            logging.error(f"Failed to configure OpenAI client: {e}")
            raise

@dataclass(frozen=True)
class AzureConfig:
    endpoint: str = field(default_factory=lambda: os.getenv("AZURE_ENDPOINT", "AZURE_ENDPOINT_URL"))
    deployment: str = field(default_factory=lambda: os.getenv("AZURE_DEPLOYMENT", "AZURE_DEVELOPMENT"))
    api_key: str = field(default_factory=lambda: os.getenv("AZURE_API_KEY", "AZURE_API_KEY"))
    api_version: str = "API_VERSION"

    @lru_cache(maxsize=1)
    def create_client(self):
        if not self.endpoint or not self.api_key:
            logging.error("Azure endpoint or API key is not set.")
            raise ValueError("Azure endpoint or API key is not set.")

        from openai import AzureOpenAI
        try:
            client = AzureOpenAI(
                azure_endpoint=self.endpoint,
                api_key=self.api_key,
                api_version=self.api_version
            )
            logging.info("AzureOpenAI client created successfully.")
            return client
        except Exception as e:
            logging.error(f"Failed to create AzureOpenAI client: {e}")
            raise

@dataclass(frozen=True)
class EndpointConfig:
    use_azure: bool = False  # Set to False to use OpenAI instead of Azure
    azure_config: AzureConfig = field(default_factory=AzureConfig)
    openai_config: OpenAIConfig = field(default_factory=OpenAIConfig)

    def create_client(self):
        if self.use_azure:
            return self.azure_config.create_client()
        else:
            return self.openai_config.create_client()

@dataclass
class ProcessingConfig:
    MAX_WORKERS: int = os.cpu_count() or 4
    BATCH_SIZE: int = 4
    OUTPUT_DIR: Path = Path("output")
    LOG_LEVEL: str = "DEBUG"

class DocumentType(Enum):
    PDF = "application/pdf"
    DOCX = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
    DOC = "application/msword"
    XLSX = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    XLS = "application/vnd.ms-excel"

class ContentAnalyzer:
    def __init__(self, endpoint_config: EndpointConfig):
        self.client = endpoint_config.create_client()
        self.deployment = endpoint_config.azure_config.deployment if endpoint_config.use_azure else None
        self.analysis_templates = {
            "contract_details": {
                "system": """Analyze the contract and extract the following information:
                1. Start date: in the format DD/MM/YYYY
                2. End date: in the format DD/MM/YYYY
                3. Contract duration
                4. Key deliverables: per each key deliverable start in new line and with "- "

                Make sure that you always reply for each point within the same line.
                Avoid markdown, as the result is used for further processing.
                """,
                "max_tokens": 300
            },
            "vendor_assessment": {
                "system": """Evaluate the vendor document and provide:
                1. Vendor capabilities
                2. Risk factors
                3. Compliance status
                4. Past performance metrics
                5. Key strengths and weaknesses""",
                "max_tokens": 400
            },
            "technical_specs": {
                "system": """Extract and analyze technical specifications including:
                1. Technical requirements
                2. Performance metrics
                3. Compatibility requirements
                4. Implementation timeline
                5. Maintenance requirements""",
                "max_tokens": 350
            }
        }

    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
    async def analyze_content(self, text: str, analysis_type: str) -> Dict[str, Any]:
        template = self.analysis_templates.get(analysis_type)
        if not template:
            raise ValueError(f"Unknown analysis type: {analysis_type}")

        try:
            completion = await asyncio.to_thread(
                self.client.chat.completions.create,
                model=self.deployment,
                messages=[
                    {"role": "system", "content": template["system"]},
                    {"role": "user", "content": text}
                ],
                max_tokens=template["max_tokens"],
                temperature=0.7
            )
        except Exception as e:
            logging.error(f"Error analyzing content: {e}")
            raise

        return {
            "analysis_type": analysis_type,
            "result": completion.choices[0].message.content
        }

    def format_information(self, result_content: str, analysis_type: str) -> Dict[str, Any]:
        formatted_info = {}
        key_deliverables = []
        lines = result_content.split('\n')
        key_deliverables_flag = False

        for line in lines:
            if "Start date" in line:
                formatted_info["start_date"] = line.split(":")[1].strip()
            elif "End date" in line:
                formatted_info["end_date"] = line.split(":")[1].strip()
            elif "Contract duration" in line:
                formatted_info["contract_duration"] = line.split(":")[1].strip()
            elif "Key deliverables" in line:
                key_deliverables_flag = True
            elif "Vendor capabilities" in line:
                formatted_info["vendor_capabilities"] = line.split(":")[1].strip()
                key_deliverables_flag = False
            elif "Risk factors" in line:
                formatted_info["risk_factors"] = line.split(":")[1].strip()
                key_deliverables_flag = False
            elif "Compliance status" in line:
                formatted_info["compliance_status"] = line.split(":")[1].strip()
                key_deliverables_flag = False
            elif "Past performance metrics" in line:
                formatted_info["past_performance_metrics"] = line.split(":")[1].strip()
                key_deliverables_flag = False
            elif "Key strengths and weaknesses" in line:
                formatted_info["key_strengths_weaknesses"] = line.split(":")[1].strip()
                key_deliverables_flag = False
            elif "Technical requirements" in line:
                formatted_info["technical_requirements"] = line.split(":")[1].strip()
                key_deliverables_flag = False
            elif "Performance metrics" in line:
                formatted_info["performance_metrics"] = line.split(":")[1].strip()
                key_deliverables_flag = False
            elif "Compatibility requirements" in line:
                formatted_info["compatibility_requirements"] = line.split(":")[1].strip()
                key_deliverables_flag = False
            elif "Implementation timeline" in line:
                formatted_info["implementation_timeline"] = line.split(":")[1].strip()
                key_deliverables_flag = False
            elif "Maintenance requirements" in line:
                formatted_info["maintenance_requirements"] = line.split(":")[1].strip()
                key_deliverables_flag = False
            elif key_deliverables_flag:
                if line.strip().startswith("-"):
                    key_deliverables.append(line.strip().lstrip("-").strip())

        if key_deliverables:
            formatted_info["key_deliverables"] = key_deliverables

        return formatted_info

class DocumentProcessor:
    def __init__(self, config: ProcessingConfig, azure_config: AzureConfig):
        self.config = config
        self.azure_config = azure_config
        self.analyzer = ContentAnalyzer(azure_config)
        self.executor = ThreadPoolExecutor(max_workers=config.MAX_WORKERS)

        # Define processors for different MIME types
        self.processors = {
            DocumentType.PDF.value: self._extract_pdf_content,
            DocumentType.DOCX.value: self._extract_word_content,
            DocumentType.DOC.value: self._extract_word_content,
            DocumentType.XLSX.value: self._extract_excel_content,
            DocumentType.XLS.value: self._extract_excel_content
        }

    async def process_document(self, file_path: Path) -> Dict[str, Any]:
        mime_type, _ = mimetypes.guess_type(str(file_path))
        processor = self.processors.get(mime_type)

        if not processor:
            logging.error(f"Unsupported file type: {mime_type}")
            return {"error": f"Unsupported file type: {mime_type}"}

        try:
            content = await processor(file_path, mime_type)
            if not content:
                logging.error(f"Error extracting content from {file_path}")
                return {"error": f"Error extracting content from {file_path}"}

            analyses = await asyncio.gather(
                self.analyzer.analyze_content(content, "contract_details"),
                self.analyzer.analyze_content(content, "vendor_assessment"),
                self.analyzer.analyze_content(content, "technical_specs")
            )

            return {
                "file_path": str(file_path),
                "metadata": {"mime_type": mime_type},
                "analyses": analyses
            }
        except Exception as e:
            logging.error(f"Processing error for {file_path}: {e}")
            return {"error": f"Processing error for {file_path}: {e}"}

    async def _extract_pdf_content(self, file_path: Path, mime_type: str) -> Optional[str]:
        try:
            reader = PdfReader(str(file_path))
            text = ""
            for page in reader.pages:
                text += page.extract_text()
            return text
        except Exception as e:
            logging.error(f"Error extracting PDF content from {file_path}: {e}")
            return None

    async def _extract_word_content(self, file_path: Path, mime_type: str) -> Optional[str]:
        try:
            if mime_type == DocumentType.DOCX.value:
                doc = docx.Document(str(file_path))
                return "\n".join([para.text for para in doc.paragraphs])
            elif mime_type == DocumentType.DOC.value:
                with open(file_path, "rb") as doc_file:
                    result = mammoth.extract_raw_text(doc_file)
                    return result.value
            return None
        except Exception as e:
            logging.error(f"Error extracting Word content from {file_path}: {e}")
            return None

    async def _extract_excel_content(self, file_path: Path) -> Optional[str]:
        try:
            # Implement extraction logic for Excel files (XLSX, XLS)
            # This is a placeholder implementation and should be replaced with actual logic
            return "Extracted Excel content"
        except Exception as e:
            logging.error(f"Error extracting Excel content from {file_path}: {e}")
            return None

    async def process_directory(self, directory: Path, update_progress, selected_file_types: List[str], update_status, batch_size: int = 10) -> List[Dict[str, Any]]:
        files = [f for f in directory.rglob("*") if f.is_file() and
                 mimetypes.guess_type(str(f))[0] in selected_file_types]

        results = []
        total_files = len(files)
        processed = 0

        output_path = self.config.OUTPUT_DIR / "processed_results.xlsx"
        output_path.parent.mkdir(parents=True, exist_ok=True)

        for batch in [files[i:i + batch_size] for i in range(0, len(files), batch_size)]:
            batch_results = await asyncio.gather(
                *[self.process_document(f) for f in batch]
            )
            batch_results = [r for r in batch_results if r]  # Filter out None results

            await ResultsExporter.to_excel(batch_results, output_path, self.analyzer, update_progress)

            results.extend(batch_results)
            processed += len(batch)
            update_progress(processed, total_files)

        return results

class EnhancedDocumentProcessorApp(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("Enhanced Document Processor with SQL & Knowledge Graph")
        self.geometry("1200x800")
        self.configure(bg='#f0f0f0')

        # Initialize settings manager
        self.settings_manager = SettingsManager()

        # Initialize database manager
        self.db_manager = None
        self.init_database()

        # Initialize memory monitor
        self.memory_config = MemoryConfig(
            memory_limit_mb=self.settings_manager.settings.memory_limit_mb,
            max_batch_size=self.settings_manager.settings.batch_size
        )
        self.memory_monitor = MemoryMonitor(self.memory_config)
        self.memory_monitor.add_callback(self._update_memory_display)

        # Processing state
        self.processing = False
        self.cancel_requested = False

        # File type selection variables
        self.file_types = {}
        for file_type, enabled in self.settings_manager.settings.file_types.items():
            self.file_types[file_type] = tk.BooleanVar(value=enabled)

        self._setup_styles()
        self._create_widgets()
        self._setup_layout()
        self._create_menu()

        # Start monitoring
        self.memory_monitor.start_monitoring()
        self._start_monitoring()

        # Bind closing event
        self.protocol("WM_DELETE_WINDOW", self._on_closing)

    def init_database(self):
        """Initialize database connection"""
        try:
            db_config = self.settings_manager.get_database_config()
            self.db_manager = DatabaseManager(db_config)
            self.db_manager.initialize()
            logging.info("Database initialized successfully")
        except Exception as e:
            logging.error(f"Failed to initialize database: {e}")
            messagebox.showwarning("Database Warning",
                                 "Failed to initialize database. SQL features will be disabled.")

    def _setup_styles(self):
        style = ttk.Style()
        style.configure('TFrame', background='#f0f0f0')
        style.configure('Header.TLabel',
                        font=('Helvetica', 16, 'bold'),
                        padding=10,
                        background='#f0f0f0')
        style.configure('Status.TLabel',
                        font=('Helvetica', 10),
                        padding=5,
                        background='#f0f0f0')
        style.configure('Action.TButton',
                        font=('Helvetica', 10, 'bold'),
                        padding=5)
        style.configure('TCheckbutton',
                        background='#f0f0f0',
                        font=('Helvetica', 10))
        style.configure('FileType.TLabelframe.Label',
                        font=('Helvetica', 12, 'bold'))

    def _create_menu(self):
        """Create application menu"""
        menubar = tk.Menu(self)
        self.config(menu=menubar)

        # File menu
        file_menu = tk.Menu(menubar, tearoff=0)
        menubar.add_cascade(label="File", menu=file_menu)
        file_menu.add_command(label="Settings", command=self._open_settings)
        file_menu.add_separator()
        file_menu.add_command(label="Exit", command=self._on_closing)

        # Database menu
        db_menu = tk.Menu(menubar, tearoff=0)
        menubar.add_cascade(label="Database", menu=db_menu)
        db_menu.add_command(label="View Processed Documents", command=self._view_processed_documents)
        db_menu.add_command(label="Export to SQL", command=self._export_to_sql)
        db_menu.add_separator()
        db_menu.add_command(label="Build Knowledge Graph", command=self._build_knowledge_graph)

        # Tools menu
        tools_menu = tk.Menu(menubar, tearoff=0)
        menubar.add_cascade(label="Tools", menu=tools_menu)
        tools_menu.add_command(label="Memory Usage Report", command=self._show_memory_report)
        tools_menu.add_command(label="Clear Cache", command=self._clear_cache)

        # Help menu
        help_menu = tk.Menu(menubar, tearoff=0)
        menubar.add_cascade(label="Help", menu=help_menu)
        help_menu.add_command(label="About", command=self._show_about)

    def _create_widgets(self):
        # Create notebook for tabbed interface
        self.notebook = ttk.Notebook(self)
        self.notebook.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        # Main processing tab
        self.main_tab = ttk.Frame(self.notebook)
        self.notebook.add(self.main_tab, text="Document Processing")

        self.main_container = ttk.Frame(self.main_tab)
        self.main_container.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        self.header = ttk.Label(
            self.main_container,
            text="Document Processing System with SQL & Knowledge Graph",
            style='Header.TLabel'
        )

        # Directory selection with recent directories
        self.dir_frame = ttk.Frame(self.main_container)
        self.dir_label = ttk.Label(
            self.dir_frame,
            text="Select Directory:",
            style='Status.TLabel'
        )
        self.dir_var = tk.StringVar()
        self.dir_combo = ttk.Combobox(self.dir_frame, textvariable=self.dir_var, width=60)
        self.dir_combo['values'] = self.settings_manager.recent_directories
        self.browse_button = ttk.Button(
            self.dir_frame,
            text="Browse",
            style='Action.TButton',
            command=self._browse_directory
        )

        # Company metadata frame
        self.metadata_frame = ttk.LabelFrame(self.main_container, text="Document Metadata", padding=10)

        ttk.Label(self.metadata_frame, text="CW Number:").grid(row=0, column=0, sticky=tk.W, padx=5)
        self.cw_var = tk.StringVar(value=self.settings_manager.settings.default_cw_prefix)
        ttk.Entry(self.metadata_frame, textvariable=self.cw_var, width=20).grid(row=0, column=1, padx=5)

        ttk.Label(self.metadata_frame, text="Company ID:").grid(row=0, column=2, sticky=tk.W, padx=5)
        self.company_id_var = tk.StringVar(value=self.settings_manager.settings.default_company_id)
        ttk.Entry(self.metadata_frame, textvariable=self.company_id_var, width=20).grid(row=0, column=3, padx=5)

        ttk.Label(self.metadata_frame, text="Company Group:").grid(row=1, column=0, sticky=tk.W, padx=5)
        self.company_group_var = tk.StringVar(value=self.settings_manager.settings.default_company_group)
        ttk.Entry(self.metadata_frame, textvariable=self.company_group_var, width=20).grid(row=1, column=1, padx=5)

        # File type selection
        self.file_type_frame = ttk.LabelFrame(self.main_container, text="File Types to Process", style='FileType.TLabelframe')
        self.checkbuttons = {}
        for file_type, var in self.file_types.items():
            cb = ttk.Checkbutton(self.file_type_frame, text=file_type, variable=var, style='TCheckbutton')
            cb.pack(side=tk.LEFT, padx=10, pady=5)
            self.checkbuttons[file_type] = cb

        # Enhanced Progress Frame with detailed stats
        self.progress_container = ttk.LabelFrame(self.main_container, text="Processing Progress", padding=10)

        # Status and progress
        self.status_frame = ttk.Frame(self.progress_container)
        self.status_label = ttk.Label(
            self.status_frame,
            text="Ready",
            style='Status.TLabel'
        )

        # Progress bars
        self.progress_frame = ttk.Frame(self.progress_container)

        # Overall progress
        ttk.Label(self.progress_frame, text="Overall:").grid(row=0, column=0, sticky=tk.W, padx=5)
        self.progress_bar = ttk.Progressbar(
            self.progress_frame,
            mode='determinate',
            length=400
        )
        self.progress_bar.grid(row=0, column=1, padx=5, pady=2)
        self.progress_label = ttk.Label(
            self.progress_frame,
            text="0%",
            style='Status.TLabel'
        )
        self.progress_label.grid(row=0, column=2, padx=5)

        # Batch progress
        ttk.Label(self.progress_frame, text="Current Batch:").grid(row=1, column=0, sticky=tk.W, padx=5)
        self.batch_progress_bar = ttk.Progressbar(
            self.progress_frame,
            mode='determinate',
            length=400
        )
        self.batch_progress_bar.grid(row=1, column=1, padx=5, pady=2)
        self.batch_label = ttk.Label(
            self.progress_frame,
            text="Batch: 0/0",
            style='Status.TLabel'
        )
        self.batch_label.grid(row=1, column=2, padx=5)

        # Statistics frame
        self.stats_frame = ttk.Frame(self.progress_container)
        self.processed_label = ttk.Label(self.stats_frame, text="Processed: 0", style='Status.TLabel')
        self.processed_label.grid(row=0, column=0, padx=10)
        self.skipped_label = ttk.Label(self.stats_frame, text="Skipped: 0", style='Status.TLabel')
        self.skipped_label.grid(row=0, column=1, padx=10)
        self.failed_label = ttk.Label(self.stats_frame, text="Failed: 0", style='Status.TLabel')
        self.failed_label.grid(row=0, column=2, padx=10)
        self.time_label = ttk.Label(self.stats_frame, text="Time: 00:00:00", style='Status.TLabel')
        self.time_label.grid(row=0, column=3, padx=10)

        # System monitoring
        self.monitor_frame = ttk.Frame(self.main_container)
        self.cpu_label = ttk.Label(
            self.monitor_frame,
            text="CPU: 0%",
            style='Status.TLabel'
        )
        self.memory_label = ttk.Label(
            self.monitor_frame,
            text="Memory: 0%",
            style='Status.TLabel'
        )
        self.memory_process_label = ttk.Label(
            self.monitor_frame,
            text="Process: 0MB",
            style='Status.TLabel'
        )
        self.db_status_label = ttk.Label(
            self.monitor_frame,
            text="DB: " + ("Connected" if self.db_manager else "Disconnected"),
            style='Status.TLabel'
        )

        # Log display
        self.log_frame = ttk.Frame(self.main_container)
        self.log_text = scrolledtext.ScrolledText(
            self.log_frame,
            height=8,
            width=70
        )

        # Buttons
        self.button_frame = ttk.Frame(self.main_container)
        self.process_button = ttk.Button(
            self.button_frame,
            text="Start Processing",
            style='Action.TButton',
            command=self._start_processing
        )
        self.cancel_button = ttk.Button(
            self.button_frame,
            text="Cancel",
            style='Action.TButton',
            state=tk.DISABLED,
            command=self._cancel_processing
        )
        self.settings_button = ttk.Button(
            self.button_frame,
            text="Settings",
            style='Action.TButton',
            command=self._open_settings
        )

        # Ontology Management Tab
        self.ontology_tab = ttk.Frame(self.notebook)
        self.notebook.add(self.ontology_tab, text="Contract Ontology")

        # Ontology tree view
        self.ontology_paned = ttk.PanedWindow(self.ontology_tab, orient=tk.HORIZONTAL)
        self.ontology_paned.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        # Left panel - Tree view
        self.tree_frame = ttk.LabelFrame(self.ontology_paned, text="Ontology Hierarchy")
        self.ontology_tree = ttk.Treeview(self.tree_frame, selectmode='browse')
        self.ontology_tree.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        tree_scroll = ttk.Scrollbar(self.tree_frame, orient=tk.VERTICAL, command=self.ontology_tree.yview)
        tree_scroll.pack(side=tk.RIGHT, fill=tk.Y)
        self.ontology_tree.configure(yscrollcommand=tree_scroll.set)

        self.ontology_paned.add(self.tree_frame, weight=2)

        # Right panel - Details and actions
        self.ontology_details_frame = ttk.LabelFrame(self.ontology_paned, text="Category Details")

        # Category info
        details_grid = ttk.Frame(self.ontology_details_frame)
        details_grid.pack(fill=tk.X, padx=10, pady=10)

        ttk.Label(details_grid, text="Code:").grid(row=0, column=0, sticky=tk.W, padx=5, pady=2)
        self.ont_code_var = tk.StringVar()
        ttk.Entry(details_grid, textvariable=self.ont_code_var, width=30).grid(row=0, column=1, padx=5, pady=2)

        ttk.Label(details_grid, text="Name:").grid(row=1, column=0, sticky=tk.W, padx=5, pady=2)
        self.ont_name_var = tk.StringVar()
        ttk.Entry(details_grid, textvariable=self.ont_name_var, width=30).grid(row=1, column=1, padx=5, pady=2)

        ttk.Label(details_grid, text="Description:").grid(row=2, column=0, sticky=tk.W, padx=5, pady=2)
        self.ont_desc_text = tk.Text(details_grid, height=3, width=40)
        self.ont_desc_text.grid(row=2, column=1, padx=5, pady=2)

        ttk.Label(details_grid, text="Color:").grid(row=3, column=0, sticky=tk.W, padx=5, pady=2)
        self.ont_color_var = tk.StringVar()
        color_frame = ttk.Frame(details_grid)
        color_frame.grid(row=3, column=1, sticky=tk.W)
        ttk.Entry(color_frame, textvariable=self.ont_color_var, width=10).pack(side=tk.LEFT, padx=5)
        self.color_preview = tk.Label(color_frame, text="   ", bg="#999999", width=3)
        self.color_preview.pack(side=tk.LEFT)

        # Ontology buttons
        ont_button_frame = ttk.Frame(self.ontology_details_frame)
        ont_button_frame.pack(fill=tk.X, padx=10, pady=10)

        ttk.Button(ont_button_frame, text="Add Category", command=self._add_ontology_category).pack(side=tk.LEFT, padx=5)
        ttk.Button(ont_button_frame, text="Update Category", command=self._update_ontology_category).pack(side=tk.LEFT, padx=5)
        ttk.Button(ont_button_frame, text="Delete Category", command=self._delete_ontology_category).pack(side=tk.LEFT, padx=5)

        self.ontology_paned.add(self.ontology_details_frame, weight=1)

        # Load ontology tree
        self._refresh_ontology_tree()

        # Knowledge Graph tab
        self.graph_tab = ttk.Frame(self.notebook)
        self.notebook.add(self.graph_tab, text="Knowledge Graph")

        # Create matplotlib figure for graph visualization
        self.figure = plt.Figure(figsize=(10, 8), dpi=100)
        self.ax = self.figure.add_subplot(111)
        self.canvas = FigureCanvasTkAgg(self.figure, master=self.graph_tab)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)

        # Graph controls
        self.graph_controls = ttk.Frame(self.graph_tab)
        self.graph_controls.pack(fill=tk.X, padx=10, pady=5)

        ttk.Button(self.graph_controls, text="Refresh Graph",
                  command=self._refresh_knowledge_graph).pack(side=tk.LEFT, padx=5)
        ttk.Button(self.graph_controls, text="Export Graph",
                  command=self._export_graph).pack(side=tk.LEFT, padx=5)

        # Graph filter options
        self.graph_filter_frame = ttk.LabelFrame(self.graph_controls, text="Filter Options", padding=5)
        self.graph_filter_frame.pack(side=tk.LEFT, padx=20)

        self.show_documents_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(self.graph_filter_frame, text="Documents", variable=self.show_documents_var).pack(side=tk.LEFT)
        self.show_companies_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(self.graph_filter_frame, text="Companies", variable=self.show_companies_var).pack(side=tk.LEFT)
        self.show_ontologies_var = tk.BooleanVar(value=False)
        ttk.Checkbutton(self.graph_filter_frame, text="Ontologies", variable=self.show_ontologies_var).pack(side=tk.LEFT)

        self.log_handler = LogHandler(self.log_text)
        logging.getLogger().addHandler(self.log_handler)
        logging.getLogger().setLevel(getattr(logging, self.settings_manager.settings.log_level))

    def _setup_layout(self):
        self.main_container.columnconfigure(0, weight=1)
        self.header.pack(fill=tk.X, pady=(0, 10))

        self.dir_frame.pack(fill=tk.X, pady=5)
        self.dir_label.pack(side=tk.LEFT, padx=5)
        self.dir_combo.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)
        self.browse_button.pack(side=tk.LEFT, padx=5)

        self.metadata_frame.pack(fill=tk.X, pady=5)
        self.file_type_frame.pack(fill=tk.X, pady=10)

        self.progress_container.pack(fill=tk.X, pady=5)
        self.status_frame.pack(fill=tk.X, pady=5)
        self.status_label.pack(side=tk.LEFT, padx=5)

        self.progress_frame.pack(fill=tk.X, pady=5)
        self.stats_frame.pack(fill=tk.X, pady=5)

        self.monitor_frame.pack(fill=tk.X, pady=5)
        self.cpu_label.pack(side=tk.LEFT, padx=5)
        self.memory_label.pack(side=tk.LEFT, padx=5)
        self.memory_process_label.pack(side=tk.LEFT, padx=5)
        self.db_status_label.pack(side=tk.LEFT, padx=10)

        self.log_frame.pack(fill=tk.BOTH, expand=True, pady=5)
        self.log_text.pack(fill=tk.BOTH, expand=True, padx=5)

        self.button_frame.pack(fill=tk.X, pady=5)
        self.process_button.pack(side=tk.LEFT, padx=5)
        self.cancel_button.pack(side=tk.LEFT, padx=5)
        self.settings_button.pack(side=tk.RIGHT, padx=5)

    def _browse_directory(self):
        directory = filedialog.askdirectory()
        if directory:
            self.dir_var.set(directory)
            self.settings_manager.add_recent_directory(directory)
            self.dir_combo['values'] = self.settings_manager.recent_directories

    def _open_settings(self):
        """Open settings dialog"""
        dialog = SettingsDialog(self, self.settings_manager)
        self.wait_window(dialog)

        # Reinitialize components if settings changed
        self.memory_config.memory_limit_mb = self.settings_manager.settings.memory_limit_mb
        self.memory_config.max_batch_size = self.settings_manager.settings.batch_size

        # Update file types
        for file_type, enabled in self.settings_manager.settings.file_types.items():
            if file_type in self.file_types:
                self.file_types[file_type].set(enabled)

        # Reinitialize database if needed
        if self.db_manager:
            self.db_manager.close()
        self.init_database()

    def _start_processing(self):
        directory = self.dir_var.get()
        if not directory:
            messagebox.showerror("Error", "Please select a directory")
            return

        self.processing = True
        self.cancel_requested = False
        self.process_button.config(state=tk.DISABLED)
        self.cancel_button.config(state=tk.NORMAL)
        self.status_label.config(text="Processing...")

        Thread(target=lambda: asyncio.run(self._process_documents(directory))).start()

    def _cancel_processing(self):
        self.cancel_requested = True
        self.status_label.config(text="Cancelling...")

    def _start_monitoring(self):
        def update_monitoring():
            if not self.processing:
                cpu_percent = psutil.cpu_percent()
                memory_percent = psutil.virtual_memory().percent

                self.cpu_label.config(text=f"CPU: {cpu_percent}%")
                self.memory_label.config(text=f"Memory: {memory_percent}%")

            self.after(1000, update_monitoring)

        update_monitoring()

    def _update_memory_display(self, memory_info: Dict[str, Any]):
        """Update memory display from monitor callback"""
        self.memory_process_label.config(
            text=f"Process: {memory_info['process_rss_mb']:.1f}MB"
        )

    def update_progress(self, processed: int = 0, total: int = 0,
                       batch: int = 0, total_batches: int = 0,
                       skipped: int = 0, failed: int = 0,
                       batch_current: int = 0, batch_total: int = 0):
        """Enhanced progress update with detailed statistics"""
        # Update overall progress
        if total > 0:
            percentage = int((processed / total) * 100)
            self.progress_bar["value"] = percentage
            self.progress_label.config(text=f"{percentage}%")

        # Update batch progress
        if batch_total > 0:
            batch_percentage = int((batch_current / batch_total) * 100)
            self.batch_progress_bar["value"] = batch_percentage

        if total_batches > 0:
            self.batch_label.config(text=f"Batch: {batch}/{total_batches}")

        # Update statistics
        self.processed_label.config(text=f"Processed: {processed}")
        self.skipped_label.config(text=f"Skipped: {skipped}")
        self.failed_label.config(text=f"Failed: {failed}")

        self.update_idletasks()

    def _update_time_display(self, start_time):
        """Update elapsed time display"""
        if not self.processing:
            return

        elapsed = datetime.utcnow() - start_time
        hours, remainder = divmod(elapsed.seconds, 3600)
        minutes, seconds = divmod(remainder, 60)
        self.time_label.config(text=f"Time: {hours:02d}:{minutes:02d}:{seconds:02d}")

        # Schedule next update
        self.after(1000, lambda: self._update_time_display(start_time))

    async def _process_documents(self, directory):
        start_time = datetime.utcnow()
        self._update_time_display(start_time)

        try:
            # Get endpoint configuration
            endpoint_config = self.settings_manager.get_endpoint_config()

            # Create enhanced processor with memory management
            processor = EnhancedDocumentProcessor(
                processing_config=ProcessingConfig(
                    MAX_WORKERS=self.settings_manager.settings.max_workers,
                    BATCH_SIZE=self.settings_manager.settings.batch_size,
                    OUTPUT_DIR=Path(self.settings_manager.settings.output_directory)
                ),
                endpoint_config=endpoint_config,
                db_manager=self.db_manager,
                memory_config=self.memory_config,
                settings=self.settings_manager.settings
            )

            # Get selected file types
            selected_file_types = [
                DocumentType[file_type].value
                for file_type, var in self.file_types.items()
                if var.get()
            ]

            # Process documents
            results = await processor.process_directory(
                directory=Path(directory),
                update_progress=self.update_progress,
                selected_file_types=selected_file_types,
                update_status=self.update_status,
                cw_number=self.cw_var.get(),
                company_id=self.company_id_var.get(),
                company_group=self.company_group_var.get(),
                cancel_check=lambda: self.cancel_requested
            )

            if not self.cancel_requested:
                # Get final statistics from database
                if self.db_manager:
                    stats = await self.db_manager.get_processing_progress()
                    summary = (f"Processing completed!\n\n"
                             f"Total documents: {stats['total_documents']}\n"
                             f"Processed today: {stats['today_processed']}\n"
                             f"Success: {stats['today_success']}\n"
                             f"Failed: {stats['today_failed']}\n"
                             f"Average time: {stats['avg_processing_time']:.1f}s")
                else:
                    summary = f"Processing completed!\nProcessed {len(results)} documents."

                self.status_label.config(text="Completed!")
                messagebox.showinfo("Success", summary)
            else:
                self.status_label.config(text="Processing cancelled.")

        except Exception as e:
            self.status_label.config(text="Error occurred")
            logging.error(f"Processing error: {e}")
            messagebox.showerror("Error", f"An error occurred:\n{str(e)}")
        finally:
            self.processing = False
            self.process_button.config(state=tk.NORMAL)
            self.cancel_button.config(state=tk.DISABLED)

    def update_status(self, file_name):
        self.status_label.config(text=f"Processing: {file_name}")
        self.update_idletasks()

    def _view_processed_documents(self):
        """View processed documents from database"""
        if not self.db_manager:
            messagebox.showerror("Error", "Database not connected")
            return

        # Create a new window to display documents
        window = tk.Toplevel(self)
        window.title("Processed Documents")
        window.geometry("800x600")

        # Add implementation to display documents from database
        # This is a placeholder - you would implement a proper view
        label = ttk.Label(window, text="Processed Documents View (To be implemented)")
        label.pack(padx=20, pady=20)

    def _export_to_sql(self):
        """Export current results to SQL"""
        messagebox.showinfo("Export", "SQL export functionality to be implemented")

    def _build_knowledge_graph(self):
        """Build and display knowledge graph with filtering"""
        if not self.db_manager:
            messagebox.showerror("Error", "Database not connected")
            return

        try:
            # Switch to graph tab
            self.notebook.select(self.graph_tab)

            # Build graph
            G = self.db_manager.build_knowledge_graph()

            # Apply filters
            nodes_to_remove = []
            for node, data in G.nodes(data=True):
                node_type = data.get('type')
                if node_type == 'document' and not self.show_documents_var.get():
                    nodes_to_remove.append(node)
                elif node_type == 'company' and not self.show_companies_var.get():
                    nodes_to_remove.append(node)
                elif node_type == 'ontology' and not self.show_ontologies_var.get():
                    nodes_to_remove.append(node)

            # Create filtered graph
            filtered_G = G.copy()
            filtered_G.remove_nodes_from(nodes_to_remove)

            # Clear previous plot
            self.ax.clear()

            if filtered_G.number_of_nodes() == 0:
                self.ax.text(0.5, 0.5, 'No nodes to display',
                           horizontalalignment='center',
                           verticalalignment='center',
                           transform=self.ax.transAxes)
                self.canvas.draw()
                return

            # Draw graph with improved layout
            pos = nx.spring_layout(filtered_G, k=2, iterations=50)

            # Draw nodes by type with different colors and sizes
            for node_type, color, size in [
                ('document', 'lightblue', 300),
                ('company', 'lightgreen', 500),
                ('ontology', 'lightyellow', 400)
            ]:
                nodes = [n for n, d in filtered_G.nodes(data=True) if d.get('type') == node_type]
                if nodes:
                    # Use ontology colors for documents if available
                    if node_type == 'document':
                        node_colors = []
                        for n in nodes:
                            ont_color = filtered_G.nodes[n].get('ontology_color', '#87CEEB')
                            node_colors.append(ont_color)
                        nx.draw_networkx_nodes(filtered_G, pos, nodelist=nodes,
                                             node_color=node_colors, node_size=size, ax=self.ax)
                    else:
                        nx.draw_networkx_nodes(filtered_G, pos, nodelist=nodes,
                                             node_color=color, node_size=size, ax=self.ax)

            # Draw edges with different styles for different relationships
            edge_colors = {
                'amendment': 'red',
                'renewal': 'green',
                'supersedes': 'orange',
                'references': 'blue',
                'vendor': 'purple',
                'client': 'brown',
                'parent_child': 'gray'
            }

            for rel_type, color in edge_colors.items():
                edges = [(u, v) for u, v, d in filtered_G.edges(data=True)
                        if d.get('relationship_type') == rel_type]
                if edges:
                    nx.draw_networkx_edges(filtered_G, pos, edgelist=edges,
                                         edge_color=color, alpha=0.6, ax=self.ax)

            # Draw labels
            labels = {}
            for n, d in filtered_G.nodes(data=True):
                name = d.get('name', str(n))
                # Truncate long names
                if len(name) > 20:
                    name = name[:17] + '...'
                labels[n] = name

            nx.draw_networkx_labels(filtered_G, pos, labels, font_size=8, ax=self.ax)

            # Add legend
            legend_elements = []
            if self.show_documents_var.get():
                legend_elements.append(plt.Line2D([0], [0], marker='o', color='w',
                                                markerfacecolor='lightblue', markersize=10, label='Documents'))
            if self.show_companies_var.get():
                legend_elements.append(plt.Line2D([0], [0], marker='o', color='w',
                                                markerfacecolor='lightgreen', markersize=10, label='Companies'))
            if self.show_ontologies_var.get():
                legend_elements.append(plt.Line2D([0], [0], marker='o', color='w',
                                                markerfacecolor='lightyellow', markersize=10, label='Ontologies'))

            if legend_elements:
                self.ax.legend(handles=legend_elements, loc='upper right')

            self.ax.set_title(f"Document Knowledge Graph ({filtered_G.number_of_nodes()} nodes, {filtered_G.number_of_edges()} edges)")
            self.ax.axis('off')
            self.canvas.draw()

            logging.info(f"Knowledge graph displayed with {filtered_G.number_of_nodes()} nodes and {filtered_G.number_of_edges()} edges")

        except Exception as e:
            logging.error(f"Error building knowledge graph: {e}")
            messagebox.showerror("Error", f"Failed to build knowledge graph: {str(e)}")

    def _refresh_knowledge_graph(self):
        """Refresh the knowledge graph display"""
        self._build_knowledge_graph()

    def _export_graph(self):
        """Export knowledge graph"""
        file_path = filedialog.asksaveasfilename(
            defaultextension=".png",
            filetypes=[("PNG files", "*.png"), ("PDF files", "*.pdf"), ("All files", "*.*")]
        )
        if file_path:
            self.figure.savefig(file_path, dpi=300, bbox_inches='tight')
            messagebox.showinfo("Success", f"Graph exported to {file_path}")

    def _show_memory_report(self):
        """Show detailed memory usage report"""
        memory_info = get_system_memory_info()
        process_info = self.memory_monitor.get_memory_info()

        report = f"""System Memory Report

System Memory:
- Total: {memory_info['total_mb']:.1f} MB
- Available: {memory_info['available_mb']:.1f} MB
- Used: {memory_info['percent']:.1f}%

Process Memory:
- RSS: {process_info['process_rss_mb']:.1f} MB
- VMS: {process_info['process_vms_mb']:.1f} MB
- Process %: {process_info['process_percent']:.1f}%

Settings:
- Memory Limit: {self.memory_config.memory_limit_mb} MB
- Current Batch Size: {self.memory_config.max_batch_size}
"""

        messagebox.showinfo("Memory Report", report)

    def _clear_cache(self):
        """Clear application caches"""
        gc.collect()
        if self.db_manager and self.db_manager._cache:
            self.db_manager._cache.clear()
        messagebox.showinfo("Cache Cleared", "Application caches have been cleared")

    def _show_about(self):
        """Show about dialog"""
        about_text = """Enhanced Document Processor

Version: 2.0
Author: Martin Bacigal
License: MIT License

Features:
- SQL Database Integration
- Memory-Efficient Processing
- Knowledge Graph Generation
- Document Relationship Tracking
- Settings Persistence
        """
        messagebox.showinfo("About", about_text)

    def _on_closing(self):
        """Handle application closing"""
        if self.processing:
            if not messagebox.askokcancel("Quit", "Processing is in progress. Do you want to quit?"):
                return

        # Stop monitoring
        self.memory_monitor.stop_monitoring()

        # Close database
        if self.db_manager:
            self.db_manager.close()

        # Save settings
        self.settings_manager.save_settings()

        self.destroy()


class EnhancedDocumentProcessor(DocumentProcessor):
    """Enhanced document processor with SQL and memory management"""

    def __init__(self, processing_config, endpoint_config, db_manager, memory_config, settings):
        # Initialize parent with endpoint config instead of azure config
        self.config = processing_config
        self.endpoint_config = endpoint_config
        self.analyzer = ContentAnalyzer(endpoint_config)
        self.executor = ThreadPoolExecutor(max_workers=processing_config.MAX_WORKERS)

        # Define processors for different MIME types
        self.processors = {
            DocumentType.PDF.value: self._extract_pdf_content,
            DocumentType.DOCX.value: self._extract_word_content,
            DocumentType.DOC.value: self._extract_word_content,
            DocumentType.XLSX.value: self._extract_excel_content,
            DocumentType.XLS.value: self._extract_excel_content
        }

        self.db_manager = db_manager
        self.memory_config = memory_config
        self.settings = settings
        self.memory_processor = MemoryEfficientDocumentProcessor(
            memory_config,
            lambda file_path: asyncio.run(self.process_document(file_path))
        )

    async def process_directory(self, directory: Path, update_progress, selected_file_types: List[str],
                              update_status, cw_number: str = None, company_id: str = None,
                              company_group: str = None, cancel_check=None) -> List[Dict[str, Any]]:
        """Process directory with SQL integration and memory management"""
        files = [f for f in directory.rglob("*") if f.is_file() and
                 mimetypes.guess_type(str(f))[0] in selected_file_types]

        results = []
        total_files = len(files)
        processed = 0
        skipped = 0
        failed = 0

        logging.info(f"Found {total_files} files to process")

        # SQL-focused processing
        if self.db_manager:
            # Pre-calculate all file hashes and check which need processing
            files_to_process = []

            for file_path in files:
                if cancel_check and cancel_check():
                    break

                # Calculate hash BEFORE processing
                file_hash = self.db_manager.calculate_file_hash(file_path)

                # Check if already processed
                existing = await self.db_manager.check_document_exists(file_hash)

                if existing['exists'] and not existing['needs_reprocessing'] and self.settings.skip_processed_files:
                    skipped += 1
                    logging.info(f"Skipping already processed file: {file_path.name} (hash: {file_hash[:8]}...)")
                    update_progress(
                        processed=processed + skipped,
                        total=total_files,
                        skipped=skipped,
                        failed=failed
                    )
                else:
                    files_to_process.append((file_path, file_hash, existing))

            logging.info(f"Files to process: {len(files_to_process)}, Skipped: {skipped}")

            # Process in batches for memory efficiency
            batch_size = self.settings.batch_size
            total_batches = (len(files_to_process) + batch_size - 1) // batch_size

            for batch_idx, i in enumerate(range(0, len(files_to_process), batch_size)):
                if cancel_check and cancel_check():
                    break

                batch = files_to_process[i:i + batch_size]
                batch_results = []

                # Update batch progress
                update_progress(
                    processed=processed,
                    total=total_files,
                    batch=batch_idx + 1,
                    total_batches=total_batches,
                    skipped=skipped,
                    failed=failed,
                    batch_current=0,
                    batch_total=len(batch)
                )

                for idx, (file_path, file_hash, existing_info) in enumerate(batch):
                    if cancel_check and cancel_check():
                        break

                    update_status(f"Processing: {file_path.name}")

                    try:
                        start_time = datetime.utcnow()
                        memory_start = psutil.Process().memory_info().rss / 1024 / 1024

                        # Create or update document record
                        document, is_new = await self.db_manager.create_or_update_document(
                            file_path=file_path,
                            file_hash=file_hash,
                            cw_number=f"{cw_number}-{file_path.stem}" if cw_number else None,
                            created_by="system"
                        )

                        # Update processing status
                        await self.db_manager.update_document_processing(
                            document.id, 'processing'
                        )

                        # Process document content
                        result = await self.process_document(file_path)

                        # Calculate processing metrics
                        duration = int((datetime.utcnow() - start_time).total_seconds())
                        memory_used = int(psutil.Process().memory_info().rss / 1024 / 1024 - memory_start)

                        if 'error' not in result:
                            # Update with results
                            await self.db_manager.update_document_processing(
                                document.id,
                                'completed',
                                analysis_results=result,
                                duration_seconds=duration,
                                memory_used_mb=memory_used
                            )

                            # Link to company if provided
                            if company_id:
                                await self.db_manager.link_document_to_company(
                                    document.id, company_id, role='primary', is_primary=True
                                )

                            # Auto-assign to ontology based on content
                            if self.settings.auto_detect_relationships:
                                await self._auto_assign_ontology(document.id, result)

                            batch_results.append(result)
                            processed += 1
                        else:
                            # Mark as failed
                            await self.db_manager.update_document_processing(
                                document.id,
                                'failed',
                                message=result.get('error', 'Unknown error'),
                                duration_seconds=duration,
                                memory_used_mb=memory_used
                            )
                            failed += 1

                    except Exception as e:
                        logging.error(f"Error processing {file_path}: {e}")
                        failed += 1

                    # Update progress for each file in batch
                    update_progress(
                        processed=processed,
                        total=total_files,
                        batch=batch_idx + 1,
                        total_batches=total_batches,
                        skipped=skipped,
                        failed=failed,
                        batch_current=idx + 1,
                        batch_total=len(batch)
                    )

                # Save batch to Excel if enabled
                if self.settings.save_to_excel and batch_results:
                    output_path = self.config.OUTPUT_DIR / "processed_results.xlsx"
                    output_path.parent.mkdir(parents=True, exist_ok=True)
                    await ResultsExporter.to_excel(batch_results, output_path, self.analyzer, update_progress)

                results.extend(batch_results)

                # Force garbage collection between batches
                gc.collect()

                # Log batch completion
                logging.info(f"Completed batch {batch_idx + 1}/{total_batches}")

        else:
            # Fallback to non-SQL processing
            logging.warning("No database connection - processing without SQL storage")
            results = await super().process_directory(
                directory, update_progress, selected_file_types, update_status, batch_size
            )

        return results

    async def _auto_assign_ontology(self, document_id: int, result: Dict[str, Any]):
        """Automatically assign document to ontology based on content analysis"""
        try:
            # Extract key terms from analysis
            analyses = result.get('analyses', [])
            content_terms = []

            for analysis in analyses:
                if analysis.get('analysis_type') == 'contract_details':
                    # Extract terms from contract details
                    result_text = str(analysis.get('result', ''))
                    content_terms.extend(result_text.lower().split())

            # Simple keyword-based assignment (can be enhanced with ML)
            if any(term in content_terms for term in ['service', 'services', 'consulting']):
                await self.db_manager.assign_document_to_ontology(
                    document_id, 'PROCUREMENT.SERVICES', confidence=0.8,
                    is_primary=True, assigned_by='ai'
                )
            elif any(term in content_terms for term in ['supply', 'goods', 'equipment']):
                await self.db_manager.assign_document_to_ontology(
                    document_id, 'PROCUREMENT.GOODS', confidence=0.8,
                    is_primary=True, assigned_by='ai'
                )
            elif any(term in content_terms for term in ['nda', 'confidentiality', 'non-disclosure']):
                await self.db_manager.assign_document_to_ontology(
                    document_id, 'LEGAL.NDA', confidence=0.9,
                    is_primary=True, assigned_by='ai'
                )
            elif any(term in content_terms for term in ['license', 'licensing', 'software']):
                await self.db_manager.assign_document_to_ontology(
                    document_id, 'LEGAL.IP', confidence=0.8,
                    is_primary=True, assigned_by='ai'
                )
            else:
                # Default to operational if no specific match
                await self.db_manager.assign_document_to_ontology(
                    document_id, 'OPERATIONAL', confidence=0.5,
                    is_primary=True, assigned_by='ai'
                )

        except Exception as e:
            logging.error(f"Error in auto-assigning ontology: {e}")

    async def _detect_relationships(self, document_id: int, result: Dict[str, Any]):
        """Auto-detect document relationships based on content"""
        try:
            # Extract contract numbers and references from content
            analyses = result.get('analyses', [])

            for analysis in analyses:
                result_text = str(analysis.get('result', ''))

                # Look for amendment references
                if 'amendment' in result_text.lower():
                    # Try to find parent contract reference
                    # This is simplified - real implementation would use NLP
                    pass

                # Look for renewal references
                if 'renewal' in result_text.lower() or 'extension' in result_text.lower():
                    # Try to find original contract
                    pass

        except Exception as e:
            logging.error(f"Error detecting relationships: {e}")

    def _refresh_ontology_tree(self):
        """Refresh the ontology tree view"""
        if not self.db_manager:
            return

        # Clear existing tree
        for item in self.ontology_tree.get_children():
            self.ontology_tree.delete(item)

        try:
            # Get ontology tree from database
            tree_data = asyncio.run(self.db_manager.get_ontology_tree())

            # Build tree
            def add_nodes(parent_id, nodes):
                for node in nodes:
                    item_id = self.ontology_tree.insert(
                        parent_id, 'end',
                        text=node['name'],
                        values=(node['code'], node['description'], node['color']),
                        tags=(node['id'],)
                    )
                    if node['children']:
                        add_nodes(item_id, node['children'])

            add_nodes('', tree_data)

            # Bind selection event
            self.ontology_tree.bind('<<TreeviewSelect>>', self._on_ontology_select)

        except Exception as e:
            logging.error(f"Error loading ontology tree: {e}")

    def _on_ontology_select(self, event):
        """Handle ontology tree selection"""
        selection = self.ontology_tree.selection()
        if selection:
            item = self.ontology_tree.item(selection[0])
            values = item['values']

            self.ont_code_var.set(values[0] if len(values) > 0 else '')
            self.ont_name_var.set(item['text'])
            self.ont_desc_text.delete(1.0, tk.END)
            self.ont_desc_text.insert(1.0, values[1] if len(values) > 1 else '')
            self.ont_color_var.set(values[2] if len(values) > 2 else '#999999')

            # Update color preview
            try:
                self.color_preview.configure(bg=self.ont_color_var.get())
            except:
                pass

    def _add_ontology_category(self):
        """Add new ontology category"""
        # Implementation for adding ontology category
        messagebox.showinfo("Add Category", "Add ontology category functionality to be implemented")

    def _update_ontology_category(self):
        """Update selected ontology category"""
        # Implementation for updating ontology category
        messagebox.showinfo("Update Category", "Update ontology category functionality to be implemented")

    def _delete_ontology_category(self):
        """Delete selected ontology category"""
        # Implementation for deleting ontology category
        messagebox.showinfo("Delete Category", "Delete ontology category functionality to be implemented")

class ResultsExporter:
    @staticmethod
    async def to_excel(batch_results: List[Dict[str, Any]], output_path: Path, analyzer: ContentAnalyzer, update_progress):
        try:
            flattened_results = []
            total_tasks = len(batch_results)
            progress = 0

            def progress_callback():
                nonlocal progress
                progress += 1
                update_progress(progress, total_tasks)

            formatted_infos = []
            for result in batch_results:
                for analysis in result.get('analyses', []):
                    logging.debug(f"Processing analysis: {analysis}")
                    analysis_type = analysis.get('analysis_type')
                    result_content = analysis.get('result')
                    if analysis_type is None:
                        logging.error(f"Missing analysis_type in analysis: {analysis}")
                        continue
                    formatted_info = analyzer.format_information(result_content, analysis_type)
                    formatted_infos.append((formatted_info, result, analysis))

            for formatted_info, result, analysis in formatted_infos:
                analysis_type = analysis.get('analysis_type')
                logging.debug(f"Flattening result for analysis_type: {analysis_type}")
                result_content = analysis.get('result')
                flattened_result = {
                    'file_path': result.get('file_path'),
                    'metadata': str(result.get('metadata')),  # Convert metadata to string
                    'analysis_type': analysis_type,
                    'result': str(result_content),  # Ensure result content is a string
                    'start_date': formatted_info.get('start_date'),
                    'end_date': formatted_info.get('end_date'),
                    'contract_duration': formatted_info.get('contract_duration'),
                    'key_deliverables': str(formatted_info.get('key_deliverables')),  # Convert list to string
                    'vendor_capabilities': str(formatted_info.get('vendor_capabilities')),
                    'risk_factors': str(formatted_info.get('risk_factors')),
                    'compliance_status': str(formatted_info.get('compliance_status')),
                    'past_performance_metrics': str(formatted_info.get('past_performance_metrics')),
                    'key_strengths_weaknesses': str(formatted_info.get('key_strengths_weaknesses')),
                    'technical_requirements': str(formatted_info.get('technical_requirements')),
                    'performance_metrics': str(formatted_info.get('performance_metrics')),
                    'compatibility_requirements': str(formatted_info.get('compatibility_requirements')),
                    'implementation_timeline': str(formatted_info.get('implementation_timeline')),
                    'maintenance_requirements': str(formatted_info.get('maintenance_requirements'))
                }
                flattened_results.append(flattened_result)
                progress_callback()

            new_data = pd.DataFrame(flattened_results)
            logging.debug(f"Flattened results: {flattened_results}")
            logging.debug(f"New data DataFrame: {new_data}")

            if output_path.exists():
                wb = load_workbook(output_path)
                with pd.ExcelWriter(output_path, engine='openpyxl', mode='a') as writer:
                    for sheet_name, data in new_data.groupby('analysis_type'):
                        logging.debug(f"Appending data to sheet: {sheet_name}")
                        if sheet_name in wb.sheetnames:
                            ws = wb[sheet_name]
                            for r in dataframe_to_rows(data, index=False, header=False):
                                ws.append(r)
                        else:
                            data.to_excel(writer, sheet_name=sheet_name, index=False)
                wb.save(output_path)
            else:
                with pd.ExcelWriter(output_path, engine='openpyxl', mode='w') as writer:
                    for sheet_name, data in new_data.groupby('analysis_type'):
                        data.to_excel(writer, sheet_name=sheet_name, index=False)

            logging.info(f"Results exported to {output_path}")
        except Exception as e:
            logging.error(f"Error in to_excel: {e}")
            raise

def main():
    # Enhanced Logging Setup
    logging.basicConfig(level=logging.DEBUG)
    logger = logging.getLogger()
    text_widget = None  # Replace with your actual text widget
    if text_widget:
        handler = LogHandler(text_widget)
        logger.addHandler(handler)

    # Ensure output directory exists
    output_dir = ProcessingConfig.OUTPUT_DIR
    output_dir.mkdir(exist_ok=True)

    # Start the application
    app = EnhancedDocumentProcessorApp()

    # Configure window styling
    app.option_add('*TButton*Padding', 5)
    app.option_add('*TButton*Font', 'Helvetica 10')
    app.option_add('*TLabel*Font', 'Helvetica 10')

    # Center the window on screen
    window_width = 800
    window_height = 600
    screen_width = app.winfo_screenwidth()
    screen_height = app.winfo_screenheight()
    center_x = int(screen_width/2 - window_width/2)
    center_y = int(screen_height/2 - window_height/2)
    app.geometry(f'{window_width}x{window_height}+{center_x}+{center_y}')
    app.mainloop()

if __name__ == "__main__":
    main()