import { loggableError } from "@mcpx/toolkit-core/logging";
import { withPolling } from "@mcpx/toolkit-core/time";
import {
  UnauthorizedError,
  discoverOAuthProtectedResourceMetadata,
  discoverAuthorizationServerMetadata,
} from "@modelcontextprotocol/sdk/client/auth.js";
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 { Logger } from "winston";
import {
  SSETargetServer,
  StreamableHttpTargetServer,
} from "../model/target-servers.js";
import { McpxOAuthProviderI } from "../oauth-providers/model.js";
import { DEVICE_FLOW_COMPLETE } from "../oauth-providers/device-flow.js";
import { OAuthSessionManagerI } from "../server/oauth-session-manager.js";
import { ExtendedClientBuilderI, ExtendedClientI } from "./client-extension.js";
import { buildClient } from "./target-server-connection-factory.js";

const OAUTH_POLLING_INTERVAL_MS = 1000;
const OAUTH_URL_POLLING_MAX_ATTEMPTS = 30;
// Device codes typically expire in 15 minutes, so poll for up to 15 minutes
const DEVICE_FLOW_TIMEOUT_MS = 15 * 60 * 1000;
const DEVICE_FLOW_MAX_POLL_ATTEMPTS = Math.ceil(
  DEVICE_FLOW_TIMEOUT_MS / OAUTH_POLLING_INTERVAL_MS,
);

type RemoteTargetServer = SSETargetServer | StreamableHttpTargetServer;
type RemoteTransport = SSEClientTransport | StreamableHTTPClientTransport;

/**
 * Represents a pending OAuth flow waiting for authorization code
 */
interface PendingOAuthFlow {
  provider: McpxOAuthProviderI;
  transport: RemoteTransport;
  client: Client;
  targetServer: RemoteTargetServer;
  /** For device flows: callback to invoke when auto-completion finishes */
  onComplete?: (client: ExtendedClientI) => void | Promise<void>;
}

/**
 * Result of initiating an OAuth flow
 */
export interface InitiateOAuthResult {
  authorizationUrl: string;
  state: string;
  userCode?: string;
}

/**
 * Handles OAuth authentication flows for target server connections
 */
export class OAuthConnectionHandler {
  // Store pending OAuth flows by server name for two-phase completion
  private pendingFlows: Map<string, PendingOAuthFlow> = new Map();

  constructor(
    private oauthSessionManager: OAuthSessionManagerI,
    private extendedClientBuilder: ExtendedClientBuilderI,
    private logger: Logger,
  ) {
    this.logger = logger.child({ component: "OAuthConnectionHandler" });
  }

  /**
   * Tries to connect using stored OAuth tokens.
   *
   * Returns the connected client on success, or `undefined` on any failure.
   *
   * On a 401 response the stored tokens are deleted here — they are invalid and
   * a fresh OAuth flow is required. For any other failure (network, timeout, etc.)
   * tokens are preserved so the watchdog can retry once the server is reachable again.
   */
  async safeTryWithExistingTokens(
    targetServer: RemoteTargetServer,
  ): Promise<ExtendedClientI | undefined> {
    const targetServerTypeStr =
      targetServer.type === "sse" ? "SSE" : "StreamableHTTP";

    // Get OAuth provider from session manager (coordinated flow)
    const authProvider = this.oauthSessionManager.getOrCreateOAuthProvider({
      serverName: targetServer.name,
      serverUrl: targetServer.url,
    });

    // Check if we already have valid tokens
    let existingTokens: Awaited<ReturnType<typeof authProvider.tokens>>;
    try {
      existingTokens = await authProvider.tokens();
    } catch (error) {
      this.logger.warn("Failed to read tokens", {
        name: targetServer.name,
        error: loggableError(error),
      });
      return undefined;
    }
    if (!existingTokens) {
      return undefined;
    }
    this.logger.info("Found existing tokens, attempting to use them", {
      name: targetServer.name,
    });

    // Create transport with existing auth provider
    const transport =
      targetServer.type === "sse"
        ? new SSEClientTransport(new URL(targetServer.url), {
            authProvider,
          })
        : new StreamableHTTPClientTransport(new URL(targetServer.url), {
            authProvider,
          });

    // Try to use the existing tokens
    try {
      // If successful, create and return the client
      const client = buildClient(targetServer.name);
      await client.connect(transport);

      const extendedClient = await this.extendedClientBuilder.build({
        name: targetServer.name,
        originalClient: client,
      });

      this.logger.info(
        `${targetServerTypeStr} Client connected with existing OAuth tokens`,
        {
          name: targetServer.name,
          url: targetServer.url,
        },
      );
      return extendedClient;
    } catch (error) {
      if (isAuthenticationError(error)) {
        // Tokens were rejected — delete them so the next attempt triggers a fresh OAuth flow.
        this.logger.info(
          "Existing tokens rejected by server (401), clearing tokens",
          { name: targetServer.name, error: loggableError(error) },
        );
        await this.deleteOAuthTokensForServer(targetServer.name);
      } else {
        // Network / transport error — tokens may still be valid once the server is back.
        this.logger.info(
          "Server unreachable with existing tokens, preserving for retry",
          {
            name: targetServer.name,
            error: loggableError(error),
          },
        );
      }
      return undefined;
    }
  }

  /**
   * Phase 1: Initiate OAuth flow, get authorization URL, return immediately.
   * The server stays in pending-auth state until completeOAuth() is called.
   * This enables non-blocking OAuth flows for both UI and external orchestration.
   *
   * For device flows, an optional `onComplete` callback can be provided.
   * The device flow polling will auto-complete and invoke this callback
   * when the user authorizes the device.
   */
  async initiateOAuth(
    targetServer: RemoteTargetServer,
    options?: {
      callbackUrl?: string;
      /** Callback invoked when device flow auto-completes */
      onComplete?: (client: ExtendedClientI) => void | Promise<void>;
    },
  ): Promise<InitiateOAuthResult> {
    const { callbackUrl, onComplete } = options ?? {};
    const authProvider = this.oauthSessionManager.getOrCreateOAuthProvider({
      serverName: targetServer.name,
      serverUrl: targetServer.url,
      callbackUrl,
    });

    const state = authProvider.state();
    this.oauthSessionManager.startOAuthFlow(
      targetServer.name,
      targetServer.url,
      state,
    );

    this.logger.info("Initiating OAuth flow", {
      name: targetServer.name,
      callbackUrl,
    });

    // Discover auth server metadata and request offline_access if supported,
    // so servers that support refresh tokens will issue one.
    await this.applyDiscoveredScope(targetServer.url, authProvider);

    // Create transport with auth provider - this will trigger OAuth flow
    const transport =
      targetServer.type === "sse"
        ? new SSEClientTransport(new URL(targetServer.url), { authProvider })
        : new StreamableHTTPClientTransport(new URL(targetServer.url), {
            authProvider,
          });

    // Start the transport WITHOUT awaiting - this triggers OAuth flow in the background.
    // The transport will block waiting for auth to complete, but we want to return
    // the authorization URL immediately. The URL is set synchronously in the auth
    // provider's redirectToAuthorization() before it starts waiting for the code.
    transport.start().catch((_e: unknown) => {
      this.logger.debug("expected transport.start() error, continuing");
    });

    if (targetServer.type === "streamable-http") {
      // Hack inspired by `mcp-remote` - also don't await
      const testTransport = new StreamableHTTPClientTransport(
        new URL(targetServer.url),
        { authProvider },
      );
      const testClient = new Client(
        { name: "mcpx-fallback-test", version: "0.0.0" },
        { capabilities: {} },
      );
      testClient.connect(testTransport).catch((_e: unknown) => {
        this.logger.debug(
          "expected client.connect() error on mcpx-fallback-test, continuing",
          { error: loggableError(_e) },
        );
      });
    }

    // Poll for authorization URL (quick, should be available fast)
    const authorizationUrl = await withPolling({
      maxAttempts: OAUTH_URL_POLLING_MAX_ATTEMPTS,
      sleepTimeMs: OAUTH_POLLING_INTERVAL_MS,
      getValue: () => authProvider.getAuthorizationUrl(),
      found: (url): url is URL => Boolean(url),
    });

    if (!authorizationUrl) {
      throw new Error(
        `Failed to obtain authorization URL for server: ${targetServer.name}`,
      );
    }

    // Store for later completion
    const client = buildClient(targetServer.name);
    this.pendingFlows.set(targetServer.name, {
      provider: authProvider,
      transport,
      client,
      targetServer,
      onComplete,
    });

    this.logger.debug("OAuth flow initiated, waiting for authorization", {
      name: targetServer.name,
      state,
    });

    // For device flows, start background polling to auto-complete when user authorizes
    if (authProvider.type === "device_flow") {
      this.startDeviceFlowPoller(targetServer.name);
    }

    return {
      authorizationUrl: authorizationUrl.toString(),
      state,
      userCode:
        authProvider.type === "device_flow"
          ? (authProvider.getUserCode() ?? undefined)
          : undefined,
    };
  }

  /**
   * Phase 2: Complete OAuth flow with authorization code.
   * Called when the authorization code is received (via callback or external system).
   *
   * On error, the pending flow is cleaned up and the error is rethrown.
   * The caller (UpstreamHandler) is responsible for updating server state on failure.
   */
  async completeOAuth(
    serverName: string,
    authorizationCode: string,
  ): Promise<ExtendedClientI> {
    const pending = this.pendingFlows.get(serverName);
    if (!pending) {
      throw new Error(`No pending OAuth flow for server: ${serverName}`);
    }

    const { provider, transport, client, targetServer } = pending;

    this.logger.debug("Completing OAuth flow", { serverName });

    try {
      // Complete authorization (resolves the provider's promise)
      provider.completeAuthorization(authorizationCode);

      // Handle device flow differently - tokens are already obtained
      if (
        provider.type === "device_flow" &&
        authorizationCode === DEVICE_FLOW_COMPLETE
      ) {
        this.logger.info("Device flow completed, using saved tokens", {
          name: serverName,
        });

        const freshTransport =
          targetServer.type === "sse"
            ? new SSEClientTransport(new URL(targetServer.url), {
                authProvider: provider,
              })
            : new StreamableHTTPClientTransport(new URL(targetServer.url), {
                authProvider: provider,
              });

        await client.connect(freshTransport);
      } else {
        // Standard OAuth authorization code flow - exchange code for tokens
        await transport.finishAuth(authorizationCode);

        const postAuthTransport =
          targetServer.type === "sse"
            ? new SSEClientTransport(new URL(targetServer.url), {
                authProvider: provider,
              })
            : new StreamableHTTPClientTransport(new URL(targetServer.url), {
                authProvider: provider,
              });

        await client.connect(postAuthTransport);
      }

      const extendedClient = await this.extendedClientBuilder.build({
        name: serverName,
        originalClient: client,
      });

      this.logger.info("OAuth flow completed successfully", {
        name: serverName,
        url: targetServer.url,
      });

      return extendedClient;
    } finally {
      this.pendingFlows.delete(serverName);
    }
  }

  /**
   * Cancel a pending OAuth flow (cleanup)
   */
  cancelPendingOAuth(serverName: string): boolean {
    const existed = this.pendingFlows.delete(serverName);
    if (existed) {
      this.logger.info("Cancelled pending OAuth flow", { serverName });
    }
    return existed;
  }

  /**
   * Discovers the auth server's metadata and calls setDiscoveredScope("offline_access")
   * on the provider when the server lists it in scopes_supported.
   * Non-fatal: errors are logged and the flow continues without the extra scope.
   */
  private async applyDiscoveredScope(
    serverUrl: string,
    provider: McpxOAuthProviderI,
  ): Promise<void> {
    try {
      const resourceMeta =
        await discoverOAuthProtectedResourceMetadata(serverUrl);
      const authServerUrl =
        resourceMeta?.authorization_servers?.[0] ?? serverUrl;
      const authMeta = await discoverAuthorizationServerMetadata(authServerUrl);
      if (authMeta?.scopes_supported?.includes("offline_access")) {
        provider.setDiscoveredScope("offline_access");
        this.logger.debug("Requested offline_access scope", {
          serverUrl,
          provider: provider.serverName,
        });
      }
    } catch (error) {
      // Non-fatal — proceed without offline_access
      this.logger.debug("Could not discover auth server scopes, continuing", {
        serverUrl,
        error: loggableError(error),
      });
    }
  }

  /**
   * Deletes stored OAuth tokens for a server and removes it from the provider cache.
   * Also cancels any pending OAuth flow for the server.
   */
  async deleteOAuthTokensForServer(serverName: string): Promise<void> {
    this.cancelPendingOAuth(serverName);
    await this.oauthSessionManager.deleteOAuthTokensForServer(serverName);
    this.logger.info("Deleted OAuth tokens for server", { serverName });
  }

  /**
   * Start background polling for device flow authorization.
   * Device flows require repeated calls to getAuthorizationCode() to trigger
   * internal token polling. When the user authorizes, the provider returns
   * DEVICE_FLOW_COMPLETE and we auto-complete the OAuth flow.
   */
  private startDeviceFlowPoller(serverName: string): void {
    const startedAt = Date.now();
    let attempts = 0;

    const poll = async (): Promise<void> => {
      const pending = this.pendingFlows.get(serverName);
      if (!pending) {
        // Flow was cancelled or completed externally
        return;
      }

      attempts++;
      if (attempts > DEVICE_FLOW_MAX_POLL_ATTEMPTS) {
        this.logger.warn("Device flow polling timed out", {
          serverName,
          attempts,
          elapsedMs: Date.now() - startedAt,
        });
        this.cancelPendingOAuth(serverName);
        return;
      }

      const code = pending.provider.getAuthorizationCode();
      if (code === DEVICE_FLOW_COMPLETE) {
        this.logger.debug(
          "Device flow authorization detected, auto-completing",
          { serverName },
        );
        try {
          const extendedClient = await this.completeOAuth(
            serverName,
            DEVICE_FLOW_COMPLETE,
          );
          await pending.onComplete?.(extendedClient);
        } catch (error) {
          this.logger.error("Failed to auto-complete device flow", {
            serverName,
            error: loggableError(error),
          });
        }
        return;
      }

      // Schedule next poll attempt
      setTimeout(poll, OAUTH_POLLING_INTERVAL_MS);
    };

    // Start polling (don't await - runs in background)
    poll().catch((error) => {
      this.logger.error("Device flow poller error", {
        serverName,
        error: loggableError(error),
      });
    });
  }

  /**
   * Returns true if the server has an OAuth provider registered
   * (i.e., it has been through OAuth at some point in this session).
   */
  isOAuthServer(serverName: string): boolean {
    return this.oauthSessionManager.hasOAuthProvider(serverName);
  }

  /**
   * Returns true if the server's stored token is expired or missing.
   * Used by the background expiry monitor.
   */
  async isTokenExpiredForServer(
    targetServer: RemoteTargetServer,
  ): Promise<boolean> {
    const provider = this.oauthSessionManager.getExistingOAuthProvider(
      targetServer.name,
    );
    if (!provider) return false;
    try {
      const tokens = await provider.tokens();
      return tokens === undefined;
    } catch (error) {
      this.logger.warn("Failed to read tokens, not treating as expired", {
        name: targetServer.name,
        error: loggableError(error),
      });
      return false;
    }
  }

  /**
   * Lookup server name by OAuth state parameter
   */
  getServerNameByState(state: string): string | null {
    const flow = this.oauthSessionManager.getOAuthFlow(state);
    return flow?.serverName ?? null;
  }

  /**
   * Cleanup OAuth flow tracking by state
   */
  completeOAuthFlowCleanup(state: string): void {
    this.oauthSessionManager.completeOAuthFlow(state);
  }
}

/**
 * Checks if an error indicates OAuth authentication is required
 */
export function isAuthenticationError(error: unknown): boolean {
  if (error instanceof UnauthorizedError) {
    return true;
  }
  if (typeof error === "object" && error !== null) {
    const err = error as Record<string, unknown>;
    return (
      (err?.["response"] as Record<string, unknown>)?.["status"] === 401 ||
      err?.["status"] === 401 ||
      err?.["code"] === 401 ||
      (typeof err?.["message"] === "string" &&
        err["message"].includes("401")) ||
      (typeof err?.["message"] === "string" &&
        err["message"].includes("Unauthorized"))
    );
  }
  return false;
}
