import { type ZodType } from '@frontmcp/lazy-zod';
import {
  ElicitResultSchema,
  EmptyResultSchema,
  McpServer,
  type RequestId,
  type StreamableHTTPServerTransport,
} from '@frontmcp/protocol';

import { type SSEServerTransport } from '#sse-transport';

import { ServerRequestTokens, type FrontMcpLogger, type ServerResponse } from '../../common';
import {
  type ElicitationStore,
  type ElicitOptions,
  type ElicitResult,
  type McpElicitResult,
  type PendingElicit,
} from '../../elicitation';
import { ElicitationNotSupportedError, UnauthorizedError } from '../../errors';
import { type Scope } from '../../scope';
import { type AuthenticatedServerRequest, type SdkAuthInfo } from '../../server/server.types';
import { buildChannelInstructions, composeInitializeInstructions } from '../../skill/skill-instructions.helper';
import { computeTaskCapabilities } from '../../task';
import { createMcpHandlers } from '../mcp-handlers';
import { type TransportKey, type TransportType } from '../transport.types';
import { type RecreateableSSEServerTransport } from './sse-transport';
import { type RecreateableStreamableHTTPServerTransport } from './streamable-http-transport';

/**
 * Merge multiple capability fragments while preserving keys inside
 * `experimental` and `extensions` from every contributor (channels +
 * SEP-2640 skills both live there). A naive `{...a, ...b}` merge would
 * have later contributors clobber earlier `experimental` keys.
 */
function mergeExperimental(fragments: Array<Record<string, unknown>>): Record<string, unknown> {
  const out: Record<string, unknown> = {};
  const experimental: Record<string, unknown> = {};
  const extensions: Record<string, unknown> = {};
  let sawExperimental = false;
  let sawExtensions = false;

  for (const frag of fragments) {
    if (!frag) continue;
    for (const [key, value] of Object.entries(frag)) {
      if (key === 'experimental' && value && typeof value === 'object') {
        Object.assign(experimental, value as Record<string, unknown>);
        sawExperimental = true;
      } else if (key === 'extensions' && value && typeof value === 'object') {
        Object.assign(extensions, value as Record<string, unknown>);
        sawExtensions = true;
      } else {
        out[key] = value;
      }
    }
  }

  if (sawExperimental) out['experimental'] = experimental;
  if (sawExtensions) out['extensions'] = extensions;
  return out;
}

/**
 * Base transport type that includes all supported transports.
 * RecreateableStreamableHTTPServerTransport extends StreamableHTTPServerTransport
 * and RecreateableSSEServerTransport extends SSEServerTransport,
 * so they're also included in this union.
 */
export type SupportedTransport =
  | StreamableHTTPServerTransport
  | SSEServerTransport
  | RecreateableStreamableHTTPServerTransport
  | RecreateableSSEServerTransport;

export abstract class LocalTransportAdapter<T extends SupportedTransport> {
  protected logger: FrontMcpLogger;
  protected transport: T;

  /**
   * Pending elicitation request. Only one elicit per session is allowed.
   * New elicit requests will cancel any pending one.
   */
  protected pendingElicit?: PendingElicit;

  /**
   * Session payload fields set during MCP initialize.
   * Stored on the adapter instance so they survive across requests in SSE mode
   * (where each POST creates a fresh anonymous HTTP session) and are available
   * in distributed environments without relying on local in-memory caches.
   */
  private initSessionPayload?: Partial<{
    supportsElicitation: boolean;
    platformType: string;
    clientName: string;
    clientVersion: string;
  }>;

  /**
   * Called by the initialize request handler to persist initialization data
   * on the transport adapter instance. This data is merged into the session
   * payload on every subsequent request via ensureAuthInfo().
   */
  setInitSessionPayload(payload: NonNullable<LocalTransportAdapter<T>['initSessionPayload']>): void {
    this.initSessionPayload = payload;
  }

  #requestId = 1;
  ready: Promise<void>;
  server: McpServer;

  constructor(
    protected readonly scope: Scope,
    protected readonly key: TransportKey,
    protected readonly onDispose: () => void,
    res: ServerResponse,
  ) {
    this.logger = scope.logger.child('LocalTransportAdapter');
    this.transport = this.createTransport(key.sessionId, res);
    this.ready = this.connectServer();
  }

  abstract createTransport(sessionId: string, response: ServerResponse): T;

  abstract initialize(req: AuthenticatedServerRequest, res: ServerResponse): Promise<void>;

  /**
   * Send an elicitation request to the client.
   *
   * @param relatedRequestId - The request ID that triggered this elicit
   * @param message - Message to display to the user
   * @param requestedSchema - Zod schema for the expected response
   * @param options - Elicit options (mode, ttl, elicitationId)
   * @returns ElicitResult with status and typed content
   */
  abstract sendElicitRequest<S extends ZodType>(
    relatedRequestId: RequestId,
    message: string,
    requestedSchema: S,
    options?: ElicitOptions,
  ): Promise<ElicitResult<S extends ZodType<infer O> ? O : unknown>>;

  abstract handleRequest(req: AuthenticatedServerRequest, res: ServerResponse): Promise<void>;

  /**
   * Whether this transport has already been initialized via the MCP initialize handshake.
   * Override in subclasses that track initialization state.
   */
  get isInitialized(): boolean {
    return false;
  }

  /**
   * Marks this transport as pre-initialized for session recreation.
   * Override in subclasses that need to set the MCP SDK's _initialized flag.
   */
  markAsInitialized(): void {
    // Default no-op - override in subclasses
  }

  /**
   * Resets initialization state to allow re-initialization.
   * Override in subclasses that support session re-initialization.
   */
  resetForReinitialization(): void {
    // Default no-op - override in subclasses
  }

  /**
   * Re-register the MCP server with the notification service.
   * Called after resetForReinitialization() to restore the server mapping
   * that was removed by terminateSession → unregisterServer.
   */
  reregisterServer(): void {
    if (this.server) {
      this.scope.notifications.registerServer(this.key.sessionId, this.server);
    }
  }

  connectServer() {
    const { info, apps } = this.scope.metadata;

    // Check if there are remote apps configured (they will have tools/resources/prompts that load later)
    // Remote apps have 'urlType' property indicating they're external MCP servers
    const hasRemoteApps = apps?.some((app) => this.isRemoteApp(app)) ?? false;

    // Check if completions capability should be enabled (when prompts or resources are present)
    // Also enable for remote apps since they may have prompts/resources
    const hasPrompts = this.scope.prompts.hasAny() || hasRemoteApps;
    const hasResources = this.scope.resources.hasAny() || hasRemoteApps;
    const hasTools = this.scope.tools.hasAny() || hasRemoteApps;
    const hasAgents = this.scope.agents.hasAny();
    const completionsCapability = hasPrompts || hasResources ? { completions: {} } : {};

    // Get capabilities from registries - they now handle both local and remote capabilities
    // When hasRemoteApps is true, we pre-advertise listChanged capabilities since remote tools/resources/prompts
    // load asynchronously after connection and may emit change notifications
    const remoteCapabilities = hasRemoteApps ? this.buildRemoteCapabilities() : {};

    const elicitationCapability = this.scope.metadata.elicitation?.enabled ? { elicitation: {} } : {};

    // Channel capabilities (experimental extension for Claude Code)
    const channelCapabilities = this.scope.channels?.getCapabilities() ?? {};

    // SEP-2640 (Skills Extension) — capability declaration and optional
    // skill-URI hints. `skillsCapabilities` advertises the extension; the
    // SEP-2640 URI hint block (gated by `skillsConfig.sep2640InInstructions`)
    // is appended to the catalog summary produced by
    // `composeInitializeInstructions` so both surfaces compose cleanly.
    const skillsCapabilities = this.scope.skills?.getCapabilities() ?? {};

    // Compose `instructions` lazily on every `initialize` so dynamic skill
    // registrations after boot are reflected per-session. The static value
    // below seeds the McpServer constructor for SDK compatibility; the
    // actual response is recomputed inside the handler via
    // `composeInstructions` (see initialize-request.handler.ts).
    const composeInstructions = (): string => {
      const composed = composeInitializeInstructions({
        userInstructions: this.scope.metadata.instructions,
        channelInstructions: buildChannelInstructions(this.scope.channels),
        skillRegistry: this.scope.skills,
        policy: this.scope.metadata.skillsConfig?.injectInstructions,
      });
      const sep2640Hints = this.buildSkillInstructionHints();
      return [composed, sep2640Hints].filter((s) => s.length > 0).join('\n\n---\n\n');
    };
    const instructions = composeInstructions();

    const mergedCapabilities = mergeExperimental([
      remoteCapabilities, // Pre-advertise for remote apps (may be overwritten by local)
      this.scope.tools.getCapabilities(),
      this.scope.resources.getCapabilities(),
      this.scope.prompts.getCapabilities(),
      this.scope.agents.getCapabilities(), // Include agent capabilities (agents as tools)
      channelCapabilities, // Channel capabilities (claude/channel extension)
      skillsCapabilities, // SEP-2640 extension declaration
      completionsCapability,
      computeTaskCapabilities(this.scope),
      // MCP logging protocol support - allows clients to set log level via logging/setLevel
      { logging: {} },
      elicitationCapability,
    ]);

    // The MCP `ServerOptions.capabilities` type doesn't model the
    // forward-compat `extensions` key (reserved for SEP-2133). The runtime
    // passthrough is fine; the cast keeps TS happy.
    const serverOptions = {
      instructions,
      capabilities: mergedCapabilities as Record<string, never>,
      serverInfo: info,
    };

    this.logger.info('connectServer: advertising capabilities', {
      hasTools: hasTools || hasAgents, // Agents expose themselves as tools
      hasResources,
      hasPrompts,
      hasAgents,
      hasRemoteApps,
      capabilities: JSON.stringify(serverOptions.capabilities),
      serverInfo: JSON.stringify(serverOptions.serverInfo),
    });

    this.server = new McpServer(info, serverOptions);
    const handlers = createMcpHandlers({
      scope: this.scope,
      serverOptions,
      composeInstructions,
    });

    for (const handler of handlers) {
      this.server.setRequestHandler(handler.requestSchema, handler.handler as any);
    }

    // Register server with notification service for server→client notifications
    this.scope.notifications.registerServer(this.key.sessionId, this.server);

    return this.server.connect(this.transport);
  }

  get newRequestId(): RequestId {
    return this.#requestId++;
  }

  /**
   * Get the transport type (sse, streamable-http, etc.).
   * Used for transport-specific behavior detection.
   */
  get type(): TransportType {
    return this.key.type;
  }

  async destroy(reason?: string): Promise<void> {
    console.log('destroying transporter, reason:', reason);

    // Unregister server from notification service
    this.scope.notifications.unregisterServer(this.key.sessionId);

    try {
      // if(!this.transport.closed){
      //   this.transport.close();
      // }
    } catch {
      /* empty */
    }
    if (reason) {
      console.warn(`Destroying transporter: ${reason}`);
    } else {
      try {
        this.onDispose?.();
      } catch {
        /* empty */
      }
    }
  }

  /**
   * Ping the connected client for this transport.
   * Returns true on success, false on timeout/error.
   */
  async ping(timeoutMs = 10_000): Promise<boolean> {
    try {
      await this.ready; // ensure server.connect(...) finished
      // Preferred if your SDK exposes it:
      // await this.server.ping({ timeout: timeoutMs });

      // Works on all versions (low-level request):
      await this.server.request({ method: 'ping' }, EmptyResultSchema, { timeout: timeoutMs });
      return true;
    } catch {
      return false;
    }
  }

  protected ensureAuthInfo(req: AuthenticatedServerRequest, transport: LocalTransportAdapter<T>) {
    const { token, user, session } = req[ServerRequestTokens.auth];

    // Session must always exist at this point — created in session.verify (public mode)
    // or synced from flow state in onInitialize (reconnect). A missing session here
    // means a verified token could not be tied back to a live session (e.g. the
    // session was evicted/expired and could not be reconstructed). That is an
    // authentication-recovery condition for the client, NOT an internal bug, so we
    // surface a 401 challenge (#471) rather than letting a plain Error bubble up as a
    // generic 500/-32000. PublicMcpError-derived errors are rendered by the flow
    // runner with their statusCode and any `wwwAuthenticate` challenge.
    // The previous fallback:${Date.now()} was predictable, unencrypted, and could collide.
    if (!session?.id) {
      // RFC 6750 §3 — challenge the client to re-authenticate with a bearer token.
      throw new UnauthorizedError(
        'Session could not be reconstructed for the provided token. Re-authenticate to start a new session.',
        'Bearer',
      );
    }
    const sessionId = session.id;
    const sessionPayload = session.payload ?? { protocol: 'streamable-http' as const };

    // Enrich session payload with initialization data stored on the adapter.
    // In SSE mode, each HTTP request creates a fresh anonymous session, so
    // fields set during MCP initialize (supportsElicitation, platformType, etc.)
    // are lost. The adapter instance persists for the transport's lifetime, so
    // we merge the initialization payload here. This works in distributed mode
    // too: the adapter is always on the node that owns the transport session.
    if (this.initSessionPayload) {
      const init = this.initSessionPayload;
      if (init.supportsElicitation !== undefined && sessionPayload.supportsElicitation === undefined) {
        sessionPayload.supportsElicitation = init.supportsElicitation;
      }
      if (init.platformType !== undefined && sessionPayload.platformType === undefined) {
        (sessionPayload as Record<string, unknown>)['platformType'] = init.platformType;
      }
      if (init.clientName !== undefined && sessionPayload.clientName === undefined) {
        sessionPayload.clientName = init.clientName;
      }
      if (init.clientVersion !== undefined && sessionPayload.clientVersion === undefined) {
        sessionPayload.clientVersion = init.clientVersion;
      }
    }

    const authInfo: SdkAuthInfo = {
      token,
      user,
      sessionId,
      sessionIdPayload: sessionPayload,
      scopes: [],
      clientId: user.sub ?? '',
      transport,
    };
    req.auth = authInfo;
    return req.auth;
  }

  /**
   * Get the elicitation store for distributed elicitation support.
   * Uses Redis in distributed mode, in-memory for single-node.
   */
  protected get elicitStore(): ElicitationStore | undefined {
    return this.scope.elicitationStore;
  }

  /**
   * Get the elicitation store, throwing if not initialized.
   * Use in contexts where elicitation is required (sendElicitRequest, cancelPendingElicit).
   */
  protected requireElicitStore(): ElicitationStore {
    const store = this.elicitStore;
    if (!store) {
      throw new Error('Elicitation store not initialized');
    }
    return store;
  }

  /**
   * Cancel any pending elicitation request.
   * Called before sending a new elicit to enforce single-elicit-per-session.
   *
   * This cancels both the local pending elicit (for timeout handling)
   * and publishes cancel to the store (for distributed mode).
   *
   * Note: The local promise is resolved immediately for responsiveness,
   * then distributed state cleanup follows asynchronously. This non-atomic
   * sequence is intentional - the worst case is publishing a cancel for
   * an already-processed elicitation, which is harmless.
   */
  protected async cancelPendingElicit(): Promise<void> {
    if (this.pendingElicit) {
      clearTimeout(this.pendingElicit.timeoutHandle);
      // Resolve local promise immediately for responsiveness
      this.pendingElicit.resolve({ status: 'cancel' });

      // Publish cancel to store for distributed mode (non-atomic, intentional)
      // In distributed mode, another node may have already processed this elicitation
      const sessionId = this.key.sessionId;
      const store = this.elicitStore;
      if (store) {
        const pending = await store.getPending(sessionId);
        if (pending) {
          await store.publishResult(pending.elicitId, sessionId, { status: 'cancel' });
        }
      }

      this.pendingElicit = undefined;
      this.logger.info('Cancelled previous pending elicit');
    }
  }

  /**
   * Handle an incoming elicitation result from the client.
   * Returns true if the request was an elicit result and was handled.
   *
   * Uses ElicitationResultFlow for processing (with hook support).
   * In distributed mode, this publishes the result via the elicitation store,
   * which routes it to the correct node that's waiting for it.
   */
  handleIfElicitResult(req: AuthenticatedServerRequest): boolean {
    this.logger.info('[handleIfElicitResult] checking request', {
      hasPendingElicit: !!this.pendingElicit,
      bodyKeys: req.body ? Object.keys(req.body) : [],
      hasResult: !!req.body?.result,
    });
    // First, check for error response (client doesn't support elicitation)
    // This still needs the local pendingElicit for backwards compatibility
    if (this.pendingElicit && req.body.error) {
      const { code, message } = req.body.error;
      if (code === -32600 && message === 'Elicitation not supported') {
        clearTimeout(this.pendingElicit.timeoutHandle);
        this.pendingElicit.reject(new ElicitationNotSupportedError());
        this.pendingElicit = undefined;
        return true;
      }
    }

    // Parse the elicit result from the client
    const parsed = ElicitResultSchema.safeParse(req.body?.result);
    if (!parsed.success) {
      return false;
    }

    // Run ElicitationResultFlow for distributed routing (with hook support)
    // Flow handles: lookup pending, publish via pub/sub
    // Note: Using void to explicitly mark as fire-and-forget; errors are caught inside the method
    const sessionId = this.key.sessionId;
    void this.handleElicitResultAsync(sessionId, parsed.data);

    // Map MCP action to our result type for local handling
    const action = parsed.data.action;
    const result: ElicitResult = {
      status: action,
      ...(action === 'accept' && parsed.data.content !== undefined && { content: parsed.data.content }),
    };

    // Also handle local pending elicit (for single-node mode and error handling)
    if (this.pendingElicit) {
      clearTimeout(this.pendingElicit.timeoutHandle);
      this.pendingElicit.resolve(result);
      this.pendingElicit = undefined;
    }

    return true;
  }

  /**
   * Async handler for elicit result - uses ElicitationResultFlow for processing.
   * Called from handleIfElicitResult without awaiting to avoid blocking the response.
   */
  private async handleElicitResultAsync(sessionId: string, mcpResult: McpElicitResult): Promise<void> {
    try {
      // Run the flow for hook support and distributed routing
      const output = await this.scope.runFlow('elicitation:result', {
        sessionId,
        result: mcpResult,
      });

      if (output?.handled) {
        this.logger.verbose('ElicitationResultFlow handled result', {
          elicitId: output.elicitId,
          sessionId: sessionId.slice(0, 20),
          status: output.result?.status,
        });
      }
    } catch (error) {
      this.logger.warn('Failed to run ElicitationResultFlow', {
        sessionId: sessionId.slice(0, 20),
        error: (error as Error).message,
      });
    }
  }

  /**
   * Type predicate to check if an app configuration represents a remote MCP app.
   * Remote apps have a 'urlType' property and are not standalone.
   */
  private isRemoteApp(app: unknown): boolean {
    if (app === null || typeof app !== 'object') return false;
    const appObj = app as Record<string, unknown>;
    // Remote apps have 'urlType' property indicating they're external MCP servers
    // Standalone apps don't contribute to the main server's capabilities
    return 'urlType' in appObj && appObj['standalone'] !== true;
  }

  /**
   * SEP-2640 §Discovery — opt-in `instructions` text listing each
   * MCP-visible skill's `skill://` URI. Returns an empty string unless
   * `skillsConfig.sep2640InInstructions` is true.
   */
  private buildSkillInstructionHints(): string {
    const skillRegistry = this.scope.skills;
    if (!skillRegistry?.hasAny()) return '';

    const skillsConfig = this.scope.metadata?.skillsConfig as { sep2640InInstructions?: boolean } | undefined;
    if (!skillsConfig?.sep2640InInstructions) return '';

    const visible = skillRegistry.getSkills({ visibility: 'mcp' });
    if (visible.length === 0) return '';

    const lines = ['Available skills (load via resources/read):'];
    for (const skill of visible) {
      const uri = `skill://${skill.getSkillPath()}/SKILL.md`;
      lines.push(`- ${uri} — ${skill.metadata.description}`);
    }

    // Include any extra URIs registered programmatically.
    const extras = skillRegistry.getSep2640InstructionUris?.() ?? [];
    for (const extra of extras) {
      lines.push(`- ${extra}`);
    }

    return lines.join('\n');
  }

  /**
   * Build capabilities that should be pre-advertised for remote apps.
   * Remote apps load their tools/resources/prompts asynchronously after connection,
   * so we need to pre-advertise listChanged capabilities for proper notification support.
   */
  private buildRemoteCapabilities(): Record<string, unknown> {
    return {
      tools: { listChanged: true },
      resources: { subscribe: true, listChanged: true },
      prompts: { listChanged: true },
    };
  }
}
