Tracker/Code/Tracker/pipeline/pipeline.go

package pipeline

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/goko/jingtian-tracker/clients"
	"github.com/goko/jingtian-tracker/config"
	"github.com/goko/jingtian-tracker/db"
	"github.com/goko/jingtian-tracker/watcher"
)

// TrackerColumns defines the tracker spreadsheet columns
var TrackerColumns = []string{
	"Document Path",
	"Document Type",
	"Client Name",
	"Matter Reference",
	"TM Number",
	"Trademark Name",
	"Trademark Class",
	"Filing Date",
	"Response Deadline",
	"Hearing Date",
	"Amount",
	"Currency",
	"Status",
	"First Processed",
	"Last Updated",
	"Update Source",
	"Processing Notes",
}

type Pipeline struct {
	cfg      *config.Config
	db       *db.DB
	watcher  *watcher.Watcher
	docling  *clients.DoclingClient
	ollama   *clients.OllamaClient
	tools    *clients.ToolsClient
	errorLog *os.File
}

func New(cfg *config.Config, database *db.DB, w *watcher.Watcher) (*Pipeline, error) {
	// Open error log file
	var errorLog *os.File
	if cfg.Logging.ErrorFile != "" {
		// Ensure directory exists
		dir := filepath.Dir(cfg.Logging.ErrorFile)
		if err := os.MkdirAll(dir, 0755); err != nil {
			return nil, fmt.Errorf("failed to create log directory: %w", err)
		}

		f, err := os.OpenFile(cfg.Logging.ErrorFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			return nil, fmt.Errorf("failed to open error log: %w", err)
		}
		errorLog = f
	}

	return &Pipeline{
		cfg:      cfg,
		db:       database,
		watcher:  w,
		docling:  clients.NewDoclingClient(cfg.Services.Docling),
		ollama:   clients.NewOllamaClient(cfg.Services.Ollama, cfg.Ollama.Model, cfg.Ollama.Temperature),
		tools:    clients.NewToolsClient(cfg.Services.Tools),
		errorLog: errorLog,
	}, nil
}

func (p *Pipeline) Close() {
	if p.errorLog != nil {
		p.errorLog.Close()
	}
}

func (p *Pipeline) Run() error {
	// Check service health
	if err := p.checkServices(); err != nil {
		return fmt.Errorf("service health check failed: %w", err)
	}

	// Ensure tracker exists
	if err := p.ensureTracker(); err != nil {
		log.Printf("[pipeline] Warning: could not ensure tracker: %v", err)
		// Don't fail - tracker might be created manually
	}

	// Start watcher
	if err := p.watcher.Start(); err != nil {
		return fmt.Errorf("failed to start watcher: %w", err)
	}

	log.Println("[pipeline] Started processing loop")

	// Process events
	for event := range p.watcher.Events {
		p.processFile(event)
	}

	return nil
}

func (p *Pipeline) checkServices() error {
	log.Println("[pipeline] Checking service health...")

	if err := p.docling.Health(); err != nil {
		return fmt.Errorf("docling: %w", err)
	}
	log.Println("[pipeline] Docling: OK")

	if err := p.ollama.Health(); err != nil {
		return fmt.Errorf("ollama: %w", err)
	}
	log.Println("[pipeline] Ollama: OK")

	if err := p.tools.Health(); err != nil {
		return fmt.Errorf("tools: %w", err)
	}
	log.Println("[pipeline] Tools: OK")

	return nil
}

func (p *Pipeline) ensureTracker() error {
	// Check if tracker exists by trying to read it
	_, err := p.tools.ReadExcel(p.cfg.Tracker.Path, p.cfg.Tracker.Sheet, p.cfg.Tracker.HeaderRow)
	if err != nil {
		// Try to create it
		log.Printf("[pipeline] Creating tracker: %s", p.cfg.Tracker.Path)
		if createErr := p.tools.CreateTracker(p.cfg.Tracker.Path, p.cfg.Tracker.Sheet, TrackerColumns); createErr != nil {
			return fmt.Errorf("failed to create tracker: %w", createErr)
		}
		log.Printf("[pipeline] Tracker created successfully")
	}
	return nil
}

func (p *Pipeline) processFile(event watcher.FileEvent) {
	// Get relative path
	relPath, err := filepath.Rel(p.cfg.Watch.Root, event.Path)
	if err != nil {
		relPath = event.Path
	}

	log.Printf("[pipeline] Processing: %s", relPath)

	// Check if processing needed
	needsProcessing, err := p.db.NeedsProcessing(relPath, event.Hash)
	if err != nil {
		p.logError(relPath, fmt.Sprintf("failed to check processing status: %v", err))
		return
	}

	if !needsProcessing {
		log.Printf("[pipeline] Skipping (already processed with same hash): %s", relPath)
		return
	}

	// Upsert document record
	doc, err := p.db.UpsertDocument(relPath, event.Hash, event.Size, event.ModTime)
	if err != nil {
		p.logError(relPath, fmt.Sprintf("failed to upsert document: %v", err))
		return
	}

	// Update status to processing
	p.db.UpdateStatus(doc.ID, "processing")

	// Process based on file type
	ext := strings.ToLower(filepath.Ext(event.Path))
	var extracted map[string]interface{}

	switch ext {
	case ".xlsx", ".xls":
		// For Excel files, we don't OCR, just log that we saw it
		log.Printf("[pipeline] Excel file detected, skipping OCR: %s", relPath)
		p.db.LogAction(doc.ID, "skipped", "Excel files are not processed through OCR")
		p.db.MarkProcessed(doc.ID)
		return

	default:
		// Send to Docling for OCR/extraction
		content, doclingErr := p.docling.Convert(event.Path)
		if doclingErr != nil {
			p.logError(relPath, fmt.Sprintf("Docling failed: %v", doclingErr))
			p.db.UpdateStatus(doc.ID, "failed")
			p.db.LogAction(doc.ID, "failed", fmt.Sprintf("Docling error: %v", doclingErr))
			return
		}

		if content == "" {
			p.logError(relPath, "Docling returned empty content")
			p.db.UpdateStatus(doc.ID, "failed")
			p.db.LogAction(doc.ID, "failed", "Empty content from Docling")
			return
		}

		// Detect document type from extension
		docType := detectDocType(ext)

		// Send to Ollama for extraction
		extracted, err = p.ollama.Extract(content, docType)
		if err != nil {
			p.logError(relPath, fmt.Sprintf("Ollama failed: %v", err))
			p.db.UpdateStatus(doc.ID, "failed")
			p.db.LogAction(doc.ID, "failed", fmt.Sprintf("Ollama error: %v", err))
			return
		}
	}

	// Save extraction to database
	if err := p.db.SaveExtraction(doc.ID, extracted); err != nil {
		p.logError(relPath, fmt.Sprintf("failed to save extraction: %v", err))
	}

	// Write to tracker
	if err := p.writeToTracker(relPath, extracted); err != nil {
		p.logError(relPath, fmt.Sprintf("failed to write to tracker: %v", err))
		p.db.UpdateStatus(doc.ID, "failed")
		p.db.LogAction(doc.ID, "failed", fmt.Sprintf("Tracker write error: %v", err))
		return
	}

	// Mark as completed
	p.db.MarkProcessed(doc.ID)
	p.db.LogAction(doc.ID, "processed", "Successfully processed and added to tracker")
	log.Printf("[pipeline] Successfully processed: %s", relPath)
}

func (p *Pipeline) writeToTracker(docPath string, extracted map[string]interface{}) error {
	// Check if row exists for this document
	existing, err := p.tools.FindRow(
		p.cfg.Tracker.Path,
		p.cfg.Tracker.Sheet,
		p.cfg.Tracker.HeaderRow,
		"Document Path",
		docPath,
	)
	if err != nil {
		// Row not found is OK - we'll add a new one
		existing = &clients.ExcelFindResponse{Found: false}
	}

	now := time.Now().Format("02-01-2006 15:04:05")

	// Build row data
	row := map[string]interface{}{
		"Document Path":     docPath,
		"Document Type":     getStringField(extracted, "document_type"),
		"Client Name":       getStringField(extracted, "client_name"),
		"Matter Reference":  getStringField(extracted, "matter_reference"),
		"TM Number":         getStringField(extracted, "tm_number"),
		"Trademark Name":    getStringField(extracted, "trademark_name"),
		"Trademark Class":   getStringField(extracted, "trademark_class"),
		"Filing Date":       getStringField(extracted, "filing_date"),
		"Response Deadline": getStringField(extracted, "response_deadline"),
		"Hearing Date":      getStringField(extracted, "hearing_date"),
		"Amount":            getStringField(extracted, "amount"),
		"Currency":          getStringField(extracted, "currency"),
		"Status":            getStringField(extracted, "status"),
		"Last Updated":      now,
		"Update Source":     "Pipeline",
		"Processing Notes":  "",
	}

	if existing.Found {
		// Check if manually edited
		if updateSource, ok := existing.Row["Update Source"].(string); ok && updateSource == "Manual" {
			log.Printf("[pipeline] Skipping update (manually edited): %s", docPath)
			return nil
		}

		// Update existing row
		// Preserve First Processed
		if fp, ok := existing.Row["First Processed"]; ok {
			row["First Processed"] = fp
		}

		return p.tools.UpdateRow(
			p.cfg.Tracker.Path,
			p.cfg.Tracker.Sheet,
			p.cfg.Tracker.HeaderRow,
			existing.RowNum,
			row,
		)
	}

	// Add new row
	row["First Processed"] = now
	return p.tools.AddRow(
		p.cfg.Tracker.Path,
		p.cfg.Tracker.Sheet,
		p.cfg.Tracker.HeaderRow,
		row,
	)
}

func (p *Pipeline) logError(path, msg string) {
	errMsg := fmt.Sprintf("[%s] %s: %s", time.Now().Format("2006-01-02 15:04:05"), path, msg)
	log.Printf("[pipeline] ERROR: %s", errMsg)

	if p.errorLog != nil {
		p.errorLog.WriteString(errMsg + "\n")
	}
}

func detectDocType(ext string) string {
	switch ext {
	case ".pdf":
		return "PDF document"
	case ".docx", ".doc":
		return "Word document"
	case ".png", ".jpg", ".jpeg":
		return "Image/screenshot"
	default:
		return "Document"
	}
}

func getStringField(m map[string]interface{}, key string) string {
	v, ok := m[key]
	if !ok || v == nil {
		return ""
	}
	if s, ok := v.(string); ok {
		return s
	}
	return fmt.Sprintf("%v", v)
}