import { CatalogMCPServerItem } from "@mcpx/shared-model";
import {
  accessLogFor,
  buildLogger,
  loggableError,
  LunarLogger,
} from "@mcpx/toolkit-core/logging";
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
import { MeterProvider } from "@opentelemetry/sdk-metrics";
import express from "express";
import { createServer, Server } from "http";
import { v7 as uuidv7 } from "uuid";
import { ConfigService, InMemoryConfigStore } from "../src/config.js";
import { Config } from "../src/model/config/config.js";
import { TargetServer } from "../src/model/target-servers.js";
import { AuthGuard, noOpAuthGuard } from "../src/server/auth.js";
import { buildDownstreamTransportsRouter } from "../src/server/downstream-transports.js";
import { buildControlPlaneAppConfigRouter } from "../src/server/control-plane-app-config.js";
import { buildControlPlaneRouter } from "../src/server/control-plane.js";
import { Services } from "../src/services/services.js";
import {
  TESTKIT_SERVER_CALCULATOR,
  TESTKIT_SERVER_ECHO,
} from "../src/testkit/root.js";
import { MockHubServer, SetCatalogPayload } from "./mock-hub-server.js";
import {
  wrapInEnvelope,
  SetIdentityPayload,
} from "@mcpx/webapp-protocol/messages";
import { buildCatalogRouter } from "../src/server/catalog.js";
import { buildAdminRouter } from "../src/server/admin.js";
import { buildIdentityRouter } from "../src/server/identity.js";

const MCPX_PORT = 9000;
let nextHubPort = 3030; // Start from 3030 and increment for each harness

// Track all created loggers for cleanup
const allLoggers = new Set<LunarLogger>();

const getTestLogger: () => LunarLogger = () => {
  const logger = buildLogger({ logLevel: "debug", label: "test" });
  allLoggers.add(logger);
  return logger;
};

export const getMcpxLogger: () => LunarLogger = () => {
  const logger = buildLogger({ logLevel: "debug", label: "mcpx" });
  allLoggers.add(logger);
  return logger;
};

// Cleanup function to close all logger instances (for global teardown)
export const closeAllLoggers = (): void => {
  allLoggers.forEach((logger) => {
    try {
      logger.close();
    } catch (_e) {
      // Ignore errors when closing loggers
    }
  });
  allLoggers.clear();
};

const BASE_CONFIG: Config = {
  permissions: {
    default: { _type: "default-allow", block: [] },
    consumers: {},
    clientNames: {},
  },
  toolGroups: [],
  auth: { enabled: false },
  toolExtensions: { services: {} },
  targetServerAttributes: {},
};
export function buildConfig(props: Partial<Config> = {}): ConfigService {
  const {
    auth,
    permissions,
    toolGroups,
    toolExtensions,
    targetServerAttributes,
  } = {
    ...BASE_CONFIG,
    ...props,
  };
  return new ConfigService(
    {
      permissions,
      toolGroups,
      auth,
      toolExtensions,
      targetServerAttributes,
    },
    new InMemoryConfigStore(),
    getMcpxLogger(),
  );
}

// Runtime server configs (for actually running servers)
export const echoTargetServer: TargetServer = {
  type: "stdio",
  name: "echo-service",
  command: "node",
  args: [TESTKIT_SERVER_ECHO],
  env: {},
};

export const calculatorTargetServer: TargetServer = {
  type: "stdio",
  name: "calculator-service",
  command: "node",
  args: [TESTKIT_SERVER_CALCULATOR],
  env: {},
};

export const stdioTargetServers: TargetServer[] = [
  echoTargetServer,
  calculatorTargetServer,
];

export const oauthTargetServer: TargetServer = {
  type: "streamable-http",
  name: "oauth-mock-server",
  url: "http://localhost:9001/mcp",
};

export const allTargetServers: TargetServer[] = [
  ...stdioTargetServers,
  oauthTargetServer,
];

// Catalog items (templates for catalog testing)
export const echoCatalogItem: CatalogMCPServerItem = {
  id: uuidv7(),
  name: "echo-service",
  displayName: "Echo Service",
  config: {
    type: "stdio",
    command: "node",
    args: [TESTKIT_SERVER_ECHO],
  },
};

export const calculatorCatalogItem: CatalogMCPServerItem = {
  id: uuidv7(),
  name: "calculator-service",
  displayName: "Calculator Service",
  config: {
    type: "stdio",
    command: "node",
    args: [TESTKIT_SERVER_CALCULATOR],
  },
};

export const oauthCatalogItem: CatalogMCPServerItem = {
  id: uuidv7(),
  name: "oauth-mock-server",
  displayName: "OAuth Mock Server",
  config: {
    type: "streamable-http",
    url: "http://localhost:9001/mcp",
  },
};

export const stdioCatalogItems: CatalogMCPServerItem[] = [
  echoCatalogItem,
  calculatorCatalogItem,
];

export const allCatalogItems: CatalogMCPServerItem[] = [
  echoCatalogItem,
  calculatorCatalogItem,
  oauthCatalogItem,
];

export function catalogItemsToPayload(
  items: CatalogMCPServerItem[],
): SetCatalogPayload {
  return {
    items: items.map((item) => ({ server: item })),
  };
}
type TransportType = "SSE" | "StreamableHTTP";

export const transportTypes: TransportType[] = [
  "SSE",
  "StreamableHTTP",
] as const;

export class TestHarness {
  public clientConnectError?: Error | undefined;
  private loggers: LunarLogger[] = [];
  public readonly mockHubServer: MockHubServer;

  constructor(
    public client: Client,
    public server: Server,
    public services: Services,
    public testLogger: LunarLogger,
    mockHubServer: MockHubServer,
    private clientConnectExtraHeaders: Record<string, string> = {},
    private targetServers: TargetServer[] = stdioTargetServers,
  ) {
    // Track loggers for this harness instance
    this.loggers.push(testLogger);
    this.mockHubServer = mockHubServer;
  }

  addLogger(logger: LunarLogger): void {
    this.loggers.push(logger);
  }

  /** Emit identity to connected mcpx-server (strictness is derived from identity) */
  emitIdentity(identity: SetIdentityPayload): void {
    const envelope = wrapInEnvelope({ payload: identity });
    const socketId = this.mockHubServer.getConnectedClients()[0];
    this.mockHubServer.emitToClient(socketId!, "set-identity", envelope);
  }

  /** Emit catalog update to connected mcpx-server */
  emitCatalog(payload: SetCatalogPayload): void {
    const envelope = wrapInEnvelope({ payload });
    const socketId = this.mockHubServer.getConnectedClients()[0];
    this.mockHubServer.emitToClient(socketId!, "set-catalog", envelope);
  }

  async initialize(transportType: TransportType): Promise<void> {
    // Start mock Hub server first
    await this.mockHubServer.waitForListening();

    // Setup MCPX (Hub connection happens here)
    await this.services.initialize();

    // Wait for catalog from Hub before adding clients (enterprise mode)
    // Check if catalog already received, otherwise subscribe and trigger re-emit
    const catalogLength = this.services.catalogManager.getCatalog().length;
    this.testLogger.info("Catalog check after initialize", { catalogLength });
    if (catalogLength === 0) {
      this.testLogger.info("Catalog empty, waiting for Hub to send it");
      const catalogReceived = new Promise<void>((resolve, reject) => {
        const timeout = setTimeout(() => {
          reject(new Error("Timeout waiting for catalog from Hub"));
        }, 3000);
        const unsubscribe = this.services.catalogManager.subscribe(() => {
          clearTimeout(timeout);
          unsubscribe();
          resolve();
        });
      });
      this.mockHubServer.emitCatalogToAll();
      await catalogReceived;
      this.testLogger.info("Catalog received from Hub");
    }

    await Promise.all(
      this.targetServers.map((target) =>
        this.services.upstreamHandler.addClient(target),
      ),
    );
    await this.server.listen(MCPX_PORT, () => {
      this.testLogger.info(`Test MCPX server listening on port ${MCPX_PORT}`);
    });

    const transport = this.buildTransport(transportType);

    await this.client
      .connect(transport)
      .then(() => this.testLogger.info("End-client connected"))
      .catch((e) => {
        const error = loggableError(e);
        this.testLogger.error("Failed to connect end-client", error);
        this.clientConnectError = e;
      });
  }

  async shutdown(): Promise<void> {
    await this.client.transport?.close();
    await this.client.close();
    await this.services.shutdown();
    await new Promise<void>((resolve) => {
      // Force close all connections to speed up shutdown
      this.server.closeAllConnections();
      this.server.close(() => {
        this.testLogger.info("Test MCPX server closed");
        resolve();
      });
    });
    // Close mock Hub server before closing loggers (disconnect events need to log)
    await this.mockHubServer.close();
    // Close loggers after all servers/sockets are fully closed
    this.loggers.forEach((logger) => {
      try {
        logger.close();
      } catch (_e) {
        // Ignore errors when closing loggers
      }
    });
    this.loggers = [];
  }

  private buildTransport(transportType: TransportType): Transport {
    switch (transportType) {
      case "SSE":
        return new SSEClientTransport(
          new URL(`http://localhost:${MCPX_PORT}/sse`),
          {
            requestInit: {
              headers: this.clientConnectExtraHeaders,
            },
          },
        );
      case "StreamableHTTP":
        return new StreamableHTTPClientTransport(
          new URL(`http://localhost:${MCPX_PORT}/mcp`),
          {
            requestInit: {
              headers: this.clientConnectExtraHeaders,
            },
          },
        );
      default:
        throw new Error(`Unknown transport type: ${transportType}`);
    }
  }
}

interface TestHarnessProps {
  config?: ConfigService;
  authGuard?: AuthGuard;
  mcpxLogger?: LunarLogger;
  clientConnectExtraHeaders?: Record<string, string>;
  targetServers?: TargetServer[];
  catalogItems?: CatalogMCPServerItem[];
}
function defaultTestHarnessProps(): Required<TestHarnessProps> {
  return {
    config: buildConfig(),
    authGuard: noOpAuthGuard,
    mcpxLogger: getMcpxLogger(),
    clientConnectExtraHeaders: {},
    targetServers: stdioTargetServers,
    catalogItems: stdioCatalogItems,
  };
}
export function getTestHarness(props: TestHarnessProps = {}): TestHarness {
  // the effective values are the defaults, unless overridden by props
  const {
    config,
    authGuard,
    mcpxLogger,
    clientConnectExtraHeaders,
    targetServers,
    catalogItems,
  } = {
    ...defaultTestHarnessProps(),
    ...props,
  };
  const testLogger = getTestLogger();

  // Assign unique Hub port for this test harness
  const hubPort = nextHubPort++;

  const meterProvider = new MeterProvider();
  const services = new Services(config, meterProvider, testLogger, {
    hubUrl: `http://localhost:${hubPort}`,
  });

  const client = new Client({ name: "end-client", version: "1.0.0" });

  // Wire router on server
  const downstreamTransportsRouter = buildDownstreamTransportsRouter(
    authGuard,
    services,
    mcpxLogger,
  );
  const controlPlaneRouter = buildControlPlaneRouter(
    authGuard,
    services,
    mcpxLogger,
  );
  const controlPlaneAppConfigRouter = buildControlPlaneAppConfigRouter(
    authGuard,
    services,
    mcpxLogger,
  );

  const catalogRouter = buildCatalogRouter(authGuard, services, mcpxLogger);
  const adminRouter = buildAdminRouter(authGuard, services, mcpxLogger);
  const identityRouter = buildIdentityRouter(authGuard, services);

  const app = express();
  const httpServer = createServer(app);
  // Reduce keep-alive timeout to speed up test shutdown
  httpServer.keepAliveTimeout = 100;
  app.use(express.json());
  app.use(accessLogFor(mcpxLogger));
  app.use(downstreamTransportsRouter);
  app.use(controlPlaneRouter);
  app.use("/config", controlPlaneAppConfigRouter);
  app.use("/catalog", catalogRouter);
  app.use("/admin", adminRouter);
  app.use("/identity", identityRouter);

  // Create mock Hub server for this test (port already assigned above)
  // Emit catalog with test servers (strictness is determined by identity)
  const mockHubServer = new MockHubServer({
    port: hubPort,
    logger: testLogger,
    catalogPayload: catalogItemsToPayload(catalogItems),
  });
  // Accept the test instance ID from .it.env
  mockHubServer.setValidTokens(["it-run"]);

  const harness = new TestHarness(
    client,
    httpServer,
    services,
    testLogger,
    mockHubServer,
    clientConnectExtraHeaders,
    targetServers,
  );

  // Track the mcpxLogger in this harness instance
  harness.addLogger(mcpxLogger);

  return harness;
}
