import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";

type GatewayDescribeFileConfig = {
  endpoint: string;
};

type GatewayConfig = {
  baseUrl: string;
  models: Array<{ id: string; name: string }>;
  tools?: {
    describe_file?: GatewayDescribeFileConfig;
  };
};

type GatewayProvider = {
  id: string;
  name: string;
  gateway: string;
  apiKey?: string;
  config: GatewayConfig;
  authenticatedConfig?: GatewayConfig;
};

type RuntimeState = {
  runtimeReady: boolean;
};

type DescribeFileToolInput = {
  input_files: Array<{
    mediaType: string;
    data: string;
  }>;
  prompt?: string;
};

type GatewayDescribeFileResponse = {
  kind: "radius#describe_file";
  model: string;
  prompt: string;
  description: string;
};

const DEFAULT_GATEWAY = "https://radius.pi.dev";
const DEFAULT_DEV_GATEWAY = "http://localhost:8788";
const GATEWAY_PROVIDER_ID = "radius";
const DEV_GATEWAY_PROVIDER_ID = "radius-dev";
const GATEWAY_API_KEY_ENV = "PI_GATEWAY_API_KEY";
const DEV_GATEWAY_API_KEY_ENV = "PI_DEV_GATEWAY_API_KEY";
const DESCRIBE_FILE_TOOL_NAME = "describe_file";

export default async function (pi: ExtensionAPI) {
  const state: RuntimeState = { runtimeReady: false };
  const providers: GatewayProvider[] = [
    {
      id: GATEWAY_PROVIDER_ID,
      name: "Radius",
      gateway: process.env.PI_GATEWAY || DEFAULT_GATEWAY,
      apiKey: process.env[GATEWAY_API_KEY_ENV],
      config: createInitialGatewayConfig(process.env.PI_GATEWAY || DEFAULT_GATEWAY),
    },
  ];

  const devGateway = process.env.PI_DEV_GATEWAY || DEFAULT_DEV_GATEWAY;
  if (devGateway) {
    providers.push({
      id: DEV_GATEWAY_PROVIDER_ID,
      name: "Radius (dev)",
      gateway: devGateway,
      apiKey: process.env[DEV_GATEWAY_API_KEY_ENV],
      config: createInitialGatewayConfig(devGateway),
    });
  }

  pi.on("session_start", () => {
    state.runtimeReady = true;
    refreshTool(pi, providers, state);
  });

  pi.on("model_select", () => {
    refreshTool(pi, providers, state);
  });

  refreshTool(pi, providers, state);
}

function refreshTool(pi: ExtensionAPI, providers: GatewayProvider[], state: RuntimeState): void {
  registerDescribeFileTool(pi, providers);
  if (state.runtimeReady) {
    ensureToolActive(pi);
  }
}

function registerDescribeFileTool(pi: ExtensionAPI, providers: GatewayProvider[]): void {
  pi.registerTool({
    name: DESCRIBE_FILE_TOOL_NAME,
    label: "Describe File",
    description:
      "Describe an attached image file through Radius using Gemini 2.5 Flash Image via OpenRouter. Returns text for the active loop.",
    promptSnippet: "Describe an attached image file through Radius and return text to continue the loop",
    promptGuidelines: [
      "Use describe_file when the current model cannot inspect an attached image directly.",
      "Pass through the exact attached image bytes when available.",
      "Use the returned description as context, then continue the task normally.",
    ],
    parameters: Type.Object(
      {
        input_files: Type.Array(
          Type.Object(
            {
              mediaType: Type.String(),
              data: Type.String(),
            },
            { additionalProperties: false },
          ),
          { minItems: 1 },
        ),
        prompt: Type.Optional(Type.String({ description: "Optional describe prompt override" })),
      },
      { additionalProperties: false },
    ),
    async execute(_toolCallId, params, signal, onUpdate, ctx) {
      const provider = selectProvider(providers, ctx);
      const config = getEffectiveConfig(provider).tools?.describe_file;
      if (!config) {
        throw new Error("Radius describe_file is not configured");
      }
      onUpdate?.({
        content: [{ type: "text", text: `Describing ${params.input_files.length} file(s)…` }],
      });
      const apiKey = await getApiKey(provider, ctx);
      const response = await fetch(config.endpoint, {
        method: "POST",
        headers: {
          authorization: `Bearer ${apiKey}`,
          accept: "application/json",
          "content-type": "application/json",
        },
        body: JSON.stringify({
          input_files: params.input_files,
          prompt: params.prompt,
        }),
        signal,
      });
      if (!response.ok) {
        throw new Error(`${response.status} ${response.statusText}: ${await response.text()}`);
      }
      const payload = (await response.json()) as GatewayDescribeFileResponse;
      return {
        content: [{ type: "text", text: payload.description }],
        details: {
          provider: provider.id,
          model: payload.model,
          prompt: payload.prompt,
        },
      };
    },
    renderCall(args: DescribeFileToolInput, theme: any, context: any) {
      const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
      text.setText(
        `${theme.fg("toolTitle", theme.bold(`${DESCRIBE_FILE_TOOL_NAME} `))}${theme.fg("accent", `${args.input_files?.length ?? 0} file(s)`)}`,
      );
      return text;
    },
    renderResult(result: any, { isPartial }: any, theme: any, context: any) {
      const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
      if (isPartial) {
        text.setText(theme.fg("warning", "Describing file…"));
        return text;
      }
      const contentText = result.content?.find?.((item: any) => item.type === "text")?.text;
      text.setText(contentText || theme.fg("error", "No description returned"));
      return text;
    },
  });
}

function ensureToolActive(pi: ExtensionAPI): void {
  const activeTools = pi.getActiveTools();
  if (!activeTools.includes(DESCRIBE_FILE_TOOL_NAME)) {
    pi.setActiveTools([...activeTools, DESCRIBE_FILE_TOOL_NAME]);
  }
}

function selectProvider(providers: GatewayProvider[], ctx: any): GatewayProvider {
  const byModel = providers.find((candidate) => candidate.id === ctx.model?.provider);
  if (byModel) {
    return byModel;
  }
  const prod = providers.find((candidate) => candidate.id === GATEWAY_PROVIDER_ID);
  if (prod) {
    return prod;
  }
  throw new Error("No Radius provider configured for describe_file");
}

async function getApiKey(provider: GatewayProvider, ctx: any): Promise<string> {
  const envKey = getGatewayApiKey(provider.id) ?? provider.apiKey;
  if (envKey) {
    return envKey;
  }

  const providerKey = await ctx.modelRegistry.getApiKeyForProvider?.(provider.id);
  if (providerKey) {
    return providerKey;
  }

  const model =
    ctx.model?.provider === provider.id ? ctx.model : findAnyGatewayModel(provider, ctx.modelRegistry);
  if (model) {
    const result = await ctx.modelRegistry.getApiKeyAndHeaders(model);
    if (result.ok && result.apiKey) {
      return result.apiKey;
    }
  }

  throw new Error(
    `Missing gateway API key for ${provider.id}: run /login for Radius, or set ${GATEWAY_API_KEY_ENV} (or ${DEV_GATEWAY_API_KEY_ENV} for the dev gateway)`,
  );
}

function findAnyGatewayModel(provider: GatewayProvider, modelRegistry: any): unknown {
  for (const config of [provider.authenticatedConfig, provider.config]) {
    for (const model of config?.models ?? []) {
      const found = modelRegistry.find?.(provider.id, getGatewayDisplayModelId(model.id));
      if (found) {
        return found;
      }
    }
  }
  return undefined;
}

function getGatewayDisplayModelId(modelId: string): string {
  return modelId.startsWith("byok/") ? modelId.slice("byok/".length) : modelId;
}

function getGatewayApiKey(provider: string): string | undefined {
  if (provider === DEV_GATEWAY_PROVIDER_ID) {
    return process.env[DEV_GATEWAY_API_KEY_ENV];
  }
  if (provider === GATEWAY_PROVIDER_ID) {
    return process.env[GATEWAY_API_KEY_ENV];
  }
  return undefined;
}

function getEffectiveConfig(provider?: GatewayProvider): GatewayConfig {
  return (
    provider?.authenticatedConfig ?? provider?.config ?? createInitialGatewayConfig(DEFAULT_GATEWAY)
  );
}

function createInitialGatewayConfig(gateway: string): GatewayConfig {
  return {
    baseUrl: `${gateway.replace(/\/+$/u, "")}/v1`,
    models: [],
    tools: {
      describe_file: {
        endpoint: `${gateway.replace(/\/+$/u, "")}/v1/describe-file`,
      },
    },
  };
}