// Copyright 2026 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

import * as Common from '../../../core/common/common.js';
import * as Host from '../../../core/host/host.js';
import * as i18n from '../../../core/i18n/i18n.js';
import * as SDK from '../../../core/sdk/sdk.js';
import type {FunctionHandlerOptions} from '../agents/AiAgent.js';
import {areOriginsEquivalent, extractContextOrigin, isOpaqueOrigin} from '../AiOrigins.js';

import {MAX_TARGET_ORIGINS, resolveDOMStorages} from './DOMStorageUtils.js';
import {
  type BaseToolCapability,
  type DataHandlerResult,
  type DataTool,
  type OriginLockCapability,
  type ServerLoggingCapability,
  ToolAnnotation,
  type ToolArgs,
  ToolName,
} from './Tool.js';

const lockedString = i18n.i18n.lockedString;

// Maximum character length allowed per storage value to prevent large values (e.g. huge JSON blobs)
// from exceeding the LLM context window.
export const MAX_NUM_CHAR_LENGTH = 10000;

export interface GetStorageValuesArgs extends ToolArgs {
  type: 'localStorage'|'sessionStorage';
  keys: string[];
  origins: string[];
  storageKey?: string;
}

export interface GetStorageValuesResult {
  storageValuesByOrigin: Record<string, {
    items: Array<{
      storageKey: string,
      values: Record<string, string>,
    }>,
  }>;
}

export class GetStorageValuesTool implements DataTool<GetStorageValuesArgs, GetStorageValuesResult,
                                                      BaseToolCapability&OriginLockCapability&ServerLoggingCapability> {
  readonly name = ToolName.GET_STORAGE_VALUES;
  readonly description = 'Retrieve specific string values from storage partitions for requested keys across origins.';

  readonly annotations = [ToolAnnotation.REDACT_FROM_HISTORY];

  readonly parameters: Host.AidaClient.FunctionObjectParam<keyof GetStorageValuesArgs> = {
    type: Host.AidaClient.ParametersTypes.OBJECT,
    description: '',
    nullable: false,
    properties: {
      type: {
        type: Host.AidaClient.ParametersTypes.STRING,
        description: 'Storage type: localStorage or sessionStorage',
        nullable: false,
      },
      keys: {
        type: Host.AidaClient.ParametersTypes.ARRAY,
        description: 'A list of keys to retrieve values for.',
        items: {type: Host.AidaClient.ParametersTypes.STRING, description: 'A storage key.'},
        nullable: false,
      },
      origins: {
        type: Host.AidaClient.ParametersTypes.ARRAY,
        description: 'List of origins to get values for.',
        items: {type: Host.AidaClient.ParametersTypes.STRING, description: 'An origin URL.'},
        nullable: false,
      },
      storageKey: {
        type: Host.AidaClient.ParametersTypes.STRING,
        description:
            'Optional. Specific storageKey partition to get values for. Only applies if single origin is provided.',
        nullable: true,
      },
    },
    required: ['type', 'keys', 'origins'],
  };

  displayInfoFromArgs(args: GetStorageValuesArgs): {
    title: string,
    action: string,
  } {
    return {
      title: lockedString('Reading storage values'),
      action: `getStorageValues('${args.type}', ${JSON.stringify(args.keys)}, ${JSON.stringify(args.origins)})`,
    };
  }

  async handler(
      args: GetStorageValuesArgs,
      context: BaseToolCapability&OriginLockCapability&ServerLoggingCapability,
      options?: FunctionHandlerOptions,
      ): Promise<DataHandlerResult<GetStorageValuesResult>> {
    context.setLoggingEnabled(false);

    // eslint-disable-next-line @devtools/no-instance-of-migrated-singletons
    const targetManager = SDK.TargetManager.TargetManager.instance();
    const primaryPageTarget = targetManager.primaryPageTarget();

    const allowedOrigin = context.getEstablishedOrigin();
    if (!allowedOrigin || isOpaqueOrigin(allowedOrigin)) {
      return {error: 'No origin available or not allowed.'};
    }

    if (!primaryPageTarget) {
      return {error: 'No origin available or not allowed.'};
    }

    const pageOrigin = Common.ParsedURL.ParsedURL.extractOrigin(primaryPageTarget.inspectedURL());
    if (!pageOrigin || !areOriginsEquivalent(pageOrigin, allowedOrigin)) {
      return {error: 'No origin available or not allowed.'};
    }

    const rawList = (args.origins && args.origins.length > 0) ? args.origins : [allowedOrigin];
    const validOrigins = rawList.map(origin => extractContextOrigin(origin))
                             .filter(origin => areOriginsEquivalent(origin, allowedOrigin));
    const targetOrigins = Array.from(new Set(validOrigins)).slice(0, MAX_TARGET_ORIGINS);
    if (targetOrigins.length === 0) {
      return {error: 'No valid origins found.'};
    }

    const storageKey = (targetOrigins.length === 1 && args.storageKey) ? args.storageKey : undefined;

    const allStoragesMap: Record<string, SDK.DOMStorageModel.DOMStorage[]> = {};
    let totalStoragesCount = 0;
    for (const origin of targetOrigins) {
      const storages = resolveDOMStorages(origin, args.type, targetManager, primaryPageTarget, storageKey);
      if (storages.length > 0) {
        allStoragesMap[origin] = storages;
        totalStoragesCount += storages.length;
      }
    }

    if (totalStoragesCount === 0) {
      return {error: 'No matching storage partitions found.'};
    }

    if (options?.approved !== true) {
      const keyString = args.keys.map(k => `\`${k}\``).join(', ');
      const targetsDesc = Object.keys(allStoragesMap).join(', ');

      return {
        requiresApproval: true,
        description:
            lockedString(`The AI wants to access the value(s) of ${args.type} keys ${keyString} on ${targetsDesc}.`),
      };
    }

    const storageValuesByOrigin: GetStorageValuesResult['storageValuesByOrigin'] = {};

    await Promise.all(targetOrigins.map(async origin => {
      const storages = allStoragesMap[origin] || [];
      const itemsResult = [];

      const keyAndItems = await Promise.all(storages.map(async storage => {
        // Silently catch CDP errors (e.g. target navigation, detachment, or frame crash)
        // so that transient target loss does not fail the entire batch.
        const items = await storage.getItems().catch(() => null);
        return {storageKey: storage.storageKey, items};
      }));

      for (const {storageKey: partitionKey, items} of keyAndItems) {
        if (!items || !partitionKey) {
          continue;
        }
        const itemMap = new Map<string, string>();
        for (const [key, value] of items) {
          itemMap.set(key, value);
        }
        const storageValues: Record<string, string> = {};

        for (const key of args.keys) {
          const value = itemMap.get(key);
          if (value === undefined) {
            continue;
          }
          const truncatedValue =
              value.length > MAX_NUM_CHAR_LENGTH ? value.substring(0, MAX_NUM_CHAR_LENGTH) + '... <truncated>' : value;
          storageValues[key] = truncatedValue;
        }

        itemsResult.push({storageKey: partitionKey, values: storageValues});
      }
      storageValuesByOrigin[origin] = {items: itemsResult};
    }));

    return {result: {storageValuesByOrigin}};
  }
}