import type { Api, Model } from "@earendil-works/pi-ai";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

import {
  getAgentConfig,
  getConfig,
  getMemoryTools,
  getReadOnlyMemoryTools,
  getToolsForType,
} from "./agent-types.js";
import { detectEnv } from "./env.js";
import { buildMemoryBlock, buildReadOnlyMemoryBlock } from "./memory.js";
import { buildAgentPrompt, type PromptExtras } from "./prompts.js";
import { preloadSkills } from "./skill-loader.js";
import type { AgentConfig, SubagentType, ThinkingLevel } from "./types.js";

/** Names of tools registered by this extension that subagents must NOT inherit. */
export const EXCLUDED_TOOL_NAMES = [
  "Agent",
  "get_subagent_result",
  "steer_subagent",
  "reply_to_subagent",
  "get_subagent_message",
  "workflow",
];

/**
 * Try to find the right model for an agent type.
 * Priority: explicit option > config.model > parent model.
 */
export function resolveDefaultModel(
  parentModel: Model<Api> | undefined,
  registry: {
    find(provider: string, modelId: string): Model<Api> | undefined;
    getAvailable?(): Model<Api>[];
  },
  configModel?: string
): Model<Api> | undefined {
  if (configModel) {
    const slashIdx = configModel.indexOf("/");
    if (slashIdx !== -1) {
      const provider = configModel.slice(0, slashIdx);
      const modelId = configModel.slice(slashIdx + 1);

      // Build a set of available model keys for fast lookup
      const available = registry.getAvailable?.();
      const availableKeys = available
        ? new Set(available.map((m) => `${m.provider}/${m.id}`))
        : undefined;
      const isAvailable = (p: string, id: string) =>
        !availableKeys || availableKeys.has(`${p}/${id}`);

      const found = registry.find(provider, modelId);
      if (found && isAvailable(provider, modelId)) {
        return found;
      }
    }
  }

  return parentModel;
}

export interface SelectedAgentLaunchOptions {
  pi: ExtensionAPI;
  cwd?: string;
  parentCwd?: string;
  parentSystemPrompt?: string;
  parentModel?: Model<Api>;
  modelRegistry: {
    find(provider: string, modelId: string): Model<Api> | undefined;
    getAvailable?(): Model<Api>[];
  };
  model?: Model<Api>;
  maxTurns?: number;
  isolated?: boolean;
  thinkingLevel?: ThinkingLevel;
}

export interface SelectedAgentLaunchConfig {
  agentConfig: AgentConfig;
  effectiveCwd: string;
  systemPrompt: string;
  builtinTools: ReturnType<typeof getToolsForType>;
  builtinToolNames: string[];
  disallowedTools?: string[];
  excludedToolNames: string[];
  extensions: true | string[] | false;
  skills: true | string[] | false;
  noSkills: boolean;
  model?: Model<Api>;
  modelInput?: string;
  thinkingLevel?: ThinkingLevel;
  maxTurns?: number;
  explicitMaxTurns: boolean;
}

export async function buildSelectedAgentLaunchConfig(
  type: SubagentType,
  options: SelectedAgentLaunchOptions
): Promise<SelectedAgentLaunchConfig> {
  const agentConfig = getAgentConfig(type);
  if (!agentConfig || agentConfig.enabled === false) {
    throw new Error(
      `Agent type "${type}" is not defined. Create it in .pi/agents/${type}.md or ~/.pi/agent/agents/${type}.md.`
    );
  }
  const config = getConfig(type);
  const effectiveCwd = options.cwd ?? options.parentCwd ?? process.cwd();
  const env = await detectEnv(options.pi, effectiveCwd);
  const extras: PromptExtras = {};

  const extensions = options.isolated ? false : config.extensions;
  const skills = options.isolated ? false : config.skills;

  if (Array.isArray(skills)) {
    const loaded = preloadSkills(skills, effectiveCwd);
    if (loaded.length > 0) {
      extras.skillBlocks = loaded;
    }
  }

  let builtinTools = [...getToolsForType(type, effectiveCwd)];

  if (agentConfig.memory) {
    const existingNames = new Set(builtinTools.map((tool) => tool.name));
    const denied = agentConfig.disallowedTools
      ? new Set(agentConfig.disallowedTools)
      : undefined;
    const effectivelyHas = (name: string) =>
      existingNames.has(name) && !denied?.has(name);
    const hasWriteTools = effectivelyHas("write") || effectivelyHas("edit");

    if (hasWriteTools) {
      const memTools = getMemoryTools(effectiveCwd, existingNames);
      if (memTools.length > 0) {
        builtinTools = [...builtinTools, ...memTools];
      }
      extras.memoryBlock = buildMemoryBlock(
        agentConfig.name,
        agentConfig.memory,
        effectiveCwd
      );
    } else {
      if (!existingNames.has("read")) {
        const readTools = getReadOnlyMemoryTools(effectiveCwd, existingNames);
        if (readTools.length > 0) {
          builtinTools = [...builtinTools, ...readTools];
        }
      }
      extras.memoryBlock = buildReadOnlyMemoryBlock(
        agentConfig.name,
        agentConfig.memory,
        effectiveCwd
      );
    }
  }

  return {
    agentConfig,
    effectiveCwd,
    systemPrompt: buildAgentPrompt(
      agentConfig,
      effectiveCwd,
      env,
      options.parentSystemPrompt,
      extras
    ),
    builtinTools,
    builtinToolNames: builtinTools.map((tool) => tool.name),
    disallowedTools: agentConfig.disallowedTools,
    excludedToolNames: EXCLUDED_TOOL_NAMES,
    extensions,
    skills,
    noSkills: skills === false || Array.isArray(skills),
    model:
      options.model ??
      resolveDefaultModel(
        options.parentModel,
        options.modelRegistry,
        agentConfig.model
      ),
    modelInput: agentConfig.model,
    thinkingLevel: options.thinkingLevel ?? agentConfig.thinking,
    maxTurns: options.maxTurns ?? agentConfig.maxTurns,
    explicitMaxTurns: options.maxTurns != null || agentConfig.maxTurns != null,
  };
}