import type { ShadeEventEmitter, ShadeEvent, ShadeSessionManager } from '@shade/core';
import type { PrekeyServerEvents, PrekeyServerEvent, PrekeyStore } from '@shade/server';

/**
 * Aggregated observer state, updated as events flow in.
 *
 * The observer maintains a rolling snapshot of:
 *  - Identity (fingerprint, registration ID)
 *  - Active sessions (per address: message counts, last activity)
 *  - Prekey stock
 *  - Server stats (registered identities, fetches, replenishes)
 *  - Recent events ring buffer
 */
export interface ObserverSnapshot {
  identity: {
    fingerprint: string | null;
    registrationId: number | null;
    lastInitialized: number | null;
    lastRotated: number | null;
  };
  sessions: Array<{
    address: string;
    remoteIdentityKeyHash: string;
    messageCountSent: number;
    messageCountReceived: number;
    lastActivity: number;
    dhRatchetSteps: number;
  }>;
  prekeys: {
    oneTimeRemaining: number;
    lastGenerated: number | null;
    lastConsumed: number | null;
    signedPreKeyId: number | null;
    signedPreKeyLastRotated: number | null;
  };
  retiredIdentities: number;
  server: {
    registeredIdentities: Set<string>;
    totalBundleFetches: number;
    totalReplenishes: number;
    totalDeleted: number;
    totalRateLimited: number;
  };
}

interface SessionStats {
  remoteIdentityKeyHash: string;
  messageCountSent: number;
  messageCountReceived: number;
  lastActivity: number;
  dhRatchetSteps: number;
}

export class StateAggregator {
  private identity: ObserverSnapshot['identity'] = {
    fingerprint: null,
    registrationId: null,
    lastInitialized: null,
    lastRotated: null,
  };
  private sessions = new Map<string, SessionStats>();
  private prekeys: ObserverSnapshot['prekeys'] = {
    oneTimeRemaining: 0,
    lastGenerated: null,
    lastConsumed: null,
    signedPreKeyId: null,
    signedPreKeyLastRotated: null,
  };
  private retiredIdentities = 0;
  private serverStats = {
    registeredIdentities: new Set<string>(),
    totalBundleFetches: 0,
    totalReplenishes: 0,
    totalDeleted: 0,
    totalRateLimited: 0,
  };

  constructor(
    clientEvents?: ShadeEventEmitter,
    serverEvents?: PrekeyServerEvents,
    _manager?: ShadeSessionManager,
    _store?: PrekeyStore,
  ) {
    if (clientEvents) {
      // Replay any events that fired before subscription, then subscribe
      for (const e of clientEvents.getBufferedSince(0)) {
        this.handleClientEvent(e);
      }
      clientEvents.on((e) => this.handleClientEvent(e));
    }
    if (serverEvents) {
      for (const e of serverEvents.getBufferedSince(0)) {
        this.handleServerEvent(e);
      }
      serverEvents.on((e) => this.handleServerEvent(e));
    }
  }

  private handleClientEvent(e: ShadeEvent): void {
    switch (e.name) {
      case 'identity.initialized':
        this.identity.fingerprint = e.data.fingerprint;
        this.identity.registrationId = e.data.registrationId;
        this.identity.lastInitialized = e.timestamp;
        break;
      case 'identity.rotated':
        this.identity.fingerprint = e.data.newFingerprint;
        this.identity.lastRotated = e.timestamp;
        this.retiredIdentities++;
        break;
      case 'session.created':
        this.sessions.set(e.data.address, {
          remoteIdentityKeyHash: e.data.remoteIdentityKeyHash,
          messageCountSent: 0,
          messageCountReceived: 0,
          lastActivity: e.timestamp,
          dhRatchetSteps: 0,
        });
        break;
      case 'session.removed':
        this.sessions.delete(e.data.address);
        break;
      case 'message.encrypted': {
        const s = this.sessions.get(e.data.address);
        if (s) {
          s.messageCountSent++;
          s.lastActivity = e.timestamp;
        }
        break;
      }
      case 'message.decrypted': {
        const s = this.sessions.get(e.data.address);
        if (s) {
          s.messageCountReceived++;
          s.lastActivity = e.timestamp;
        }
        break;
      }
      case 'ratchet.dh_step': {
        const s = this.sessions.get(e.data.address);
        if (s) s.dhRatchetSteps++;
        break;
      }
      case 'prekey.generated':
        this.prekeys.oneTimeRemaining = e.data.totalAfter;
        this.prekeys.lastGenerated = e.timestamp;
        break;
      case 'prekey.consumed':
        if (this.prekeys.oneTimeRemaining > 0) this.prekeys.oneTimeRemaining--;
        this.prekeys.lastConsumed = e.timestamp;
        break;
      case 'signed_prekey.rotated':
        this.prekeys.signedPreKeyId = e.data.newKeyId;
        this.prekeys.signedPreKeyLastRotated = e.timestamp;
        break;
      // trust.* don't directly affect snapshot but appear in event feed
    }
  }

  private handleServerEvent(e: PrekeyServerEvent): void {
    switch (e.name) {
      case 'server.identity_registered':
        this.serverStats.registeredIdentities.add(e.data.address);
        break;
      case 'server.bundle_fetched':
        this.serverStats.totalBundleFetches++;
        break;
      case 'server.prekeys_replenished':
        this.serverStats.totalReplenishes++;
        break;
      case 'server.identity_deleted':
        this.serverStats.registeredIdentities.delete(e.data.address);
        this.serverStats.totalDeleted++;
        break;
      case 'server.rate_limited':
        this.serverStats.totalRateLimited++;
        break;
    }
  }

  /** Get current snapshot */
  snapshot(): ObserverSnapshot {
    return {
      identity: { ...this.identity },
      sessions: Array.from(this.sessions.entries()).map(([address, s]) => ({
        address,
        ...s,
      })),
      prekeys: { ...this.prekeys },
      retiredIdentities: this.retiredIdentities,
      server: {
        registeredIdentities: new Set(this.serverStats.registeredIdentities),
        totalBundleFetches: this.serverStats.totalBundleFetches,
        totalReplenishes: this.serverStats.totalReplenishes,
        totalDeleted: this.serverStats.totalDeleted,
        totalRateLimited: this.serverStats.totalRateLimited,
      },
    };
  }

  /** Snapshot with serializable JSON (Set → array) */
  toJSON(): any {
    const s = this.snapshot();
    return {
      ...s,
      server: {
        ...s.server,
        registeredIdentities: Array.from(s.server.registeredIdentities),
      },
    };
  }
}
