import {
  PublicClientApplication,
  type DeviceCodeRequest,
  type AuthenticationResult,
} from '@azure/msal-node';

import { randomUUID } from 'crypto';

import {
  loadCache,
  saveCache,
  clearCache,
} from './token-cache.ts';

import {
  saveSession,
  loadSession,
  clearSessionFile,
} from './session-cache.ts';

export interface DeviceCodeInfo {
  message: string;
  userCode: string;
  verificationUrl: string;
  sessionId: string;
}

interface SessionCredential {
  tenantId: string;
  clientId: string;
  scopes: string[];
}

export default class AuthManager {
  private defaultScopes: string[];

  private pendingDeviceFlows: Map<
    string,
    {
      promise: Promise<string>;
      deviceCodeInfo: DeviceCodeInfo;
      userId?: string;
      error?: unknown;
    }
  > = new Map();

  private sessionToUser: Map<string, string> = new Map();

  // 🔥 NEW: MSAL client per session
  private sessionToMsal: Map<string, PublicClientApplication> =
    new Map();
  // store credential info per session (tenantId/clientId/scopes)
  private sessionToCredential: Map<string, { tenantId?: string; clientId?: string; scopes?: string[] }> =
    new Map();

  constructor(scopes: string[]) {
    this.defaultScopes = scopes;
  }

  private createMsalClient(
    credential?: SessionCredential
  ): PublicClientApplication {
    const tenantId =
      credential?.tenantId ??
      process.env.TENANT_ID ??
      'common';

    const clientId =
      credential?.clientId ??
      process.env.CLIENT_ID!;

    return new PublicClientApplication({
      auth: {
        clientId,
        authority: `https://login.microsoftonline.com/${tenantId}/v2.0`,
      },
    });
  }

  private async loadUserCache(
    msalClient: PublicClientApplication,
    userId: string
  ) {
    const cacheData = await loadCache(userId);
    if (cacheData) {
      msalClient.getTokenCache().deserialize(cacheData);
    }
  }

  private async saveUserCache(
    msalClient: PublicClientApplication,
    userId: string
  ) {
    await saveCache(
      userId,
      msalClient.getTokenCache().serialize()
    );
  }

  async initiateDeviceCodeFlow(
    credential?: SessionCredential
  ): Promise<DeviceCodeInfo> {
    const sessionId = randomUUID();

    const scopes =
      credential?.scopes ?? this.defaultScopes;

    const msalClient =
      this.createMsalClient(credential);

    // persist credential per session so we can use the same scopes later
    this.sessionToCredential.set(sessionId, {
      tenantId: credential?.tenantId,
      clientId: credential?.clientId,
      scopes,
    });

    // 🔥 store MSAL per session
    this.sessionToMsal.set(sessionId, msalClient);

    let resolveDeviceInfo!: (
      value: Omit<DeviceCodeInfo, 'sessionId'>
    ) => void;

    const deviceInfoPromise = new Promise<
      Omit<DeviceCodeInfo, 'sessionId'>
    >((resolve) => {
      resolveDeviceInfo = resolve;
    });

    const deviceCodeRequest: DeviceCodeRequest = {
      deviceCodeCallback: (resp) => {
        resolveDeviceInfo({
          message: resp.message,
          userCode: resp.userCode,
          verificationUrl: resp.verificationUri,
        });
      },
      scopes,
    };

    const flowPromise: Promise<string> =
      (async () => {
        try {
          const response: AuthenticationResult | null =
            await msalClient.acquireTokenByDeviceCode(
              deviceCodeRequest
            );

          if (!response) {
            throw new Error(
              'Failed to acquire token'
            );
          }

          const userId =
            response.account?.homeAccountId ??
            response.uniqueId;

          if (!userId) {
            throw new Error(
              'Failed to get user ID'
            );
          }

          this.sessionToUser.set(
            sessionId,
            userId
          );

          await saveSession(
            sessionId,
            userId
          );

          const flow =
            this.pendingDeviceFlows.get(
              sessionId
            );
          if (flow) {
            flow.userId = userId;
          }

          await this.saveUserCache(
            msalClient,
            userId
          );

          return response.accessToken;
        } catch (error) {
          const flow =
            this.pendingDeviceFlows.get(
              sessionId
            );
          if (flow) {
            flow.error = error;
          }
          throw error;
        } finally {
          setTimeout(() => {
            this.pendingDeviceFlows.delete(
              sessionId
            );
          }, 5 * 60 * 1000);
        }
      })();

    flowPromise.catch(() => { });

    const deviceCodeInfo = await Promise.race([
      deviceInfoPromise,
      new Promise<never>((_, reject) =>
        setTimeout(
          () =>
            reject(
              new Error(
                'Timeout waiting for device code'
              )
            ),
          10000
        )
      ),
    ]);

    const fullDeviceCodeInfo: DeviceCodeInfo = {
      ...deviceCodeInfo,
      sessionId,
    };

    this.pendingDeviceFlows.set(sessionId, {
      promise: flowPromise,
      deviceCodeInfo: fullDeviceCodeInfo,
    });

    return fullDeviceCodeInfo;
  }

  async acquireTokenBySessionId(
    sessionId: string,
    forceRefresh = false
  ): Promise<string> {
    const pendingFlow =
      this.pendingDeviceFlows.get(sessionId);

    if (pendingFlow && !pendingFlow.userId) {
      await pendingFlow.promise;
    }

    let userId =
      this.sessionToUser.get(sessionId);

    if (!userId) {
      const loadedUserId =
        await loadSession(sessionId);

      if (loadedUserId) {
        userId = loadedUserId;
        this.sessionToUser.set(
          sessionId,
          userId
        );
      }
    }

    if (!userId) {
      throw new Error(
        'Session not found or expired.'
      );
    }

    const msalClient =
      this.sessionToMsal.get(sessionId);

    if (!msalClient) {
      throw new Error(
        'MSAL client not found for session.'
      );
    }

    const credentialInfo = this.sessionToCredential.get(sessionId);
    const scopes = credentialInfo?.scopes ?? this.defaultScopes;

    return this.acquireTokenByUserId(
      msalClient,
      userId,
      forceRefresh,
      scopes
    );
  }

  private async acquireTokenByUserId(
    msalClient: PublicClientApplication,
    userId: string,
    forceRefresh = false,
    scopes?: string[]
  ): Promise<string> {
    await this.loadUserCache(
      msalClient,
      userId
    );

    const accounts =
      await msalClient
        .getTokenCache()
        .getAllAccounts();

    const account = accounts.find(
      (a) =>
        a.homeAccountId === userId
    );

    if (!account) {
      throw new Error(
        'User not authenticated.'
      );
    }

    const requestedScopes = scopes ?? this.defaultScopes;
    try {
      const response =
        await msalClient.acquireTokenSilent({
          account,
          scopes: requestedScopes,
          forceRefresh,
        });

      return response.accessToken;
    } catch {
      if (forceRefresh) {
        throw new Error('Token refresh failed.');
      }

      return this.acquireTokenByUserId(
        msalClient,
        userId,
        true,
        requestedScopes
      );
    }
  }

  async clearSession(
    sessionId: string
  ): Promise<void> {
    const userId =
      this.sessionToUser.get(sessionId) ??
      (await loadSession(sessionId));

    if (userId) {
      await clearCache(userId);
    }

    this.sessionToUser.delete(
      sessionId
    );
    this.sessionToMsal.delete(
      sessionId
    );
    this.sessionToCredential.delete(sessionId);
    this.pendingDeviceFlows.delete(
      sessionId
    );

    await clearSessionFile(
      sessionId
    );
  }

  hasPendingDeviceFlow(
    sessionId: string
  ): boolean {
    return this.pendingDeviceFlows.has(
      sessionId
    );
  }

  getUserIdFromSession(
    sessionId: string
  ): string | undefined {
    return this.sessionToUser.get(
      sessionId
    );
  }

  getPendingFlowError(
    sessionId: string
  ): unknown {
    const flow =
      this.pendingDeviceFlows.get(
        sessionId
      );
    return flow?.error;
  }
}