database_manager.py

"""
Database Manager for Contract Processing System
Author: Martin Bacigal, 01/2025 @ https://procureai.tech
License: MIT License
"""

import asyncio
import logging
from datetime import datetime, date
from typing import Optional, Dict, List, Any, Union, Tuple
from pathlib import Path
import hashlib
import json
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
import uuid
from decimal import Decimal

from sqlalchemy import (
    create_engine, Column, Integer, String, DateTime, Boolean,
    Text, JSON, ForeignKey, Table, Float, Index, UniqueConstraint,
    Date, DECIMAL, select, text, func, and_, or_
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship, Session
from sqlalchemy.pool import QueuePool
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.dialects.postgresql import UUID as PostgresUUID
from sqlalchemy.types import TypeDecorator, CHAR
from redis import Redis
from cachetools import TTLCache
import networkx as nx
from py2neo import Graph, Node, Relationship as Neo4jRelationship

Base = declarative_base()

# Custom UUID type that works across databases
class UUID(TypeDecorator):
    """Platform-independent GUID type."""
    impl = CHAR
    cache_ok = True

    def load_dialect_impl(self, dialect):
        if dialect.name == 'postgresql':
            return dialect.type_descriptor(PostgresUUID())
        else:
            return dialect.type_descriptor(CHAR(36))

    def process_bind_param(self, value, dialect):
        if value is None:
            return value
        elif dialect.name == 'postgresql':
            return str(value)
        else:
            if not isinstance(value, uuid.UUID):
                return str(uuid.UUID(value))
            else:
                return str(value)

    def process_result_value(self, value, dialect):
        if value is None:
            return value
        else:
            if not isinstance(value, uuid.UUID):
                value = uuid.UUID(value)
            return value

# Association tables for many-to-many relationships
document_relationships = Table(
    'document_relationships',
    Base.metadata,
    Column('id', Integer, primary_key=True),
    Column('parent_document_id', Integer, ForeignKey('documents.id', ondelete='CASCADE')),
    Column('child_document_id', Integer, ForeignKey('documents.id', ondelete='CASCADE')),
    Column('relationship_type', String(50), nullable=False),
    Column('confidence', DECIMAL(3, 2), default=Decimal('1.00')),
    Column('detected_by', String(50)),
    Column('metadata', JSON),
    Column('created_at', DateTime, default=datetime.utcnow)
)

document_companies = Table(
    'document_companies',
    Base.metadata,
    Column('id', Integer, primary_key=True),
    Column('document_id', Integer, ForeignKey('documents.id', ondelete='CASCADE')),
    Column('company_id', Integer, ForeignKey('companies.id')),
    Column('role', String(50), nullable=False),
    Column('is_primary', Boolean, default=False),
    Column('confidence_score', DECIMAL(3, 2), default=Decimal('1.00')),
    Column('detected_from', String(100)),
    Column('created_at', DateTime, default=datetime.utcnow)
)

@dataclass
class DatabaseConfig:
    """Database configuration"""
    db_type: str = "postgresql"
    host: str = "localhost"
    port: int = 5432
    database: str = "contract_processor"
    username: str = "postgres"
    password: str = "password"
    pool_size: int = 20
    max_overflow: int = 40
    pool_pre_ping: bool = True
    pool_recycle: int = 1800  # seconds to recycle stale connections
    echo: bool = False

    # Redis config for caching
    redis_host: str = "localhost"
    redis_port: int = 6379
    redis_db: int = 0

    # Neo4j config for knowledge graph
    neo4j_uri: str = "bolt://localhost:7687"
    neo4j_username: str = "neo4j"
    neo4j_password: str = "password"

    def get_connection_string(self) -> str:
        """Get database connection string"""
        if self.db_type == "sqlite":
            return f"sqlite:///{self.database}.db"
        elif self.db_type == "postgresql":
            return f"postgresql://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
        elif self.db_type == "mysql":
            return f"mysql+pymysql://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
        else:
            raise ValueError(f"Unsupported database type: {self.db_type}")

    def get_async_connection_string(self) -> str:
        """Get async database connection string"""
        if self.db_type == "postgresql":
            return f"postgresql+asyncpg://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
        else:
            return self.get_connection_string()

class Company(Base):
    """Company entity model"""
    __tablename__ = 'companies'

    id = Column(Integer, primary_key=True)
    company_id = Column(String(50), unique=True, nullable=False, index=True)
    name = Column(String(255), nullable=False)
    company_group = Column(String(100), index=True)
    industry = Column(String(100))
    country = Column(String(100))
    parent_company_id = Column(Integer, ForeignKey('companies.id'))
    metadata = Column(JSON)
    active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationships
    documents = relationship("Document", secondary=document_companies, back_populates="companies")
    parent_company = relationship("Company", remote_side=[id], backref="subsidiaries")

class ContractType(Base):
    """Contract types enumeration"""
    __tablename__ = 'contract_types'

    id = Column(Integer, primary_key=True)
    type_code = Column(String(50), unique=True, nullable=False)
    type_name = Column(String(100), nullable=False)
    description = Column(Text)
    typical_duration_days = Column(Integer)
    metadata = Column(JSON)
    created_at = Column(DateTime, default=datetime.utcnow)

    # Relationships
    documents = relationship("Document", back_populates="contract_type")

class Document(Base):
    """Document entity model with enhanced fields"""
    __tablename__ = 'documents'

    id = Column(Integer, primary_key=True)
    document_id = Column(UUID, default=uuid.uuid4, unique=True, nullable=False)

    # File information
    file_path = Column(String(500), nullable=False)
    file_name = Column(String(255), nullable=False)
    file_type = Column(String(50))
    file_size_bytes = Column(Integer)

    # Contract identifiers
    cw_number = Column(String(50), index=True)
    contract_number = Column(String(100), index=True)
    contract_type_id = Column(Integer, ForeignKey('contract_types.id'))

    # Processing status
    processed = Column(Boolean, default=False, index=True)
    processing_started_at = Column(DateTime)
    processing_completed_at = Column(DateTime)
    processing_status = Column(String(50), default='pending')
    processing_version = Column(Integer, default=1)
    error_message = Column(Text)

    # Current hash (latest version)
    current_file_hash = Column(String(64), nullable=False, index=True)

    # Contract details (extracted)
    contract_title = Column(String(500))
    contract_start_date = Column(Date)
    contract_end_date = Column(Date)
    contract_duration = Column(String(100))
    contract_value = Column(DECIMAL(15, 2))
    currency = Column(String(3))

    # AI Analysis results
    analysis_results = Column(JSON)
    key_deliverables = Column(JSON)
    risk_score = Column(DECIMAL(3, 2))
    compliance_score = Column(DECIMAL(3, 2))

    # Metadata
    tags = Column(JSON)
    custom_fields = Column(JSON)
    created_at = Column(DateTime, default=datetime.utcnow, index=True)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    created_by = Column(String(100))

    # Relationships
    contract_type = relationship("ContractType", back_populates="documents")
    companies = relationship("Company", secondary=document_companies, back_populates="documents")
    parent_documents = relationship(
        "Document",
        secondary=document_relationships,
        primaryjoin=id == document_relationships.c.child_document_id,
        secondaryjoin=id == document_relationships.c.parent_document_id,
        backref="child_documents"
    )
    processing_logs = relationship("ProcessingLog", back_populates="document", cascade="all, delete-orphan")
    file_hashes = relationship("FileHash", back_populates="document", cascade="all, delete-orphan")
    versions = relationship("DocumentVersion", back_populates="document", cascade="all, delete-orphan")
    ontology_mappings = relationship("DocumentOntologyMapping", back_populates="document", cascade="all, delete-orphan")

    __table_args__ = (
        Index('idx_processed_status', 'processed', 'processing_status'),
        Index('idx_contract_dates', 'contract_start_date', 'contract_end_date'),
    )

class FileHash(Base):
    """File hash tracking for version control"""
    __tablename__ = 'file_hashes'

    id = Column(Integer, primary_key=True)
    document_id = Column(Integer, ForeignKey('documents.id', ondelete='CASCADE'), nullable=False)
    file_hash = Column(String(64), nullable=False, index=True)
    hash_algorithm = Column(String(20), default='SHA256')
    file_size_bytes = Column(Integer)
    calculated_at = Column(DateTime, default=datetime.utcnow)
    is_current = Column(Boolean, default=True)

    # Relationships
    document = relationship("Document", back_populates="file_hashes")

    __table_args__ = (
        UniqueConstraint('document_id', 'file_hash'),
        Index('idx_document_current', 'document_id', 'is_current'),
    )

class DocumentVersion(Base):
    """Document versions tracking"""
    __tablename__ = 'document_versions'

    id = Column(Integer, primary_key=True)
    document_id = Column(Integer, ForeignKey('documents.id', ondelete='CASCADE'), nullable=False)
    version_number = Column(Integer, nullable=False)
    file_hash = Column(String(64), nullable=False)

    # What changed
    change_type = Column(String(50))  # content_update, metadata_update, reprocessed, manual_edit
    change_description = Column(Text)

    # Processing info for this version
    processed = Column(Boolean, default=False)
    processing_completed_at = Column(DateTime)
    analysis_results = Column(JSON)

    # Version metadata
    created_at = Column(DateTime, default=datetime.utcnow)
    created_by = Column(String(100))

    # Relationships
    document = relationship("Document", back_populates="versions")

    __table_args__ = (
        UniqueConstraint('document_id', 'version_number'),
        Index('idx_version_hash', 'file_hash'),
    )

class ContractOntology(Base):
    """Contract ontology categories for classification"""
    __tablename__ = 'contract_ontology'

    id = Column(Integer, primary_key=True)
    category_code = Column(String(50), unique=True, nullable=False)
    category_name = Column(String(100), nullable=False)
    parent_category_id = Column(Integer, ForeignKey('contract_ontology.id'))
    level = Column(Integer, nullable=False, default=0)
    description = Column(Text)
    keywords = Column(JSON)
    rules = Column(JSON)  # Rules for automatic classification
    color_hex = Column(String(7))  # For visualization
    icon = Column(String(50))
    sort_order = Column(Integer, default=0)
    active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationships
    parent_category = relationship("ContractOntology", remote_side=[id], backref="subcategories")
    document_mappings = relationship("DocumentOntologyMapping", back_populates="ontology")

    __table_args__ = (
        Index('idx_parent_category', 'parent_category_id'),
        Index('idx_level', 'level'),
    )

class DocumentOntologyMapping(Base):
    """Mapping documents to ontology categories"""
    __tablename__ = 'document_ontology_mapping'

    id = Column(Integer, primary_key=True)
    document_id = Column(Integer, ForeignKey('documents.id', ondelete='CASCADE'), nullable=False)
    ontology_id = Column(Integer, ForeignKey('contract_ontology.id'), nullable=False)
    confidence_score = Column(DECIMAL(3, 2), default=Decimal('1.00'))
    is_primary = Column(Boolean, default=False)
    assigned_by = Column(String(50))  # 'ai', 'user', 'rule'
    assigned_at = Column(DateTime, default=datetime.utcnow)

    # Relationships
    document = relationship("Document", back_populates="ontology_mappings")
    ontology = relationship("ContractOntology", back_populates="document_mappings")

    __table_args__ = (
        UniqueConstraint('document_id', 'ontology_id'),
        Index('idx_document_ontology', 'document_id', 'is_primary'),
    )

class ProcessingLog(Base):
    """Processing log for audit trail"""
    __tablename__ = 'processing_logs'

    id = Column(Integer, primary_key=True)
    document_id = Column(Integer, ForeignKey('documents.id', ondelete='CASCADE'))
    action = Column(String(50), nullable=False)
    status = Column(String(50))
    message = Column(Text)
    details = Column(JSON)
    duration_seconds = Column(Integer)
    memory_used_mb = Column(Integer)
    created_at = Column(DateTime, default=datetime.utcnow)
    created_by = Column(String(100))

    # Relationships
    document = relationship("Document", back_populates="processing_logs")

    __table_args__ = (
        Index('idx_document_action', 'document_id', 'action'),
        Index('idx_created_at', 'created_at'),
        Index('idx_action', 'action'),
    )

class ProcessingStatistics(Base):
    """Daily processing statistics for dashboards"""
    __tablename__ = 'processing_statistics'

    id = Column(Integer, primary_key=True)
    date = Column(Date, nullable=False, unique=True, index=True)
    total_documents_processed = Column(Integer, default=0)
    total_processing_time_seconds = Column(Integer, default=0)
    average_processing_time_seconds = Column(DECIMAL(10, 2))
    success_count = Column(Integer, default=0)
    failure_count = Column(Integer, default=0)
    total_file_size_mb = Column(DECIMAL(15, 2))
    unique_companies = Column(Integer, default=0)
    new_relationships_found = Column(Integer, default=0)
    created_at = Column(DateTime, default=datetime.utcnow)

class DatabaseManager:
    """Enhanced database manager with version tracking and ontology support"""

    def __init__(self, config: DatabaseConfig):
        self.config = config
        self._engine = None
        self._async_engine = None
        self._session_factory = None
        self._async_session_factory = None
        self._redis_client = None
        self._neo4j_graph = None
        self._cache = TTLCache(maxsize=1000, ttl=300)

    def initialize(self):
        """Initialize database connections with resilient pooling."""
        # Create synchronous engine
        self._engine = create_engine(
            self.config.get_connection_string(),
            poolclass=QueuePool,
            pool_size=self.config.pool_size,
            max_overflow=self.config.max_overflow,
            pool_pre_ping=self.config.pool_pre_ping,
            pool_recycle=self.config.pool_recycle,
            echo=self.config.echo
        )
        self._session_factory = sessionmaker(bind=self._engine)

        # Create asynchronous engine
        self._async_engine = create_async_engine(
            self.config.get_async_connection_string(),
            pool_size=self.config.pool_size,
            max_overflow=self.config.max_overflow,
            pool_pre_ping=self.config.pool_pre_ping,
            pool_recycle=self.config.pool_recycle,
            echo=self.config.echo
        )
        self._async_session_factory = async_sessionmaker(bind=self._async_engine)

        # Initialize Redis
        try:
            self._redis_client = Redis(
                host=self.config.redis_host,
                port=self.config.redis_port,
                db=self.config.redis_db,
                decode_responses=True
            )
            self._redis_client.ping()
            logging.info("Redis connection established")
        except Exception as e:
            logging.warning(f"Redis connection failed: {e}. Caching will use in-memory cache only.")
            self._redis_client = None

        # Initialize Neo4j
        try:
            self._neo4j_graph = Graph(
                self.config.neo4j_uri,
                auth=(self.config.neo4j_username, self.config.neo4j_password)
            )
            logging.info("Neo4j connection established")
        except Exception as e:
            logging.warning(f"Neo4j connection failed: {e}. Knowledge graph features will be limited.")
            self._neo4j_graph = None

        # Create tables
        Base.metadata.create_all(self._engine)

        # Initialize default data
        self._initialize_default_data()

        logging.info("Database initialized successfully")

    def _initialize_default_data(self):
        """Initialize default contract types and ontologies"""
        with self.get_session() as session:
            # Default contract types
            default_types = [
                ('SERVICE', 'Service Agreement', 'General service contracts'),
                ('SUPPLY', 'Supply Agreement', 'Product supply contracts'),
                ('NDA', 'Non-Disclosure Agreement', 'Confidentiality agreements'),
                ('MSA', 'Master Service Agreement', 'Framework agreements'),
                ('SOW', 'Statement of Work', 'Project-specific work agreements'),
                ('LICENSE', 'License Agreement', 'Software/IP licensing'),
                ('PURCHASE', 'Purchase Agreement', 'One-time purchase contracts')
            ]

            for type_code, type_name, description in default_types:
                if not session.query(ContractType).filter_by(type_code=type_code).first():
                    contract_type = ContractType(
                        type_code=type_code,
                        type_name=type_name,
                        description=description
                    )
                    session.add(contract_type)

            # Default ontologies
            default_ontologies = [
                ('ROOT', 'All Contracts', 0, None, 'Root category', '#000000'),
                ('PROCUREMENT', 'Procurement', 1, 'ROOT', 'Procurement related contracts', '#1E90FF'),
                ('PROCUREMENT.GOODS', 'Goods Procurement', 2, 'PROCUREMENT', 'Purchase of physical goods', '#4169E1'),
                ('PROCUREMENT.SERVICES', 'Services Procurement', 2, 'PROCUREMENT', 'Purchase of services', '#0000CD'),
                ('LEGAL', 'Legal', 1, 'ROOT', 'Legal agreements', '#FF6347'),
                ('LEGAL.NDA', 'Confidentiality', 2, 'LEGAL', 'Non-disclosure and confidentiality', '#DC143C'),
                ('LEGAL.IP', 'Intellectual Property', 2, 'LEGAL', 'IP and licensing agreements', '#8B0000'),
                ('OPERATIONAL', 'Operational', 1, 'ROOT', 'Day-to-day operational contracts', '#32CD32'),
                ('OPERATIONAL.FACILITIES', 'Facilities', 2, 'OPERATIONAL', 'Facility management contracts', '#228B22'),
                ('OPERATIONAL.IT', 'IT Services', 2, 'OPERATIONAL', 'Information technology contracts', '#006400')
            ]

            # Create ontology map for parent lookups
            ontology_map = {}

            for code, name, level, parent_code, desc, color in default_ontologies:
                if not session.query(ContractOntology).filter_by(category_code=code).first():
                    parent_id = None
                    if parent_code and parent_code in ontology_map:
                        parent_id = ontology_map[parent_code]

                    ontology = ContractOntology(
                        category_code=code,
                        category_name=name,
                        level=level,
                        parent_category_id=parent_id,
                        description=desc,
                        color_hex=color
                    )
                    session.add(ontology)
                    session.flush()
                    ontology_map[code] = ontology.id

            session.commit()

    def get_session(self) -> Session:
        """Get a database session"""
        return self._session_factory()

    @asynccontextmanager
    async def get_async_session(self) -> AsyncSession:
        """Get an async database session"""
        async with self._async_session_factory() as session:
            yield session

    def calculate_file_hash(self, file_path: Path) -> str:
        """Calculate SHA-256 hash of a file"""
        sha256_hash = hashlib.sha256()
        with open(file_path, "rb") as f:
            for byte_block in iter(lambda: f.read(4096), b""):
                sha256_hash.update(byte_block)
        return sha256_hash.hexdigest()

    async def check_document_exists(self, file_hash: str) -> Optional[Dict[str, Any]]:
        """Check if document with this hash exists and get its info"""
        cache_key = f"doc_hash:{file_hash}"

        # Check memory cache
        if cache_key in self._cache:
            return self._cache[cache_key]

        # Check Redis cache
        if self._redis_client:
            cached = self._redis_client.get(cache_key)
            if cached:
                result = json.loads(cached)
                self._cache[cache_key] = result
                return result

        # Check database
        async with self.get_async_session() as session:
            result = await session.execute(
                select(Document, FileHash).join(
                    FileHash,
                    and_(
                        Document.id == FileHash.document_id,
                        FileHash.is_current == True
                    )
                ).where(FileHash.file_hash == file_hash)
            )
            row = result.first()

            if row:
                document, file_hash_obj = row
                doc_data = {
                    "exists": True,
                    "document_id": str(document.document_id),
                    "processing_status": document.processing_status,
                    "processed": document.processed,
                    "version": document.processing_version,
                    "file_path": document.file_path,
                    "needs_reprocessing": False
                }

                # Check if file has changed
                current_path = Path(document.file_path)
                if current_path.exists():
                    current_hash = self.calculate_file_hash(current_path)
                    if current_hash != file_hash:
                        doc_data["needs_reprocessing"] = True

                # Cache the result
                self._cache[cache_key] = doc_data
                if self._redis_client:
                    self._redis_client.setex(cache_key, 300, json.dumps(doc_data))

                return doc_data

        return {"exists": False}

    async def create_or_update_document(
        self,
        file_path: Path,
        file_hash: str,
        cw_number: Optional[str] = None,
        created_by: Optional[str] = None
    ) -> Tuple[Document, bool]:
        """Create new document or update existing one if file changed"""
        async with self.get_async_session() as session:
            # Check if document exists
            result = await session.execute(
                select(Document).where(Document.file_path == str(file_path))
            )
            document = result.scalar_one_or_none()

            is_new = document is None
            start_time = datetime.utcnow()

            if is_new:
                # Create new document
                document = Document(
                    document_id=uuid.uuid4(),
                    file_path=str(file_path),
                    file_name=file_path.name,
                    file_type=file_path.suffix,
                    file_size_bytes=file_path.stat().st_size,
                    current_file_hash=file_hash,
                    cw_number=cw_number,
                    processing_status='pending',
                    created_by=created_by
                )
                session.add(document)
                await session.flush()

                # Add initial hash record
                file_hash_record = FileHash(
                    document_id=document.id,
                    file_hash=file_hash,
                    file_size_bytes=document.file_size_bytes,
                    is_current=True
                )
                session.add(file_hash_record)

                # Add initial version
                version = DocumentVersion(
                    document_id=document.id,
                    version_number=1,
                    file_hash=file_hash,
                    change_type='initial',
                    change_description='Initial document upload',
                    created_by=created_by
                )
                session.add(version)

                # Log creation
                log_entry = ProcessingLog(
                    document_id=document.id,
                    action='created',
                    status='success',
                    message=f'Document created: {file_path.name}',
                    created_by=created_by,
                    duration_seconds=int((datetime.utcnow() - start_time).total_seconds())
                )
                session.add(log_entry)

            else:
                # Check if file hash changed
                if document.current_file_hash != file_hash:
                    # Update version
                    new_version = document.processing_version + 1

                    # Update document
                    document.current_file_hash = file_hash
                    document.processing_version = new_version
                    document.processing_status = 'reprocessing'
                    document.processed = False
                    document.file_size_bytes = file_path.stat().st_size

                    # Mark old hash as not current
                    await session.execute(
                        text("UPDATE file_hashes SET is_current = false WHERE document_id = :doc_id"),
                        {"doc_id": document.id}
                    )

                    # Add new hash record
                    file_hash_record = FileHash(
                        document_id=document.id,
                        file_hash=file_hash,
                        file_size_bytes=document.file_size_bytes,
                        is_current=True
                    )
                    session.add(file_hash_record)

                    # Add new version
                    version = DocumentVersion(
                        document_id=document.id,
                        version_number=new_version,
                        file_hash=file_hash,
                        change_type='content_update',
                        change_description='File content changed',
                        created_by=created_by
                    )
                    session.add(version)

                    # Log update
                    log_entry = ProcessingLog(
                        document_id=document.id,
                        action='version_updated',
                        status='success',
                        message=f'Document updated to version {new_version}',
                        details={'old_hash': document.current_file_hash, 'new_hash': file_hash},
                        created_by=created_by,
                        duration_seconds=int((datetime.utcnow() - start_time).total_seconds())
                    )
                    session.add(log_entry)

            await session.commit()
            await session.refresh(document)

            return document, is_new

    async def update_document_processing(
        self,
        document_id: int,
        status: str,
        message: Optional[str] = None,
        analysis_results: Optional[Dict[str, Any]] = None,
        duration_seconds: Optional[int] = None,
        memory_used_mb: Optional[int] = None
    ):
        """Update document processing status with enhanced tracking"""
        async with self.get_async_session() as session:
            result = await session.execute(
                select(Document).where(Document.id == document_id)
            )
            document = result.scalar_one_or_none()

            if document:
                document.processing_status = status

                if status == 'processing':
                    document.processing_started_at = datetime.utcnow()
                elif status == 'completed':
                    document.processing_completed_at = datetime.utcnow()
                    document.processed = True

                    if analysis_results:
                        document.analysis_results = analysis_results
                        self._extract_analysis_fields(document, analysis_results)

                        # Update current version as processed
                        await session.execute(
                            text("""
                                UPDATE document_versions
                                SET processed = true,
                                    processing_completed_at = :completed_at,
                                    analysis_results = :results
                                WHERE document_id = :doc_id
                                AND version_number = :version
                            """),
                            {
                                "completed_at": document.processing_completed_at,
                                "results": json.dumps(analysis_results),
                                "doc_id": document_id,
                                "version": document.processing_version
                            }
                        )

                elif status == 'failed':
                    document.error_message = message

                # Add log entry
                log_entry = ProcessingLog(
                    document_id=document_id,
                    action=f'processing_{status}',
                    status=status,
                    message=message,
                    details={'analysis_results': analysis_results} if analysis_results else None,
                    duration_seconds=duration_seconds,
                    memory_used_mb=memory_used_mb
                )
                session.add(log_entry)

                await session.commit()

                # Update daily statistics
                await self._update_processing_statistics(session, status, duration_seconds)

                # Invalidate cache
                await self._invalidate_document_cache(document.current_file_hash)

    def _extract_analysis_fields(self, document: Document, analysis_results: Dict[str, Any]):
        """Extract specific fields from analysis results"""
        for analysis in analysis_results.get('analyses', []):
            if analysis.get('analysis_type') == 'contract_details':
                result = analysis.get('result', {})
                if isinstance(result, dict):
                    document.contract_title = result.get('title')
                    document.contract_start_date = self._parse_date(result.get('start_date'))
                    document.contract_end_date = self._parse_date(result.get('end_date'))
                    document.contract_duration = result.get('contract_duration')
                    document.key_deliverables = result.get('key_deliverables')

                    # Extract contract value if present
                    if 'value' in result:
                        try:
                            document.contract_value = Decimal(str(result['value']))
                            document.currency = result.get('currency', 'USD')
                        except:
                            pass

    def _parse_date(self, date_str: Optional[str]) -> Optional[date]:
        """Parse date string in DD/MM/YYYY format"""
        if not date_str:
            return None
        try:
            return datetime.strptime(date_str, '%d/%m/%Y').date()
        except:
            # Try other formats
            for fmt in ['%Y-%m-%d', '%m/%d/%Y', '%d.%m.%Y']:
                try:
                    return datetime.strptime(date_str, fmt).date()
                except:
                    continue
            return None

    async def _update_processing_statistics(self, session: AsyncSession, status: str, duration_seconds: Optional[int]):
        """Update daily processing statistics"""
        today = date.today()

        # Get or create today's statistics
        result = await session.execute(
            select(ProcessingStatistics).where(ProcessingStatistics.date == today)
        )
        stats = result.scalar_one_or_none()

        if not stats:
            stats = ProcessingStatistics(date=today)
            session.add(stats)

        # Update counts
        if status == 'completed':
            stats.success_count += 1
            stats.total_documents_processed += 1
        elif status == 'failed':
            stats.failure_count += 1

        if duration_seconds:
            stats.total_processing_time_seconds += duration_seconds

        # Recalculate average
        if stats.total_documents_processed > 0:
            stats.average_processing_time_seconds = Decimal(
                stats.total_processing_time_seconds / stats.total_documents_processed
            )

    async def _invalidate_document_cache(self, file_hash: str):
        """Invalidate document cache entries"""
        cache_key = f"doc_hash:{file_hash}"
        if cache_key in self._cache:
            del self._cache[cache_key]
        if self._redis_client:
            self._redis_client.delete(cache_key)

    async def assign_document_to_ontology(
        self,
        document_id: int,
        ontology_code: str,
        confidence: float = 1.0,
        is_primary: bool = False,
        assigned_by: str = 'user'
    ):
        """Assign document to ontology category"""
        async with self.get_async_session() as session:
            # Get ontology
            result = await session.execute(
                select(ContractOntology).where(ContractOntology.category_code == ontology_code)
            )
            ontology = result.scalar_one_or_none()

            if not ontology:
                raise ValueError(f"Ontology category not found: {ontology_code}")

            # Check if mapping exists
            result = await session.execute(
                select(DocumentOntologyMapping).where(
                    and_(
                        DocumentOntologyMapping.document_id == document_id,
                        DocumentOntologyMapping.ontology_id == ontology.id
                    )
                )
            )
            mapping = result.scalar_one_or_none()

            if not mapping:
                mapping = DocumentOntologyMapping(
                    document_id=document_id,
                    ontology_id=ontology.id,
                    confidence_score=Decimal(str(confidence)),
                    is_primary=is_primary,
                    assigned_by=assigned_by
                )
                session.add(mapping)
            else:
                mapping.confidence_score = Decimal(str(confidence))
                mapping.is_primary = is_primary
                mapping.assigned_by = assigned_by
                mapping.assigned_at = datetime.utcnow()

            # If this is primary, unset other primary mappings
            if is_primary:
                await session.execute(
                    text("""
                        UPDATE document_ontology_mapping
                        SET is_primary = false
                        WHERE document_id = :doc_id AND ontology_id != :ont_id
                    """),
                    {"doc_id": document_id, "ont_id": ontology.id}
                )

            await session.commit()

    async def get_ontology_tree(self) -> List[Dict[str, Any]]:
        """Get complete ontology tree structure"""
        async with self.get_async_session() as session:
            result = await session.execute(
                select(ContractOntology).order_by(ContractOntology.level, ContractOntology.sort_order)
            )
            ontologies = result.scalars().all()

            # Build tree structure
            tree = []
            node_map = {}

            for ont in ontologies:
                node = {
                    'id': ont.id,
                    'code': ont.category_code,
                    'name': ont.category_name,
                    'level': ont.level,
                    'description': ont.description,
                    'color': ont.color_hex,
                    'children': []
                }
                node_map[ont.id] = node

                if ont.parent_category_id:
                    parent = node_map.get(ont.parent_category_id)
                    if parent:
                        parent['children'].append(node)
                else:
                    tree.append(node)

            return tree

    async def create_document_relationship(
        self,
        parent_doc_id: int,
        child_doc_id: int,
        relationship_type: str,
        confidence: float = 1.0,
        detected_by: str = 'user',
        metadata: Optional[Dict[str, Any]] = None
    ):
        """Create a relationship between documents with enhanced metadata"""
        async with self.get_async_session() as session:
            # Check if relationship exists
            result = await session.execute(
                select(func.count()).select_from(document_relationships).where(
                    and_(
                        document_relationships.c.parent_document_id == parent_doc_id,
                        document_relationships.c.child_document_id == child_doc_id,
                        document_relationships.c.relationship_type == relationship_type
                    )
                )
            )

            if result.scalar() == 0:
                # Insert new relationship
                stmt = document_relationships.insert().values(
                    parent_document_id=parent_doc_id,
                    child_document_id=child_doc_id,
                    relationship_type=relationship_type,
                    confidence=Decimal(str(confidence)),
                    detected_by=detected_by,
                    metadata=metadata
                )
                await session.execute(stmt)
                await session.commit()

                # Update statistics
                await self._update_relationship_statistics(session)

                # Update Neo4j if available
                if self._neo4j_graph:
                    try:
                        parent_node = Node("Document", document_id=parent_doc_id)
                        child_node = Node("Document", document_id=child_doc_id)
                        rel = Neo4jRelationship(parent_node, relationship_type.upper(), child_node, **metadata or {})
                        self._neo4j_graph.create(rel)
                    except Exception as e:
                        logging.error(f"Failed to update Neo4j: {e}")

    async def _update_relationship_statistics(self, session: AsyncSession):
        """Update relationship discovery statistics"""
        today = date.today()
        result = await session.execute(
            select(ProcessingStatistics).where(ProcessingStatistics.date == today)
        )
        stats = result.scalar_one_or_none()

        if stats:
            stats.new_relationships_found += 1

    async def link_document_to_company(
        self,
        document_id: int,
        company_id: str,
        role: str = 'vendor',
        confidence: float = 1.0,
        is_primary: bool = False
    ):
        """Link a document to a company with enhanced metadata"""
        async with self.get_async_session() as session:
            # Get or create company
            result = await session.execute(
                select(Company).where(Company.company_id == company_id)
            )
            company = result.scalar_one_or_none()

            if not company:
                company = Company(
                    company_id=company_id,
                    name=f"Company {company_id}"
                )
                session.add(company)
                await session.flush()

            # Check if link exists
            result = await session.execute(
                select(func.count()).select_from(document_companies).where(
                    and_(
                        document_companies.c.document_id == document_id,
                        document_companies.c.company_id == company.id,
                        document_companies.c.role == role
                    )
                )
            )

            if result.scalar() == 0:
                # Create link
                stmt = document_companies.insert().values(
                    document_id=document_id,
                    company_id=company.id,
                    role=role,
                    confidence_score=Decimal(str(confidence)),
                    is_primary=is_primary
                )
                await session.execute(stmt)

                # If this is primary, unset other primary links
                if is_primary:
                    await session.execute(
                        text("""
                            UPDATE document_companies
                            SET is_primary = false
                            WHERE document_id = :doc_id
                            AND company_id != :comp_id
                        """),
                        {"doc_id": document_id, "comp_id": company.id}
                    )

                await session.commit()

                # Update company statistics
                await self._update_company_statistics(session)

    async def _update_company_statistics(self, session: AsyncSession):
        """Update unique company count in statistics"""
        today = date.today()

        # Count unique companies
        result = await session.execute(
            select(func.count(func.distinct(Company.id))).select_from(Company)
        )
        unique_companies = result.scalar()

        # Update statistics
        result = await session.execute(
            select(ProcessingStatistics).where(ProcessingStatistics.date == today)
        )
        stats = result.scalar_one_or_none()

        if stats:
            stats.unique_companies = unique_companies

    def build_knowledge_graph(self) -> nx.DiGraph:
        """Build an enhanced NetworkX knowledge graph from the database"""
        G = nx.DiGraph()

        with self.get_session() as session:
            # Add all documents as nodes with detailed attributes
            documents = session.query(Document).all()
            for doc in documents:
                # Get primary ontology
                primary_ont = session.query(
                    DocumentOntologyMapping, ContractOntology
                ).join(
                    ContractOntology
                ).filter(
                    DocumentOntologyMapping.document_id == doc.id,
                    DocumentOntologyMapping.is_primary == True
                ).first()

                node_attrs = {
                    'type': 'document',
                    'name': doc.file_name,
                    'cw_number': doc.cw_number,
                    'processed': doc.processed,
                    'status': doc.processing_status,
                    'version': doc.processing_version,
                    'contract_type': doc.contract_type.type_name if doc.contract_type else None,
                    'start_date': doc.contract_start_date.isoformat() if doc.contract_start_date else None,
                    'end_date': doc.contract_end_date.isoformat() if doc.contract_end_date else None,
                    'value': float(doc.contract_value) if doc.contract_value else None,
                    'ontology': primary_ont[1].category_name if primary_ont else None,
                    'ontology_color': primary_ont[1].color_hex if primary_ont else '#999999'
                }
                G.add_node(f"doc_{doc.id}", **node_attrs)

            # Add all companies as nodes
            companies = session.query(Company).all()
            for comp in companies:
                G.add_node(
                    f"company_{comp.id}",
                    type='company',
                    name=comp.name,
                    company_id=comp.company_id,
                    company_group=comp.company_group,
                    industry=comp.industry,
                    country=comp.country
                )

            # Add document relationships with metadata
            relationships = session.execute(
                select(document_relationships)
            ).fetchall()

            for rel in relationships:
                G.add_edge(
                    f"doc_{rel.parent_document_id}",
                    f"doc_{rel.child_document_id}",
                    relationship_type=rel.relationship_type,
                    confidence=float(rel.confidence),
                    detected_by=rel.detected_by,
                    metadata=rel.metadata
                )

            # Add company-document relationships
            company_docs = session.execute(
                select(document_companies)
            ).fetchall()

            for cd in company_docs:
                G.add_edge(
                    f"company_{cd.company_id}",
                    f"doc_{cd.document_id}",
                    relationship_type=cd.role,
                    confidence=float(cd.confidence_score),
                    is_primary=cd.is_primary
                )

            # Add ontology structure as subgraph
            ontologies = session.query(ContractOntology).all()
            for ont in ontologies:
                G.add_node(
                    f"ont_{ont.id}",
                    type='ontology',
                    name=ont.category_name,
                    code=ont.category_code,
                    level=ont.level,
                    color=ont.color_hex,
                    description=ont.description
                )

                if ont.parent_category_id:
                    G.add_edge(
                        f"ont_{ont.parent_category_id}",
                        f"ont_{ont.id}",
                        relationship_type='parent_child'
                    )

        return G

    async def get_processing_progress(self) -> Dict[str, Any]:
        """Get current processing progress statistics"""
        async with self.get_async_session() as session:
            # Total documents
            total_result = await session.execute(
                select(func.count()).select_from(Document)
            )
            total_documents = total_result.scalar()

            # Status breakdown
            status_result = await session.execute(
                select(
                    Document.processing_status,
                    func.count().label('count')
                ).group_by(Document.processing_status)
            )
            status_breakdown = {row.processing_status: row.count for row in status_result}

            # Today's statistics
            today = date.today()
            stats_result = await session.execute(
                select(ProcessingStatistics).where(ProcessingStatistics.date == today)
            )
            today_stats = stats_result.scalar_one_or_none()

            return {
                'total_documents': total_documents,
                'processed': status_breakdown.get('completed', 0),
                'processing': status_breakdown.get('processing', 0),
                'pending': status_breakdown.get('pending', 0),
                'failed': status_breakdown.get('failed', 0),
                'today_processed': today_stats.total_documents_processed if today_stats else 0,
                'today_success': today_stats.success_count if today_stats else 0,
                'today_failed': today_stats.failure_count if today_stats else 0,
                'avg_processing_time': float(today_stats.average_processing_time_seconds) if today_stats else 0
            }

    def close(self):
        """Close all database connections"""
        if self._engine:
            self._engine.dispose()
        if self._async_engine:
            asyncio.run(self._async_engine.dispose())
        if self._redis_client:
            self._redis_client.close()
        if self._neo4j_graph:
            self._neo4j_graph = None
        logging.info("Database connections closed")

# Import guard for SQLAlchemy async
try:
    from sqlalchemy import select, text
except ImportError:
    logging.warning("SQLAlchemy async features not available")