import { z } from 'zod';
import { tool, DynamicStructuredTool } from '@langchain/core/tools';
import type * as _t from '@/types';

/**
 * Browser tool names - keep in sync with the browser extension
 * These tools execute locally in the browser extension, NOT on the server
 */
export const EBrowserTools = {
  CLICK: 'browser_click',
  TYPE: 'browser_type',
  NAVIGATE: 'browser_navigate',
  SCROLL: 'browser_scroll',
  EXTRACT: 'browser_extract',
  HOVER: 'browser_hover',
  WAIT: 'browser_wait',
  BACK: 'browser_back',
  SCREENSHOT: 'browser_screenshot',
  GET_PAGE_STATE: 'browser_get_page_state',
  KEYPRESS: 'browser_keypress',
  SWITCH_TAB: 'browser_switch_tab',
} as const;

export type BrowserToolName =
  (typeof EBrowserTools)[keyof typeof EBrowserTools];

/**
 * Callback function type for waiting on browser action results
 * This allows the host server to provide a callback that waits for the extension
 * to POST results back to the server before returning to the LLM.
 *
 * @param action - The browser action (click, type, navigate, etc.)
 * @param args - Arguments for the action
 * @param toolCallId - Unique ID for this tool call (from config.toolCall.id)
 * @returns Promise that resolves with the actual browser result (page state, etc.)
 */
export type BrowserToolCallback = (
  action: string,
  args: Record<string, unknown>,
  toolCallId: string
) => Promise<BrowserActionResult>;

/**
 * Result returned from browser action execution
 */
export interface BrowserActionResult {
  success: boolean;
  url?: string;
  title?: string;
  elementList?: string; // Text-based element list
  error?: string;
  screenshot?: string; // Base64 screenshot (if requested)
}

/**
 * Check if browser capability is available based on request headers or context
 * The browser extension sets these headers when connected:
 * - X-Illuma-Browser-Extension: true
 * - X-Illuma-Browser-Capable: true
 */
export function hasBrowserCapability(req?: {
  headers?: Record<string, string | string[] | undefined>;
}): boolean {
  if (!req?.headers) {
    return false;
  }

  const browserExtension = req.headers['x-illuma-browser-extension'];
  const browserCapable = req.headers['x-illuma-browser-capable'];

  return browserExtension === 'true' || browserCapable === 'true';
}

// Tool schemas
const BrowserClickSchema = z.object({
  index: z
    .number()
    .describe(
      'The [index] of the element to click. Use fieldLabel to identify the correct element. For form fields, target <input> or <textarea> elements, NOT parent <div> containers.'
    ),
  text: z
    .string()
    .optional()
    .describe(
      'The visible text of the element being clicked (e.g., "Send", "Delete", "Buy Now"). Always include this for buttons and links.'
    ),
  label: z
    .string()
    .optional()
    .describe('The fieldLabel or ariaLabel of the element, if available.'),
});

const BrowserTypeSchema = z.object({
  index: z
    .number()
    .describe(
      'The [index] of the INPUT element to type into. Target <input> or <textarea> elements. Check fieldLabel to identify the correct field.'
    ),
  text: z.string().describe('The text to type into the element'),
  pressEnter: z
    .boolean()
    .optional()
    .describe(
      'Whether to press Enter after typing (useful for search forms and submitting)'
    ),
});

const BrowserNavigateSchema = z.object({
  url: z
    .string()
    .describe('The full URL to navigate to (must include https://)'),
});

const BrowserScrollSchema = z.object({
  direction: z
    .enum(['up', 'down', 'left', 'right'])
    .describe('Direction to scroll'),
  amount: z
    .number()
    .optional()
    .describe('Pixels to scroll (default: one viewport height)'),
});

const BrowserExtractSchema = z.object({
  query: z
    .string()
    .optional()
    .describe('Optional: specific content to extract from the page'),
});

const BrowserHoverSchema = z.object({
  index: z.number().describe('The index number of the element to hover over'),
});

const BrowserWaitSchema = z.object({
  duration: z
    .number()
    .optional()
    .describe('Milliseconds to wait (default: 1000)'),
});

const BrowserBackSchema = z.object({});

const BrowserScreenshotSchema = z.object({});

const BrowserGetPageStateSchema = z.object({});

const BrowserKeypressSchema = z.object({
  keys: z
    .string()
    .describe(
      'Keyboard keys to press. Use "+" to combine modifiers (e.g., "Control+Enter", "Control+a", "Escape", "Tab", "Enter"). Common shortcuts: Control+Enter (submit forms/send), Escape (close dialogs), Tab (next field).'
    ),
});

const BrowserSwitchTabSchema = z.object({
  tabId: z
    .number()
    .describe(
      'The tab ID to switch to. Use the tab IDs shown in the tabs list from page state.'
    ),
});

/**
 * Browser tool response interface
 * This is what the extension returns after executing the action
 */
export interface BrowserToolResponse {
  requiresBrowserExecution: true;
  action: string;
  args: Record<string, unknown>;
  toolCallId?: string; // Added to help extension correlate with callback
}

/**
 * Options for creating browser tools
 */
export interface CreateBrowserToolsOptions {
  /**
   * Optional callback that waits for browser action results.
   * When provided, tools will await this callback to get actual results from the extension.
   * When not provided, tools return markers immediately (for non-server contexts).
   */
  waitForResult?: BrowserToolCallback;
}

/**
 * Format browser action result for LLM consumption
 */
function formatResultForLLM(
  result: BrowserActionResult,
  action: string
): string {
  const parts: string[] = [];

  // Include error info but DON'T discard page state — the extension often sends
  // success=false with valid page state (URL, elementList) that the LLM needs
  if (!result.success && result.error) {
    parts.push(`**Note:** Action "${action}" reported: ${result.error}`);
  }

  if (result.url != null && result.url !== '') {
    parts.push(`**Current URL:** ${result.url}`);
  }
  if (result.title != null && result.title !== '') {
    parts.push(`**Page Title:** ${result.title}`);
  }
  if (result.elementList != null && result.elementList !== '') {
    parts.push(
      `\n**Interactive Elements** (for typing: target <input> elements with fieldLabel, NOT parent <div> containers):\n${result.elementList}`
    );
  }
  if (result.screenshot != null && result.screenshot !== '') {
    parts.push('\n[Screenshot captured and displayed to user]');
  }

  if (parts.length === 0) {
    if (!result.success) {
      return `Browser action "${action}" failed: ${result.error || 'Unknown error'}`;
    }
    return `Browser action "${action}" completed successfully.`;
  }

  return parts.join('\n');
}

/**
 * Create browser tools with optional callback for waiting on results
 *
 * When waitForResult callback is provided:
 * 1. Tool returns marker that triggers extension
 * 2. Tool then awaits callback to get actual results
 * 3. Returns real page state to LLM
 *
 * When no callback:
 * 1. Tool returns marker only (for non-server contexts)
 *
 * NOTE: These tools use TEXT-BASED element lists, NOT screenshots
 * Screenshots would be 100K+ tokens each - element lists are ~100 tokens
 */
export function createBrowserTools(
  options?: CreateBrowserToolsOptions
): DynamicStructuredTool[] {
  const { waitForResult } = options || {};
  const tools: DynamicStructuredTool[] = [];

  /**
   * Helper to create tool function that optionally waits for results
   * The toolCallId is extracted from the RunnableConfig passed by LangChain
   */
  const createToolFunction = (action: string) => {
    return async (
      args: Record<string, unknown>,
      config?: { toolCall?: { id?: string } }
    ): Promise<string> => {
      const toolCallId =
        config?.toolCall?.id ??
        `tool_${Date.now()}_${Math.random().toString(36).slice(2)}`;

      // Create marker for extension
      const marker: BrowserToolResponse = {
        requiresBrowserExecution: true,
        action,
        args,
        toolCallId,
      };

      // If no callback, return marker immediately (extension handles via SSE interception)
      if (!waitForResult) {
        return JSON.stringify(marker);
      }

      // With callback: wait for actual results from extension
      // The marker is still returned initially via SSE, but we wait for the callback
      try {
        const result = await waitForResult(action, args, toolCallId);
        return formatResultForLLM(result, action);
      } catch (error) {
        const errorMessage =
          error instanceof Error ? error.message : String(error);
        return `Browser action "${action}" failed: ${errorMessage}`;
      }
    };
  };

  // browser_click
  tools.push(
    tool(createToolFunction('click'), {
      name: EBrowserTools.CLICK,
      description:
        'Click element by [index]. ALWAYS include the element text and label in your call for safety review. Use fieldLabel attribute to identify correct element. For form fields, target <input> elements NOT parent <div> containers.',
      schema: BrowserClickSchema,
    })
  );

  // browser_type
  tools.push(
    tool(createToolFunction('type'), {
      name: EBrowserTools.TYPE,
      description:
        'Type text into <input> element by [index]. CRITICAL: Always target <input> or <textarea> tags (NOT parent <div> containers). Use fieldLabel to identify correct field (e.g., fieldLabel="To recipients" for To field).',
      schema: BrowserTypeSchema,
    })
  );

  // browser_navigate
  tools.push(
    tool(createToolFunction('navigate'), {
      name: EBrowserTools.NAVIGATE,
      description:
        'Navigate to URL (include https://). Returns new page element list.',
      schema: BrowserNavigateSchema,
    })
  );

  // browser_scroll
  tools.push(
    tool(createToolFunction('scroll'), {
      name: EBrowserTools.SCROLL,
      description:
        'Scroll page (up/down/left/right). Returns updated element list.',
      schema: BrowserScrollSchema,
    })
  );

  // browser_extract
  tools.push(
    tool(createToolFunction('extract'), {
      name: EBrowserTools.EXTRACT,
      description:
        'Extract page content. Returns URL, title, and element list.',
      schema: BrowserExtractSchema,
    })
  );

  // browser_hover
  tools.push(
    tool(createToolFunction('hover'), {
      name: EBrowserTools.HOVER,
      description:
        'Hover element by [index] to reveal menus/tooltips. Returns updated element list.',
      schema: BrowserHoverSchema,
    })
  );

  // browser_wait
  tools.push(
    tool(createToolFunction('wait'), {
      name: EBrowserTools.WAIT,
      description:
        'Wait for async content to load. Returns updated element list.',
      schema: BrowserWaitSchema,
    })
  );

  // browser_back
  tools.push(
    tool(createToolFunction('back'), {
      name: EBrowserTools.BACK,
      description:
        'Go back in browser history. Returns previous page element list.',
      schema: BrowserBackSchema,
    })
  );

  // browser_screenshot
  tools.push(
    tool(createToolFunction('screenshot'), {
      name: EBrowserTools.SCREENSHOT,
      description:
        'Capture screenshot. Displayed to user. Use get_page_state for automation.',
      schema: BrowserScreenshotSchema,
    })
  );

  // browser_get_page_state
  tools.push(
    tool(createToolFunction('get_page_state'), {
      name: EBrowserTools.GET_PAGE_STATE,
      description:
        'Get page URL, title, and interactive elements with [index] for actions. Start here.',
      schema: BrowserGetPageStateSchema,
    })
  );

  // browser_keypress - for keyboard shortcuts
  tools.push(
    tool(createToolFunction('keypress'), {
      name: EBrowserTools.KEYPRESS,
      description:
        'Send keyboard shortcut or key press. Use for: Control+Enter (send email/submit), Escape (close dialog/cancel), Tab (next field), Enter (confirm). The keys are sent to the currently focused element.',
      schema: BrowserKeypressSchema,
    })
  );

  // browser_switch_tab - for switching between tabs
  tools.push(
    tool(createToolFunction('switch_tab'), {
      name: EBrowserTools.SWITCH_TAB,
      description:
        'Switch to a different browser tab by its ID. Tab IDs are shown in the page state. Use this to work with existing open tabs (e.g., use existing Gmail tab instead of opening a new one).',
      schema: BrowserSwitchTabSchema,
    })
  );

  return tools;
}