gridpilot.gg/packages/infrastructure/adapters/automation/TemplateMatchingService.ts

import { screen, imageResource, Region } from '@nut-tree-fork/nut-js';
import type { ImageTemplate } from '../../../domain/value-objects/ImageTemplate';
import type { ElementLocation, ScreenRegion, Point } from '../../../domain/value-objects/ScreenRegion';
import type { ILogger } from '../../../application/ports/ILogger';
import { NoOpLogAdapter } from '../logging/NoOpLogAdapter';
import * as fs from 'fs';
import * as path from 'path';

/**
 * Result of a template matching operation.
 */
export interface TemplateMatchResult {
  found: boolean;
  location?: ElementLocation;
  searchDurationMs: number;
  error?: string;
}

/**
 * Service for finding UI elements on screen using image template matching.
 * Uses nut.js screen.find() with @nut-tree-fork/template-matcher for image comparison.
 */
export class TemplateMatchingService {
  private logger: ILogger;
  private templateBasePath: string;
  private loadedTemplates: Map<string, string> = new Map();

  constructor(templateBasePath?: string, logger?: ILogger) {
    this.logger = logger ?? new NoOpLogAdapter();
    this.templateBasePath = templateBasePath ?? './resources/templates';
  }

  /**
   * Find an element on screen using image template matching.
   * @param template - The image template to search for
   * @returns ElementLocation if found, null if not found
   */
  async findElement(template: ImageTemplate): Promise<ElementLocation | null> {
    const startTime = Date.now();
    
    try {
      const fullPath = this.resolveTemplatePath(template.imagePath);
      
      if (!this.templateExists(fullPath)) {
        this.logger.warn('Template image not found, skipping search', { 
          templateId: template.id, 
          imagePath: fullPath,
          description: template.description 
        });
        return null;
      }

      this.logger.debug('Searching for template', { 
        templateId: template.id, 
        confidence: template.confidence,
        searchRegion: template.searchRegion 
      });

      const templateImage = await imageResource(fullPath);
      
      // Set confidence threshold for this search
      screen.config.confidence = template.confidence;
      
      let foundRegion: Region;
      
      if (template.searchRegion) {
        // Search within a specific region for better performance
        const searchArea = new Region(
          template.searchRegion.x,
          template.searchRegion.y,
          template.searchRegion.width,
          template.searchRegion.height
        );
        foundRegion = await screen.findRegion(templateImage, searchArea);
      } else {
        // Search entire screen
        foundRegion = await screen.find(templateImage);
      }

      const location = this.regionToElementLocation(foundRegion, template.confidence);
      
      const durationMs = Date.now() - startTime;
      this.logger.debug('Template found', { 
        templateId: template.id, 
        location: location.center,
        confidence: location.confidence,
        durationMs 
      });

      return location;
    } catch (error) {
      const durationMs = Date.now() - startTime;
      
      // nut.js throws when template is not found - this is expected behavior
      if (String(error).includes('No match')) {
        this.logger.debug('Template not found on screen', { 
          templateId: template.id, 
          durationMs 
        });
        return null;
      }
      
      this.logger.error('Template matching error', error instanceof Error ? error : new Error(String(error)), {
        templateId: template.id,
        durationMs
      });
      return null;
    }
  }

  /**
   * Find multiple occurrences of a template on screen.
   * @param template - The image template to search for
   * @returns Array of ElementLocation for all matches
   */
  async findAllElements(template: ImageTemplate): Promise<ElementLocation[]> {
    const startTime = Date.now();
    
    try {
      const fullPath = this.resolveTemplatePath(template.imagePath);
      
      if (!this.templateExists(fullPath)) {
        this.logger.warn('Template image not found', { templateId: template.id, imagePath: fullPath });
        return [];
      }

      const templateImage = await imageResource(fullPath);
      screen.config.confidence = template.confidence;
      
      const foundRegions = await screen.findAll(templateImage);
      
      const locations = foundRegions.map(region => this.regionToElementLocation(region, template.confidence));
      
      const durationMs = Date.now() - startTime;
      this.logger.debug('Found multiple templates', { 
        templateId: template.id, 
        count: locations.length,
        durationMs 
      });

      return locations;
    } catch (error) {
      const durationMs = Date.now() - startTime;
      
      if (String(error).includes('No match')) {
        this.logger.debug('No templates found on screen', { templateId: template.id, durationMs });
        return [];
      }
      
      this.logger.error('Template matching error', error instanceof Error ? error : new Error(String(error)), {
        templateId: template.id,
        durationMs
      });
      return [];
    }
  }

  /**
   * Wait for a template to appear on screen.
   * @param template - The image template to wait for
   * @param maxWaitMs - Maximum time to wait in milliseconds
   * @param pollIntervalMs - How often to check for the template
   * @returns TemplateMatchResult with timing information
   */
  async waitForTemplate(
    template: ImageTemplate, 
    maxWaitMs: number = 10000,
    pollIntervalMs: number = 500
  ): Promise<TemplateMatchResult> {
    const startTime = Date.now();
    
    this.logger.debug('Waiting for template', { 
      templateId: template.id, 
      maxWaitMs, 
      pollIntervalMs 
    });

    while (Date.now() - startTime < maxWaitMs) {
      const location = await this.findElement(template);
      
      if (location) {
        const durationMs = Date.now() - startTime;
        this.logger.debug('Template appeared', { templateId: template.id, durationMs });
        
        return {
          found: true,
          location,
          searchDurationMs: durationMs,
        };
      }

      await this.delay(pollIntervalMs);
    }

    const durationMs = Date.now() - startTime;
    this.logger.debug('Template did not appear within timeout', { 
      templateId: template.id, 
      durationMs,
      maxWaitMs 
    });

    return {
      found: false,
      searchDurationMs: durationMs,
      error: `Template '${template.id}' not found within ${maxWaitMs}ms`,
    };
  }

  /**
   * Check if a template exists on screen without returning location.
   * Useful for boolean state checks.
   */
  async isTemplateVisible(template: ImageTemplate): Promise<boolean> {
    const location = await this.findElement(template);
    return location !== null;
  }

  /**
   * Resolve a template path relative to the base template directory.
   */
  private resolveTemplatePath(imagePath: string): string {
    if (path.isAbsolute(imagePath)) {
      return imagePath;
    }
    return path.join(this.templateBasePath, imagePath);
  }

  /**
   * Check if a template file exists.
   */
  private templateExists(fullPath: string): boolean {
    try {
      return fs.existsSync(fullPath);
    } catch {
      return false;
    }
  }

  /**
   * Convert nut.js Region to ElementLocation.
   */
  private regionToElementLocation(region: Region, confidence: number): ElementLocation {
    const bounds: ScreenRegion = {
      x: region.left,
      y: region.top,
      width: region.width,
      height: region.height,
    };

    const center: Point = {
      x: region.left + Math.floor(region.width / 2),
      y: region.top + Math.floor(region.height / 2),
    };

    return {
      center,
      bounds,
      confidence,
    };
  }

  /**
   * Utility delay function.
   */
  private delay(ms: number): Promise<void> {
    return new Promise(resolve => setTimeout(resolve, ms));
  }

  /**
   * Get the configured template base path.
   */
  getTemplateBasePath(): string {
    return this.templateBasePath;
  }

  /**
   * Set the template base path.
   */
  setTemplateBasePath(basePath: string): void {
    this.templateBasePath = basePath;
    this.loadedTemplates.clear();
  }
}