brainy/src/importers/SmartDOCXImporter.ts

/**
 * Smart DOCX Importer
 *
 * Extracts entities and relationships from Word documents using:
 * - Mammoth parser for DOCX → HTML/text conversion
 * - Heading extraction for document structure
 * - Table extraction for structured data
 * - NeuralEntityExtractor for entity extraction from paragraphs
 * - NaturalLanguageProcessor for relationship inference
 * - Hierarchical relationship creation based on heading hierarchy
 *
 * v4.2.0: New format handler
 * NO MOCKS - Production-ready implementation
 */

import { Brainy } from '../brainy.js'
import { NeuralEntityExtractor, ExtractedEntity } from '../neural/entityExtractor.js'
import { NaturalLanguageProcessor } from '../neural/naturalLanguageProcessor.js'
import { SmartRelationshipExtractor } from '../neural/SmartRelationshipExtractor.js'
import { NounType, VerbType } from '../types/graphTypes.js'

// Mammoth type definitions (no @types package available)
interface MammothResult {
  value: string
  messages: Array<{
    type: string
    message: string
  }>
}

interface Mammoth {
  extractRawText(options: { buffer: Buffer }): Promise<MammothResult>
  convertToHtml(options: { buffer: Buffer }): Promise<MammothResult>
}

// Dynamic import for mammoth (ESM compatibility)
let mammoth: Mammoth

export interface SmartDOCXOptions {
  /** Enable neural entity extraction from paragraphs */
  enableNeuralExtraction?: boolean

  /** Enable hierarchical relationship creation based on headings */
  enableHierarchicalRelationships?: boolean

  /** Enable concept extraction for tagging */
  enableConceptExtraction?: boolean

  /** Confidence threshold for entities (0-1) */
  confidenceThreshold?: number

  /** Minimum paragraph length to process */
  minParagraphLength?: number

  /** Progress callback */
  onProgress?: (stats: {
    processed: number
    entities: number
    relationships: number
  }) => void
}

export interface ExtractedDOCXEntity {
  /** Entity ID */
  id: string

  /** Entity name */
  name: string

  /** Entity type */
  type: NounType

  /** Entity description/context */
  description: string

  /** Confidence score */
  confidence: number

  /** Weight/importance score */
  weight?: number

  /** Section/heading context */
  section: string | null

  /** Paragraph index in document */
  paragraphIndex: number

  /** Metadata */
  metadata: Record<string, any>
}

export interface ExtractedDOCXRelationship {
  from: string
  to: string
  type: VerbType
  confidence: number
  weight?: number
  evidence: string
}

export interface SmartDOCXResult {
  /** Total paragraphs processed */
  paragraphsProcessed: number

  /** Entities extracted */
  entitiesExtracted: number

  /** Relationships inferred */
  relationshipsInferred: number

  /** All extracted entities */
  entities: ExtractedDOCXEntity[]

  /** All relationships */
  relationships: ExtractedDOCXRelationship[]

  /** Entity ID mapping (index -> ID) */
  entityMap: Map<string, string>

  /** Processing time in ms */
  processingTime: number

  /** Document structure */
  structure: {
    headings: Array<{ level: number; text: string; index: number }>
    paragraphCount: number
    tableCount: number
  }

  /** Extraction statistics */
  stats: {
    byType: Record<string, number>
    bySection: Record<string, number>
    byConfidence: {
      high: number  // > 0.8
      medium: number  // 0.6-0.8
      low: number  // < 0.6
    }
  }
}

/**
 * SmartDOCXImporter - Extracts structured knowledge from Word documents
 */
export class SmartDOCXImporter {
  private brain: Brainy
  private extractor: NeuralEntityExtractor
  private nlp: NaturalLanguageProcessor
  private relationshipExtractor: SmartRelationshipExtractor
  private mammothLoaded = false

  constructor(brain: Brainy) {
    this.brain = brain
    this.extractor = new NeuralEntityExtractor(brain)
    this.nlp = new NaturalLanguageProcessor(brain)
    this.relationshipExtractor = new SmartRelationshipExtractor(brain)
  }

  /**
   * Initialize the importer
   */
  async init(): Promise<void> {
    await this.nlp.init()

    // Lazy load mammoth
    if (!this.mammothLoaded) {
      try {
        mammoth = await import('mammoth')
        this.mammothLoaded = true
      } catch (error: any) {
        throw new Error(`Failed to load mammoth parser: ${error.message}`)
      }
    }
  }

  /**
   * Extract entities and relationships from DOCX buffer
   */
  async extract(
    buffer: Buffer,
    options: SmartDOCXOptions = {}
  ): Promise<SmartDOCXResult> {
    const startTime = Date.now()

    // Ensure mammoth is loaded
    if (!this.mammothLoaded) {
      await this.init()
    }

    // Extract raw text for entity extraction
    const textResult = await mammoth.extractRawText({ buffer })

    // Extract HTML for structure analysis (headings, tables)
    const htmlResult = await mammoth.convertToHtml({ buffer })

    // Process the document
    const result = await this.extractFromContent(
      textResult.value,
      htmlResult.value,
      options
    )

    result.processingTime = Date.now() - startTime

    return result
  }

  /**
   * Extract entities and relationships from parsed DOCX content
   */
  private async extractFromContent(
    rawText: string,
    html: string,
    options: SmartDOCXOptions
  ): Promise<SmartDOCXResult> {
    const opts = {
      enableNeuralExtraction: options.enableNeuralExtraction !== false,
      enableHierarchicalRelationships: options.enableHierarchicalRelationships !== false,
      enableConceptExtraction: options.enableConceptExtraction !== false,
      confidenceThreshold: options.confidenceThreshold || 0.6,
      minParagraphLength: options.minParagraphLength || 20
    }

    const entities: ExtractedDOCXEntity[] = []
    const relationships: ExtractedDOCXRelationship[] = []
    const entityMap = new Map<string, string>()

    const stats = {
      byType: {} as Record<string, number>,
      bySection: {} as Record<string, number>,
      byConfidence: { high: 0, medium: 0, low: 0 }
    }

    // Parse document structure from HTML
    const structure = this.parseStructure(html)

    // Split into paragraphs
    const paragraphs = rawText.split(/\n\n+/).filter(p => p.trim().length >= opts.minParagraphLength)

    let currentSection = 'Introduction'
    let headingIndex = 0

    // Process each paragraph
    for (let i = 0; i < paragraphs.length; i++) {
      const paragraph = paragraphs[i].trim()

      // Check if this paragraph is a heading
      if (headingIndex < structure.headings.length) {
        const heading = structure.headings[headingIndex]
        if (paragraph.startsWith(heading.text) || heading.text.includes(paragraph.substring(0, 50))) {
          currentSection = heading.text
          headingIndex++
          stats.bySection[currentSection] = 0
          continue
        }
      }

      // Extract entities from paragraph
      if (opts.enableNeuralExtraction) {
        const extractedEntities = await this.extractor.extract(paragraph, {
          confidence: opts.confidenceThreshold
        })

        for (const extracted of extractedEntities) {
          const entityId = `para${i}:${extracted.text}`
          const entity: ExtractedDOCXEntity = {
            id: entityId,
            name: extracted.text,
            type: extracted.type,
            description: paragraph,
            confidence: extracted.confidence,
            weight: extracted.weight || 1.0,
            section: currentSection,
            paragraphIndex: i,
            metadata: {
              position: extracted.position,
              headingContext: currentSection
            }
          }

          entities.push(entity)
          entityMap.set(entityId, entityId)

          // Update stats
          stats.byType[entity.type] = (stats.byType[entity.type] || 0) + 1
          stats.bySection[currentSection] = (stats.bySection[currentSection] || 0) + 1
          if (entity.confidence > 0.8) stats.byConfidence.high++
          else if (entity.confidence >= 0.6) stats.byConfidence.medium++
          else stats.byConfidence.low++
        }
      }

      // Report progress
      if (options.onProgress && i % 10 === 0) {
        options.onProgress({
          processed: i,
          entities: entities.length,
          relationships: relationships.length
        })
      }
    }

    // Create hierarchical relationships based on sections
    if (opts.enableHierarchicalRelationships) {
      const entitiesBySection = new Map<string, ExtractedDOCXEntity[]>()

      for (const entity of entities) {
        const section = entity.section || 'Unknown'
        if (!entitiesBySection.has(section)) {
          entitiesBySection.set(section, [])
        }
        entitiesBySection.get(section)!.push(entity)
      }

      // Create relationships within sections
      for (const [section, sectionEntities] of entitiesBySection) {
        for (let i = 0; i < sectionEntities.length - 1; i++) {
          for (let j = i + 1; j < Math.min(i + 3, sectionEntities.length); j++) {
            const entityA = sectionEntities[i]
            const entityB = sectionEntities[j]

            // Infer relationship type using SmartRelationshipExtractor
            // Combine entity descriptions for better context
            const context = `In section "${section}": ${entityA.description.substring(0, 150)}... ${entityB.description.substring(0, 150)}...`

            const inferredRelationship = await this.relationshipExtractor.infer(
              entityA.name,
              entityB.name,
              context,
              {
                subjectType: entityA.type,
                objectType: entityB.type
              }
            )

            relationships.push({
              from: entityA.id,
              to: entityB.id,
              type: inferredRelationship?.type || VerbType.RelatedTo,  // Fallback to RelatedTo for co-occurrence
              confidence: inferredRelationship?.confidence || 0.7,
              weight: inferredRelationship?.weight || 0.8,
              evidence: inferredRelationship?.evidence || `Both entities appear in section: ${section}`
            })
          }
        }
      }
    }

    // Final progress report
    if (options.onProgress) {
      options.onProgress({
        processed: paragraphs.length,
        entities: entities.length,
        relationships: relationships.length
      })
    }

    return {
      paragraphsProcessed: paragraphs.length,
      entitiesExtracted: entities.length,
      relationshipsInferred: relationships.length,
      entities,
      relationships,
      entityMap,
      processingTime: 0, // Will be set by caller
      structure,
      stats
    }
  }

  /**
   * Parse document structure from HTML
   */
  private parseStructure(html: string): {
    headings: Array<{ level: number; text: string; index: number }>
    paragraphCount: number
    tableCount: number
  } {
    const headings: Array<{ level: number; text: string; index: number }> = []

    // Extract headings (h1-h6)
    const headingRegex = /<h([1-6])>(.*?)<\/h\1>/gi
    let match
    let index = 0

    while ((match = headingRegex.exec(html)) !== null) {
      const level = parseInt(match[1])
      const text = match[2].replace(/<[^>]+>/g, '').trim() // Strip HTML tags
      headings.push({ level, text, index: index++ })
    }

    // Count paragraphs
    const paragraphCount = (html.match(/<p>/g) || []).length

    // Count tables
    const tableCount = (html.match(/<table>/g) || []).length

    return {
      headings,
      paragraphCount,
      tableCount
    }
  }
}