import { inject, injectable } from "@codemation/core";
import { ApplicationTokens } from "../applicationTokens";
import type { LoggerFactory } from "../application/logging/Logger";
import { McpServerCatalog } from "./McpServerCatalog";
import { DefaultMcpClientFactory } from "./McpClientFactory";
import type { McpClientFactory } from "./McpClientFactory";
import { CredentialOAuth2MaterialReader } from "../credentials/CredentialOAuth2MaterialReader";
import type { MCPClient, McpToolSet } from "./McpConnectionPool.types";

type MutablePoolEntry = {
  client: MCPClient;
  toolsCache: McpToolSet | null;
  openedAt: Date;
};

@injectable()
export class McpConnectionPool {
  private readonly pool = new Map<string, MutablePoolEntry>();
  private readonly inFlight = new Map<string, Promise<MutablePoolEntry>>();

  constructor(
    @inject(McpServerCatalog) private readonly catalog: McpServerCatalog,
    @inject(CredentialOAuth2MaterialReader) private readonly oauth2Material: CredentialOAuth2MaterialReader,
    @inject(ApplicationTokens.LoggerFactory) private readonly loggers: LoggerFactory,
    @inject(DefaultMcpClientFactory) private readonly clientFactory: McpClientFactory,
  ) {}

  async getClient(credentialInstanceId: string, serverId: string): Promise<MCPClient> {
    const entry = await this.getOrOpenEntry(credentialInstanceId, serverId);
    return entry.client;
  }

  async getTools(credentialInstanceId: string, serverId: string): Promise<McpToolSet> {
    const entry = await this.getOrOpenEntry(credentialInstanceId, serverId);
    if (!entry.toolsCache) {
      const raw = await entry.client.tools();
      const decl = this.catalog.get(serverId);
      entry.toolsCache = this.applyOverrides(raw, decl?.toolDescriptionOverrides);
    }
    return entry.toolsCache;
  }

  async closeForCredential(credentialInstanceId: string): Promise<void> {
    const logger = this.loggers.create("McpConnectionPool");
    const prefix = `${credentialInstanceId}:`;
    const toClose: Array<[string, MutablePoolEntry]> = [];
    for (const [key, entry] of this.pool.entries()) {
      if (key.startsWith(prefix)) {
        toClose.push([key, entry]);
        this.pool.delete(key);
        logger.info(`McpConnectionPool: closed pool entry on credential revocation (key=${key})`);
      }
    }
    await Promise.allSettled(
      toClose.map(([key, entry]) =>
        entry.client.close().catch((e: unknown) => {
          logger.warn(
            `McpConnectionPool: error closing client on credential revocation (key=${key})`,
            e instanceof Error ? e : undefined,
          );
        }),
      ),
    );
  }

  async closeAll(): Promise<void> {
    await Promise.allSettled([...this.pool.values()].map((e) => e.client.close()));
    this.pool.clear();
    this.inFlight.clear();
  }

  private async getOrOpenEntry(credentialInstanceId: string, serverId: string): Promise<MutablePoolEntry> {
    const key = this.poolKey(credentialInstanceId, serverId);
    const cached = this.pool.get(key);
    if (cached) {
      return cached;
    }
    const existing = this.inFlight.get(key);
    if (existing) {
      return existing;
    }
    const openPromise = this.open(credentialInstanceId, serverId, key).finally(() => {
      this.inFlight.delete(key);
    });
    this.inFlight.set(key, openPromise);
    return openPromise;
  }

  private async open(credentialInstanceId: string, serverId: string, key: string): Promise<MutablePoolEntry> {
    const decl = this.catalog.get(serverId);
    if (!decl) {
      throw new Error(`McpConnectionPool: MCP server "${serverId}" not found in catalog`);
    }
    if (decl.transport !== "http") {
      throw new Error(
        `McpConnectionPool: MCP server "${serverId}" uses transport "${decl.transport}" which is not allowed in managed mode. ` +
          `Only "http" transport is supported. For stdio, set CODEMATION_ALLOW_STDIO_MCP=true in a self-hosted environment.`,
      );
    }

    const accessToken = await this.readAccessToken(credentialInstanceId, serverId);
    const headers: Record<string, string> = {
      ...(decl.staticHeaders ?? {}),
      authorization: `Bearer ${accessToken}`,
    };

    const client = await this.clientFactory.open({ url: decl.url, headers });
    const entry: MutablePoolEntry = { client, toolsCache: null, openedAt: new Date() };
    this.pool.set(key, entry);
    return entry;
  }

  private poolKey(credentialInstanceId: string, serverId: string): string {
    return `${credentialInstanceId}:${serverId}`;
  }

  private async readAccessToken(credentialInstanceId: string, serverId: string): Promise<string> {
    const material = await this.oauth2Material.readMaterial(credentialInstanceId);
    if (!material.accessToken) {
      throw new Error(
        `McpConnectionPool: credential instance "${credentialInstanceId}" has no access token — reconnect the credential bound to MCP server "${serverId}"`,
      );
    }
    return material.accessToken;
  }

  private applyOverrides(tools: McpToolSet, overrides?: Record<string, string>): McpToolSet {
    if (!overrides) {
      return tools;
    }
    const result = { ...tools };
    for (const [name, description] of Object.entries(overrides)) {
      if (result[name]) {
        result[name] = { ...result[name], description };
      }
    }
    return result;
  }
}