/**
 * Generic account subscription manager using Solana programNotifications.
 *
 * Framework-agnostic — no React, no Sentry, no webapp-specific deps.
 * Consumers provide error/log callbacks for observability.
 *
 * Features:
 * - Ref-counted: first listener starts the WS, last unsubscribe tears it down
 * - Slot-based ordering: rejects out-of-order notifications
 * - Buffer dedup: skips unchanged data (updates slot without firing callbacks)
 * - Auto-reconnect: retries on all WS errors with configurable delay
 * - Per-subscription onReconnect hook for state refresh after WS recovery
 */
import {
  type Address,
  type Base58EncodedBytes,
  type Commitment,
  type RpcSubscriptions,
  type Slot,
  type SolanaRpcSubscriptionsApi,
} from '@solana/kit';
import { Buffer } from 'buffer';

import { blobEquals, noop, sleep, toJson } from './utils';

export const WS_RECONNECT_DELAY_MS = 500;

export type SubscriptionCallback = (address: Address, buffer: Uint8Array, slot: Slot) => void;
export type SubscriptionMode = 'ws' | 'idle';
export type SubscriptionErrorHandler = (error: unknown) => void;
export type SubscriptionLogger = (...args: unknown[]) => void;

export type ProgramFilter =
  | Readonly<{ dataSize: bigint }>
  | Readonly<{ memcmp: Readonly<{ offset: bigint; bytes: Base58EncodedBytes; encoding: 'base58' }> }>;

export interface AccountSubscriptionManagerConfig {
  wsRpc: RpcSubscriptions<SolanaRpcSubscriptionsApi>;
  onError?: SubscriptionErrorHandler;
  onLog?: SubscriptionLogger;
  reconnectDelayMs?: number;
}

export interface SubscribeProgramAccountsOptions {
  commitment?: Commitment;
  /** Called after WS reconnects (before resuming the notification loop).
   *  Use this to refresh local state that may have gone stale during downtime.
   *
   *  Note: onReconnect is per-subscription-key, not per-listener. If multiple
   *  listeners share the same key, the last one to provide onReconnect wins. */
  onReconnect?: () => void | Promise<void>;
  /** Throttle callback invocations to at most once per N milliseconds.
   *  Useful for high-frequency accounts like reserves that update every slot.
   *  When set, the latest notification is buffered and delivered after the interval.
   *  Not recommended for obligations (low-frequency, want instant delivery).
   *
   *  Note: throttleMs is set once per subscription key (first subscriber wins).
   *  Subsequent subscribers to the same key inherit the existing throttle. */
  throttleMs?: number;
}

interface ManagedSubscription {
  key: string;
  programId: Address;
  filters: readonly ProgramFilter[];
  /** Pre-serialized filters for O(1) lookups in findSubscription */
  serializedFilters: string;
  commitment: Commitment;
  listeners: Map<number, SubscriptionCallback>;
  buffers: Map<Address, { buffer: Uint8Array; slot: Slot }>;
  mode: SubscriptionMode;
  abortController: AbortController | undefined;
  onReconnect?: () => void | Promise<void>;
  throttleMs?: number;
  throttleTimer?: ReturnType<typeof setTimeout>;
  pendingNotifications: Map<Address, { buffer: Uint8Array; slot: Slot }>;
  /** Resolves when the WS subscription is first established */
  connectedResolve?: () => void;
  connectedPromise?: Promise<void>;
}

export class AccountSubscriptionManager {
  private subscriptions = new Map<string, ManagedSubscription>();
  private nextListenerId = 0;
  private wsRpc: RpcSubscriptions<SolanaRpcSubscriptionsApi>;
  private onError: SubscriptionErrorHandler;
  private onLog: SubscriptionLogger;
  private reconnectDelayMs: number;

  constructor(config: AccountSubscriptionManagerConfig) {
    this.wsRpc = config.wsRpc;
    this.onError = config.onError ?? noop;
    this.onLog = config.onLog ?? noop;
    this.reconnectDelayMs = config.reconnectDelayMs ?? WS_RECONNECT_DELAY_MS;
  }

  /**
   * Subscribe to program account changes. Returns an unsubscribe function.
   *
   * The first subscriber for a given key starts the WS connection; the last
   * unsubscribe tears it down. Late subscribers immediately receive any
   * cached account data from previous notifications.
   */
  subscribeProgramAccounts(
    programId: Address,
    filters: readonly ProgramFilter[],
    callback: SubscriptionCallback,
    options?: SubscribeProgramAccountsOptions
  ): () => void {
    const commitment = options?.commitment ?? 'processed';
    const key = this.buildKey(programId, commitment, filters);
    const listenerId = this.nextListenerId++;

    const sub =
      this.subscriptions.get(key) ?? this.createSubscription(key, programId, filters, commitment, options?.throttleMs);
    sub.listeners.set(listenerId, callback);
    if (options?.onReconnect) sub.onReconnect = options.onReconnect;

    if (sub.listeners.size === 1) this.startWs(sub);

    for (const [addr, cached] of sub.buffers) {
      try {
        callback(addr, cached.buffer, cached.slot);
      } catch (e) {
        this.onError(e);
      }
    }

    return () => {
      const current = this.subscriptions.get(key);
      if (!current) return;
      current.listeners.delete(listenerId);
      if (current.listeners.size === 0) {
        this.teardown(current);
        this.subscriptions.delete(key);
      }
    };
  }

  /**
   * Immediately flush any throttled pending notifications for the subscription
   * matching the given programId + filters. Resets the throttle timer so the
   * next notification after flush waits the full interval again.
   *
   * No-op if the subscription doesn't exist or has no pending notifications.
   */
  /**
   * Immediately flush any throttled pending notifications.
   * When commitment is provided, flushes only that specific subscription.
   * When omitted, flushes all matching programId+filters across all commitments.
   */
  flushPending(programId: Address, filters: readonly ProgramFilter[], commitment?: Commitment): void {
    if (commitment) {
      const key = this.buildKey(programId, commitment, filters);
      const sub = this.subscriptions.get(key);
      if (sub) this.flushSub(sub);
    } else {
      const serializedFilters = toJson(filters, true);
      for (const sub of this.subscriptions.values()) {
        if (sub.programId === programId && sub.serializedFilters === serializedFilters) {
          this.flushSub(sub);
        }
      }
    }
  }

  getMode(programId: Address, filters: readonly ProgramFilter[], commitment?: Commitment): SubscriptionMode {
    return this.findSubscription(programId, filters, commitment)?.mode ?? 'idle';
  }

  /** Tear down all subscriptions. Useful for shutdown and test cleanup. */
  destroy(): void {
    for (const [key, sub] of this.subscriptions) {
      this.teardown(sub);
      this.subscriptions.delete(key);
    }
  }

  /** Resolves when the WS subscription for the given programId + filters is established. */
  async waitForConnection(
    programId: Address,
    filters: readonly ProgramFilter[],
    commitment?: Commitment
  ): Promise<void> {
    const sub = this.findSubscription(programId, filters, commitment);
    if (!sub) return;
    await sub.connectedPromise;
  }

  private buildKey(programId: Address, commitment: string, filters: readonly ProgramFilter[]): string {
    return `${programId}:${commitment}:${toJson(filters, true)}`;
  }

  /**
   * Find a subscription by programId + filters.
   * When commitment is provided, does a direct Map lookup (O(1)).
   * When omitted, scans all commitments and returns the first match.
   */
  private findSubscription(
    programId: Address,
    filters: readonly ProgramFilter[],
    commitment?: Commitment
  ): ManagedSubscription | undefined {
    if (commitment) {
      return this.subscriptions.get(this.buildKey(programId, commitment, filters));
    }
    const serializedFilters = toJson(filters, true);
    for (const sub of this.subscriptions.values()) {
      if (sub.programId === programId && sub.serializedFilters === serializedFilters) {
        return sub;
      }
    }
    return undefined;
  }

  private flushSub(sub: ManagedSubscription): void {
    if (sub.pendingNotifications.size === 0) return;
    this.onLog('[WS] Flushing', sub.pendingNotifications.size, 'throttled notifications');
    if (sub.throttleTimer) {
      clearTimeout(sub.throttleTimer);
      sub.throttleTimer = undefined;
    }
    this.flushPendingNotifications(sub);
  }

  private createSubscription(
    key: string,
    programId: Address,
    filters: readonly ProgramFilter[],
    commitment: Commitment,
    throttleMs?: number
  ): ManagedSubscription {
    let connectedResolve: (() => void) | undefined;
    const connectedPromise = new Promise<void>((resolve) => {
      connectedResolve = resolve;
    });
    const sub: ManagedSubscription = {
      key,
      programId,
      filters,
      serializedFilters: toJson(filters, true),
      commitment,
      listeners: new Map(),
      buffers: new Map(),
      mode: 'idle',
      abortController: undefined,
      throttleMs,
      pendingNotifications: new Map(),
      connectedResolve,
      connectedPromise,
    };
    this.subscriptions.set(key, sub);
    return sub;
  }

  private startWs(sub: ManagedSubscription): void {
    const abortController = new AbortController();
    sub.abortController = abortController;
    sub.mode = 'ws';

    const wsLoop = async () => {
      let isFirstConnect = true;
      while (!abortController.signal.aborted) {
        try {
          // On reconnect (not first connect), call the hook so consumers can
          // refresh local state that went stale during the WS downtime.
          // If onReconnect throws, the catch block handles it (logs via onError,
          // sleeps, retries). This is intentional — a failed state refresh should
          // not prevent the WS from reconnecting on the next loop iteration.
          if (!isFirstConnect && sub.onReconnect) {
            this.onLog('[WS] Calling onReconnect before re-subscribing');
            await sub.onReconnect();
          }
          isFirstConnect = false;

          this.onLog('[WS] Connecting programSubscribe for', sub.programId);
          const notifications = await this.wsRpc
            .programNotifications(sub.programId, {
              commitment: sub.commitment,
              encoding: 'base64',
              filters: sub.filters,
            })
            .subscribe({ abortSignal: abortController.signal });

          this.onLog('[WS] Subscription established, listening...');
          sub.connectedResolve?.();
          sub.connectedResolve = undefined;
          for await (const notification of notifications) {
            const { pubkey, account } = notification.value as {
              pubkey: Address;
              account: { data: [string, string] };
            };
            if (!account) continue;

            this.handleNotification(sub, pubkey, Buffer.from(account.data[0], 'base64'), notification.context.slot);
          }
          // Stream ended cleanly (server closed) — treat as disconnect, sleep before retry
          this.onLog('[WS] Stream ended cleanly, reconnecting...');
          await sleep(this.reconnectDelayMs);
        } catch (e: unknown) {
          if (abortController.signal.aborted) break;
          this.onError(e);
          await sleep(this.reconnectDelayMs);
        }
      }
      sub.mode = 'idle';
    };
    wsLoop().catch(this.onError);
  }

  private handleNotification(sub: ManagedSubscription, addr: Address, buffer: Uint8Array, slot: Slot): void {
    const cached = sub.buffers.get(addr);

    if (cached && slot < cached.slot) return;
    if (cached && blobEquals(cached.buffer, buffer)) {
      cached.slot = slot;
      return;
    }

    sub.buffers.set(addr, { buffer, slot });

    if (sub.throttleMs) {
      sub.pendingNotifications.set(addr, { buffer, slot });
      if (!sub.throttleTimer) {
        sub.throttleTimer = setTimeout(() => {
          sub.throttleTimer = undefined;
          this.flushPendingNotifications(sub);
        }, sub.throttleMs);
      }
    } else {
      this.fireCallbacks(sub, addr, buffer, slot);
    }
  }

  private flushPendingNotifications(sub: ManagedSubscription): void {
    for (const [addr, { buffer, slot }] of sub.pendingNotifications) {
      this.fireCallbacks(sub, addr, buffer, slot);
    }
    sub.pendingNotifications.clear();
  }

  private fireCallbacks(sub: ManagedSubscription, addr: Address, buffer: Uint8Array, slot: Slot): void {
    for (const cb of sub.listeners.values()) {
      try {
        cb(addr, buffer, slot);
      } catch (e) {
        this.onError(e);
      }
    }
  }

  private teardown(sub: ManagedSubscription): void {
    sub.abortController?.abort();
    sub.abortController = undefined;
    if (sub.throttleTimer) clearTimeout(sub.throttleTimer);
    sub.throttleTimer = undefined;
    sub.pendingNotifications.clear();
    sub.buffers.clear();
    sub.mode = 'idle';
  }
}
