locator.ts

import sharp from 'sharp';
import type { MobilewrightDriver, ViewNode, Bounds, SwipeDirection, ScreenSize } from '@mobilewright/protocol';
import { queryAll, type LocatorStrategy, type Role } from './query-engine.js';
import { sleep } from './sleep.js';
import { runStep, type StepLocation } from './stackTrace.js';

export type StepFn = (title: string, fn: () => Promise<unknown>, location: StepLocation | undefined) => Promise<unknown>;

export interface LocatorOptions {
  timeout?: number;
  pollInterval?: number;
  stabilityDelay?: number;
  /** Default timeout for expect() assertions on this locator, in ms. */
  expectTimeout?: number;
}

export interface FilterOptions {
  /** Keep only elements whose subtree contains this text. */
  hasText?: string | RegExp;
  /** Keep only elements whose subtree does NOT contain this text. */
  hasNotText?: string | RegExp;
  /** Keep only elements that contain an element matching this locator. */
  has?: Locator;
  /** Keep only elements that do NOT contain an element matching this locator. */
  hasNot?: Locator;
}

export interface ScrollIntoViewOptions {
  /** Maximum number of swipe attempts before giving up (default: 10) */
  maxSwipes?: number;
  /** Swipe gesture direction — 'up' swipes up (scrolls content down), 'down' swipes down (scrolls content up). Default: 'up' */
  direction?: 'up' | 'down';
}

const DEFAULT_TIMEOUT = 5_000;
const DEFAULT_POLL_INTERVAL = 100;
const DEFAULT_STABILITY_DELAY = 50;

export class Locator {
  /** Create a root locator that searches the entire view hierarchy. */
  static root(driver: MobilewrightDriver, options: LocatorOptions = {}): Locator {
    return new Locator(driver, { kind: 'root' }, options);
  }

  _stepFn: StepFn | null = null;

  constructor(
    protected readonly driver: MobilewrightDriver,
    protected readonly strategy: LocatorStrategy,
    protected readonly options: LocatorOptions = {},
  ) {}

  get expectTimeout(): number | undefined {
    return this.options.expectTimeout;
  }

  protected async _step<T>(title: string, fn: () => Promise<T>): Promise<T> {
    return runStep(this._stepFn, title, fn);
  }

  // ─── Chaining ────────────────────────────────────────────────

  getByLabel(label: string, opts?: { exact?: boolean }): Locator {
    return this.child({ kind: 'label', value: label, exact: opts?.exact });
  }

  getByTestId(testId: string): Locator {
    return this.child({ kind: 'testId', value: testId });
  }

  getByText(text: string | RegExp, opts?: { exact?: boolean }): Locator {
    return this.child({ kind: 'text', value: text, exact: opts?.exact });
  }

  getByType(type: string): Locator {
    return this.child({ kind: 'type', value: type });
  }

  getByRole(role: Role, opts?: { name?: string | RegExp }): Locator {
    return this.child({ kind: 'role', value: role, name: opts?.name });
  }

  getByPlaceholder(placeholder: string, opts?: { exact?: boolean }): Locator {
    return this.child({ kind: 'placeholder', value: placeholder, exact: opts?.exact });
  }

  protected child(childStrategy: LocatorStrategy): Locator {
    return this.withStrategy({ kind: 'chain', parent: this.strategy, child: childStrategy });
  }

  /** Build a sibling locator that shares this one's driver, options, and step fn. */
  protected withStrategy(strategy: LocatorStrategy): Locator {
    const loc = new Locator(this.driver, strategy, this.options);
    loc._stepFn = this._stepFn;
    return loc;
  }

  // ─── Narrowing & combining ───────────────────────────────────

  /** Narrow this locator to elements matching the given conditions. */
  filter(opts?: FilterOptions): Locator {
    return this.withStrategy({
      kind: 'filter',
      parent: this.strategy,
      hasText: opts?.hasText,
      hasNotText: opts?.hasNotText,
      has: opts?.has?.strategy,
      hasNot: opts?.hasNot?.strategy,
    });
  }

  /** Match elements that satisfy both this locator and the given locator. */
  and(locator: Locator): Locator {
    return this.withStrategy({ kind: 'and', left: this.strategy, right: locator.strategy });
  }

  /** Match elements that satisfy either this locator or the given locator. */
  or(locator: Locator): Locator {
    return this.withStrategy({ kind: 'or', left: this.strategy, right: locator.strategy });
  }

  // ─── Collection ──────────────────────────────────────────────

  first(): Locator {
    return this.nth(0);
  }

  last(): Locator {
    return this.nth(-1);
  }

  nth(index: number): Locator {
    return this.withStrategy({ kind: 'nth', parent: this.strategy, index });
  }

  async count(): Promise<number> {
    const roots = await this.driver.getViewHierarchy();
    return queryAll(roots, this.strategy).length;
  }

  async all(): Promise<Locator[]> {
    const roots = await this.driver.getViewHierarchy();
    const matches = queryAll(roots, this.strategy);
    return matches.map((_, i) => {
      const loc = new Locator(
        this.driver,
        { kind: 'nth', parent: this.strategy, index: i },
        this.options,
      );
      loc._stepFn = this._stepFn;
      return loc;
    });
  }

  // ─── Actions ─────────────────────────────────────────────────

  async tap(opts?: { timeout?: number }): Promise<void> {
    return this._step('locator.tap()', async () => {
      const node = await this.resolveActionable(opts?.timeout);
      const { x, y } = centerOf(node.bounds);
      await this.driver.tap(x, y);
    });
  }

  async doubleTap(opts?: { timeout?: number }): Promise<void> {
    return this._step('locator.doubleTap()', async () => {
      const node = await this.resolveActionable(opts?.timeout);
      const { x, y } = centerOf(node.bounds);
      await this.driver.doubleTap(x, y);
    });
  }

  async longPress(opts?: { timeout?: number; duration?: number }): Promise<void> {
    return this._step('locator.longPress()', async () => {
      const node = await this.resolveActionable(opts?.timeout);
      const { x, y } = centerOf(node.bounds);
      await this.driver.longPress(x, y, opts?.duration);
    });
  }

  async fill(text: string, opts?: { timeout?: number }): Promise<void> {
    return this._step(`locator.fill(${JSON.stringify(text)})`, async () => {
      const node = await this.resolveActionable(opts?.timeout);
      const { x, y } = centerOf(node.bounds);
      await this.driver.tap(x, y);
      await this.driver.typeText(text);
    });
  }

  async screenshot(opts?: { timeout?: number }): Promise<Buffer> {
    return this._step('locator.screenshot()', async () => {
      const node = await this.resolveVisible(opts?.timeout);
      const fullScreenshot = await this.driver.screenshot();
      return cropToElement(fullScreenshot, node.bounds, await this.driver.getScreenSize());
    });
  }

  async swipe(opts: { direction: SwipeDirection; timeout?: number }): Promise<void> {
    return this._step(`locator.swipe(${opts.direction})`, async () => {
      const node = await this.resolveActionable(opts.timeout);
      const { x, y } = centerOf(node.bounds);
      await this.driver.swipe(opts.direction, { startX: x, startY: y });
    });
  }

  async scrollIntoViewIfNeeded(opts?: ScrollIntoViewOptions): Promise<void> {
    return this._step('locator.scrollIntoViewIfNeeded()', async () => {
      const maxSwipes = opts?.maxSwipes ?? 10;
      const direction: SwipeDirection = opts?.direction ?? 'up';
      const screenSize = await this.driver.getScreenSize();
      const POST_SWIPE_SETTLE = 200;

      for (let i = 0; i < maxSwipes; i++) {
        const roots = await this.driver.getViewHierarchy();
        const node = queryAll(roots, this.strategy)[0] ?? null;

        if (node && isWithinViewport(node.bounds, screenSize)) {
          return;
        }

        const swipeDirection = node ? swipeDirectionToReveal(node.bounds, screenSize) : direction;
        await this.driver.swipe(swipeDirection);
        await sleep(POST_SWIPE_SETTLE);
      }

      throw new LocatorError(
        `Element not scrolled into view after ${maxSwipes} swipes`,
        this.strategy,
      );
    });
  }

  // ─── Queries (with auto-wait for visibility) ─────────────────

  async exists(): Promise<boolean> {
    const node = await this.resolve(0);
    return node !== null;
  }

  async isVisible(opts?: { timeout?: number }): Promise<boolean> {
    try {
      await this.waitFor({ state: 'visible', timeout: opts?.timeout ?? 0 });
      return true;
    } catch (error) {
      if (!(error instanceof LocatorError)) {
        throw error;
      }
      return false;
    }
  }

  async isEnabled(opts?: { timeout?: number }): Promise<boolean> {
    const node = await this.resolve(opts?.timeout ?? 0);
    return node !== null && node.isEnabled;
  }

  async isSelected(opts?: { timeout?: number }): Promise<boolean> {
    const node = await this.resolve(opts?.timeout ?? 0);
    return node !== null && node.isSelected === true;
  }

  async isFocused(opts?: { timeout?: number }): Promise<boolean> {
    const node = await this.resolve(opts?.timeout ?? 0);
    return node !== null && node.isFocused === true;
  }

  async isChecked(opts?: { timeout?: number }): Promise<boolean> {
    const node = await this.resolve(opts?.timeout ?? 0);
    return node !== null && node.isChecked === true;
  }

  async boundingBox(opts?: { timeout?: number }): Promise<{ x: number; y: number; width: number; height: number }> {
    const node = await this.resolveVisible(opts?.timeout);
    return { x: node.bounds.x, y: node.bounds.y, width: node.bounds.width, height: node.bounds.height };
  }

  async getText(opts?: { timeout?: number }): Promise<string> {
    const node = await this.resolveVisible(opts?.timeout);
    return node.text ?? node.label ?? node.value ?? '';
  }

  async getValue(opts?: { timeout?: number }): Promise<string> {
    const node = await this.resolveVisible(opts?.timeout);
    return node.value ?? '';
  }

  async waitFor(opts?: {
    state?: 'visible' | 'hidden' | 'enabled' | 'disabled';
    timeout?: number;
  }): Promise<void> {
    await this.pollUntilState(opts?.state ?? 'visible', opts?.timeout);
  }

  // ─── Internal resolution ─────────────────────────────────────

  /** Wait for a visible node and return it. Used by getText, screenshot. */
  private async resolveVisible(timeout?: number): Promise<ViewNode> {
    const node = await this.pollUntilState('visible', timeout);
    return node!;
  }

  /** Poll until the given state is satisfied. Returns the matched node (or null for hidden/disabled). */
  private async pollUntilState(
    state: 'visible' | 'hidden' | 'enabled' | 'disabled',
    timeout?: number,
  ): Promise<ViewNode | null> {
    const effectiveTimeout = timeout ?? this.options.timeout ?? DEFAULT_TIMEOUT;
    const pollInterval = this.options.pollInterval ?? DEFAULT_POLL_INTERVAL;
    const deadline = Date.now() + effectiveTimeout;

    while (true) {
      const roots = await this.driver.getViewHierarchy();
      const node = queryAll(roots, this.strategy)[0] ?? null;

      if (checkState(node, state)) {
        return node;
      }

      if (Date.now() >= deadline) {
        throw new LocatorError(
          `Locator timed out waiting for state "${state}" after ${effectiveTimeout}ms`,
          this.strategy,
        );
      }

      await sleep(pollInterval);
    }
  }

  /** Resolve to a single actionable node (visible, enabled, stable bounds) */
  private async resolveActionable(
    timeout?: number,
  ): Promise<ViewNode> {
    const effectiveTimeout =
      timeout ?? this.options.timeout ?? DEFAULT_TIMEOUT;
    const pollInterval =
      this.options.pollInterval ?? DEFAULT_POLL_INTERVAL;
    const stabilityDelay =
      this.options.stabilityDelay ?? DEFAULT_STABILITY_DELAY;
    const deadline = Date.now() + effectiveTimeout;

    let previousBounds: Bounds | null = null;
    let lastReason = 'no matching element found';

    while (true) {
      const roots = await this.driver.getViewHierarchy();
      const node = queryAll(roots, this.strategy)[0];

      if (!node) {
        lastReason = 'no matching element found';
      } else if (!node.isVisible) {
        lastReason = 'element found but not visible';
      } else if (!node.isEnabled) {
        lastReason = 'element found but not enabled';
      } else {
        // Stability check: bounds haven't changed since last poll
        if (previousBounds && boundsEqual(previousBounds, node.bounds)) {
          return node;
        }
        previousBounds = { ...node.bounds };
        if (Date.now() >= deadline) {
          return node; // accept without stability
        }
        await sleep(stabilityDelay);
        continue;
      }

      if (Date.now() >= deadline) {
        throw new LocatorError(
          `Locator: ${lastReason} after ${effectiveTimeout}ms`,
          this.strategy,
        );
      }
      await sleep(pollInterval);
    }
  }

  /** Resolve without waiting — returns null if not found */
  private async resolve(timeout: number): Promise<ViewNode | null> {
    const deadline = Date.now() + timeout;
    const pollInterval =
      this.options.pollInterval ?? DEFAULT_POLL_INTERVAL;

    do {
      const roots = await this.driver.getViewHierarchy();
      const matches = queryAll(roots, this.strategy);
      if (matches.length > 0) {
        return matches[0];
      }
      if (timeout <= 0) {
        return null;
      }
      await sleep(pollInterval);
    } while (Date.now() < deadline);

    return null;
  }
}

async function cropToElement(
  screenshot: Buffer,
  bounds: Bounds,
  screenSize: { width: number; height: number },
): Promise<Buffer> {
  const metadata = await sharp(screenshot).metadata();
  const scale = (metadata.width ?? 1) / screenSize.width;
  return sharp(screenshot)
    .extract({
      left: Math.round(bounds.x * scale),
      top: Math.round(bounds.y * scale),
      width: Math.round(bounds.width * scale),
      height: Math.round(bounds.height * scale),
    })
    .toBuffer();
}

function centerOf(bounds: Bounds): { x: number; y: number } {
  return {
    x: Math.round(bounds.x + bounds.width / 2),
    y: Math.round(bounds.y + bounds.height / 2),
  };
}

function boundsEqual(a: Bounds, b: Bounds): boolean {
  return (
    a.x === b.x &&
    a.y === b.y &&
    a.width === b.width &&
    a.height === b.height
  );
}

function checkState(
  node: ViewNode | null,
  state: 'visible' | 'hidden' | 'enabled' | 'disabled',
): boolean {
  switch (state) {
    case 'visible':
      return node !== null && node.isVisible;
    case 'hidden':
      return node === null || !node.isVisible;
    case 'enabled':
      return node !== null && node.isEnabled;
    case 'disabled':
      return node !== null && !node.isEnabled;
  }
}

function isWithinViewport(bounds: Bounds, screen: ScreenSize): boolean {
  return bounds.y >= 0
    && bounds.y + bounds.height <= screen.height
    && bounds.x >= 0
    && bounds.x + bounds.width <= screen.width;
}

function swipeDirectionToReveal(bounds: Bounds, screen: ScreenSize): SwipeDirection {
  const centerY = bounds.y + bounds.height / 2;
  // Element is below the viewport → swipe up to reveal it
  if (centerY > screen.height) {
    return 'up';
  }
  // Element is above the viewport → swipe down to reveal it
  return 'down';
}

export class LocatorError extends Error {
  constructor(
    message: string,
    public readonly strategy: LocatorStrategy,
  ) {
    super(message);
    this.name = 'LocatorError';
  }
}