// 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 Host from '../../core/host/host.js';
import * as SDK from '../../core/sdk/sdk.js';

import {
  AiAgent,
  type AllowedOriginResult,
  type ContextResponse,
  type ConversationContext,
  type MultimodalInputType,
  type RequestOptions,
  ResponseType
} from './agents/AiAgent.js';
import {type ExecuteJsAgentOptions, executeJsCode} from './agents/ExecuteJavascript.js';
import {ChangeManager} from './ChangeManager.js';
import {DOMNodeContext} from './contexts/DOMNodeContext.js';
import {debugLog} from './debug.js';
import {ExtensionScope} from './ExtensionScope.js';
import type {Skill, SkillName} from './skills/Skill.js';
import {SKILLS} from './skills/SkillRegistry.js';
import type {AllToolsContext, Tool, ToolArgs} from './tools/Tool.js';
import {ToolRegistry} from './tools/ToolRegistry.js';

const SKILL_DISPLAY_NAMES: Record<SkillName, string> = {
  styling: 'CSS and styling',
};

export class AiAgent2 extends AiAgent<unknown> {
  // TODO: The static preamble is a placeholder and will eventually live server-side.
  readonly preamble = 'You are a unified AI assistant in Chrome DevTools. You can learn skills to help the user.';
  readonly clientFeature = Host.AidaClient.ClientFeature.CHROME_STYLING_AGENT;  // Placeholder
  readonly userTier = 'TESTERS';

  #skillsInjected = false;
  #changes = new ChangeManager();
  #execJs: typeof executeJsCode;
  readonly #allowedOrigin?: () => AllowedOriginResult;

  get options(): RequestOptions {
    return {};
  }

  readonly #activeSkills = new Set<SkillName>();
  readonly #declaredTools = new Set<string>();

  constructor(opts: ExecuteJsAgentOptions) {
    super(opts);
    this.#execJs = opts.execJs ?? executeJsCode;
    this.#allowedOrigin = opts.allowedOrigin;
    this.#declaredTools.add('learnSkills');
    const skillsList = Object.keys(SKILLS).join(', ');
    this.declareFunction<{skills: SkillName[]}>('learnSkills', {
      description: `Load skills to help with the task. Available skills: ${skillsList}.`,
      parameters: {
        type: Host.AidaClient.ParametersTypes.OBJECT,
        description: 'Parameters for learning skills',
        properties: {
          skills: {
            type: Host.AidaClient.ParametersTypes.ARRAY,
            items: {
              type: Host.AidaClient.ParametersTypes.STRING,
              description: 'Skill name',
            },
            description: 'List of skill names to load',
          },
        },
        required: ['skills'],
      },
      displayInfoFromArgs: args => {
        const isSingular = args.skills.length === 1;
        const prefix = isSingular ? 'Learning skill' : 'Learning skills';
        const names = args.skills.map(name => SKILL_DISPLAY_NAMES[name] ?? name).join(', ');
        return {
          title: `${prefix}: ${names}`,
          action: `learnSkills(${args.skills.map(name => `'${name}'`).join(', ')})`,
        };
      },
      handler: async args => {
        const result = await this.learnSkill(args.skills);
        return {result};
      },
    });
  }

  override async enhanceQuery(
      query: string,
      selected: ConversationContext<unknown>|null = null,
      // TODO: support multimodal input in AiAgent2.
      _multimodalInputType?: MultimodalInputType,
      ): Promise<string> {
    let enhancedQuery = query;
    if (selected) {
      const promptDetails = await selected.getPromptDetails();
      if (promptDetails) {
        enhancedQuery = `${promptDetails}

# User request

QUERY: ${query}`;
      }
    }

    if (this.#skillsInjected) {
      return enhancedQuery;
    }
    this.#skillsInjected = true;
    const skillsManifest =
        Object.entries(this.getSkills()).map(([name, skill]) => `- ${name}: ${skill.description}`).join('\n');
    return `Available skills:
${skillsManifest}

You must call \`learnSkills\` to load a skill before you can use it.

User query: ${enhancedQuery}`;
  }

  override async *
      handleContextDetails(selected: ConversationContext<unknown>|null): AsyncGenerator<ContextResponse, void, void> {
    if (selected) {
      const details = await selected.getUserFacingDetails();
      if (details) {
        yield {
          type: ResponseType.CONTEXT,
          details,
        };
      }
    }
  }

  getSkills(): Record<SkillName, Skill> {
    return SKILLS;
  }

  async learnSkill(names: SkillName[]): Promise<string> {
    let response = '';
    const skills = this.getSkills();
    for (const name of names) {
      debugLog(`AiAgent2: Attempting to load skill ${name}`);
      if (this.#activeSkills.has(name)) {
        debugLog(`AiAgent2: Skill ${name} is already loaded`);
        response += `Skill ${name} is already loaded.\n`;
        continue;
      }

      const skillObj: Skill = skills[name];
      if (skillObj) {
        this.#activeSkills.add(name);
        debugLog(`AiAgent2: Skill ${name} loaded successfully`);
        response += `Skill ${name} loaded. Instructions:\n${skillObj.instructions}\n`;
        for (const toolName of skillObj.allowedTools) {
          const tool = ToolRegistry.get(toolName);
          if (tool) {
            this.#declareTool(tool);
          }
        }
      } else {
        debugLog(`AiAgent2: Failed to load skill ${name}`);
        response += `Failed to load skill ${name}. Valid skills are: ${Object.keys(skills).join(', ')}.\n`;
      }
    }
    return response.trim();
  }

  #createExtensionScope(changes: ChangeManager): {install(): Promise<void>, uninstall(): Promise<void>} {
    const selectedNode = this.context && this.context instanceof DOMNodeContext ? this.context.getItem() : null;
    return new ExtensionScope(changes, this.sessionId, selectedNode);
  }

  /**
   * Declares a tool to be available to the agent model, verifying first that
   * it hasn't already been declared to prevent duplicate declaration errors.
   */
  #declareTool(tool: Tool<ToolArgs, unknown, AllToolsContext>): void {
    if (this.#declaredTools.has(tool.name)) {
      debugLog(`AiAgent2: Tool ${tool.name} is already declared`);
      return;
    }
    this.#declaredTools.add(tool.name);
    this.declareFunction(tool.name, {
      description: tool.description,
      parameters: tool.parameters,
      displayInfoFromArgs: tool.displayInfoFromArgs,
      handler: (args, options) => {
        const context: AllToolsContext = {
          conversationContext: this.context ?? null,
          changeManager: this.#changes,
          createExtensionScope: this.#createExtensionScope.bind(this),
          execJs: this.#execJs,
          getExecutionContextNode: () => this.context instanceof DOMNodeContext ? this.context.getItem() : null,
          getTarget: () => SDK.TargetManager.TargetManager.instance().primaryPageTarget(),
          getEstablishedOrigin: () => this.#getConversationOrigin(),
        };
        return tool.handler(args, context, options);
      },
    });
  }

  #getConversationOrigin(): string|undefined {
    const allowed = this.#allowedOrigin?.();
    return allowed && 'origin' in allowed ? allowed.origin : undefined;
  }

  get activeSkills(): Set<SkillName> {
    return this.#activeSkills;
  }
}