import type { IncomingMessage, Server } from 'node:http';
import { WebSocketServer, WebSocket } from 'ws';
import { createLogger } from '@markus/shared';

const log = createLogger('ws-server');

export interface WSEvent {
  type: string;
  payload: unknown;
  timestamp: string;
}

const COALESCE_INTERVAL_MS = 200;
const COALESCE_EVENT_TYPES = new Set([
  'agent:update', 'agent:attention', 'agent:focus',
  'agent:mailbox', 'agent:decision', 'agent:triage',
]);

export class WSBroadcaster {
  private wss?: WebSocketServer;
  private clients = new Set<WebSocket>();
  /** Connections that identified with ?userId= (targeted events, e.g. notifications) */
  private userConnections = new Map<string, Set<WebSocket>>();
  /** Pending coalesced events keyed by `${type}:${agentId}` — only latest kept */
  private pendingCoalesced = new Map<string, WSEvent>();
  private coalesceTimer: ReturnType<typeof setTimeout> | null = null;

  attach(server: Server): void {
    this.wss = new WebSocketServer({ server, path: '/ws' });

    this.wss.on('connection', (ws, req: IncomingMessage) => {
      this.clients.add(ws);

      let trackedUserId: string | null = null;
      try {
        const host = req.headers.host ?? '127.0.0.1';
        const u = new URL(req.url ?? '/', `http://${host}`);
        const qUid = u.searchParams.get('userId');
        if (qUid) {
          trackedUserId = qUid;
          let set = this.userConnections.get(trackedUserId);
          if (!set) {
            set = new Set();
            this.userConnections.set(trackedUserId, set);
          }
          set.add(ws);
        }
      } catch {
        /* ignore malformed URL */
      }

      log.info('WebSocket client connected', {
        clients: this.clients.size,
        userScoped: !!trackedUserId,
      });

      const detachUser = () => {
        if (!trackedUserId) return;
        const set = this.userConnections.get(trackedUserId);
        if (!set) return;
        set.delete(ws);
        if (set.size === 0) this.userConnections.delete(trackedUserId);
      };

      ws.on('close', () => {
        this.clients.delete(ws);
        detachUser();
        log.debug('WebSocket client disconnected', { clients: this.clients.size });
      });

      ws.on('error', (err) => {
        log.error('WebSocket client error', { error: String(err) });
        this.clients.delete(ws);
        detachUser();
      });

      // Send initial ping
      ws.send(JSON.stringify({
        type: 'connected',
        payload: { message: 'Connected to Markus WebSocket' },
        timestamp: new Date().toISOString(),
      }));
    });

    log.info('WebSocket server attached');
  }

  sendToUser(userId: string, event: WSEvent): void {
    this.sendToUsers([userId], event);
  }

  sendToUsers(userIds: string[], event: WSEvent): void {
    const data = JSON.stringify(event);
    const seen = new Set<string>();
    for (const id of userIds) {
      if (seen.has(id)) continue;
      seen.add(id);
      const set = this.userConnections.get(id);
      if (!set) continue;
      for (const ws of set) {
        if (ws.readyState === WebSocket.OPEN) {
          ws.send(data);
        }
      }
    }
  }

  broadcast(event: WSEvent): void {
    const payload = event.payload as Record<string, unknown> | undefined;
    const agentId = payload?.['agentId'] as string | undefined;

    if (agentId && COALESCE_EVENT_TYPES.has(event.type)) {
      const key = `${event.type}:${agentId}`;
      this.pendingCoalesced.set(key, event);
      if (!this.coalesceTimer) {
        this.coalesceTimer = setTimeout(() => this.flushCoalesced(), COALESCE_INTERVAL_MS);
      }
      return;
    }

    this.broadcastImmediate(event);
  }

  private flushCoalesced(): void {
    this.coalesceTimer = null;
    const events = [...this.pendingCoalesced.values()];
    this.pendingCoalesced.clear();
    for (const event of events) {
      this.broadcastImmediate(event);
    }
  }

  private broadcastImmediate(event: WSEvent): void {
    const data = JSON.stringify(event);
    for (const ws of this.clients) {
      if (ws.readyState === WebSocket.OPEN) {
        ws.send(data);
      }
    }
  }

  broadcastAgentUpdate(agentId: string, status: string, extra?: Record<string, unknown>): void {
    this.broadcast({
      type: 'agent:update',
      payload: { agentId, status, ...extra },
      timestamp: new Date().toISOString(),
    });
  }

  broadcastTaskUpdate(taskId: string, status: string, extra?: Record<string, unknown>): void {
    this.broadcast({
      type: 'task:update',
      payload: { taskId, status, ...extra },
      timestamp: new Date().toISOString(),
    });
  }

  broadcastTaskCreate(taskId: string, title: string, extra?: Record<string, unknown>): void {
    this.broadcast({
      type: 'task:create',
      payload: { taskId, title, ...extra },
      timestamp: new Date().toISOString(),
    });
  }

  broadcastTeamUpdate(teamId: string, extra?: Record<string, unknown>): void {
    this.broadcast({
      type: 'team:update',
      payload: { teamId, ...extra },
      timestamp: new Date().toISOString(),
    });
  }

  broadcastDeliverableUpdate(deliverableId: string, action: 'created' | 'updated' | 'removed', extra?: Record<string, unknown>): void {
    this.broadcast({
      type: `deliverable:${action}`,
      payload: { deliverableId, ...extra },
      timestamp: new Date().toISOString(),
    });
  }

  broadcastChat(agentId: string, message: string, sender: 'user' | 'agent'): void {
    this.broadcast({
      type: 'chat:message',
      payload: { agentId, message, sender },
      timestamp: new Date().toISOString(),
    });
  }

  broadcastProactiveMessage(agentId: string, agentName: string, sessionId: string, messageId: string, message: string, metadata?: Record<string, unknown>, targetUserId?: string): void {
    const event = {
      type: 'chat:proactive_message' as const,
      payload: { agentId, agentName, sessionId, messageId, message, metadata, targetUserId },
      timestamp: new Date().toISOString(),
    };
    if (targetUserId) {
      this.sendToUser(targetUserId, event);
    } else {
      this.broadcast(event);
    }
  }

  broadcastUnreadUpdate(conversationKey: string, messageId: string): void {
    this.broadcast({
      type: 'chat:unread_update',
      payload: { conversationKey, messageId },
      timestamp: new Date().toISOString(),
    });
  }

  broadcastExecutionLog(entry: Record<string, unknown>): void {
    const ts = new Date().toISOString();
    this.broadcast({ type: 'execution:log', payload: entry, timestamp: ts });
  }

  broadcastExecutionLogDelta(sourceType: string, sourceId: string, agentId: string, text: string): void {
    const ts = new Date().toISOString();
    this.broadcast({ type: 'execution:log:delta', payload: { sourceType, sourceId, agentId, text }, timestamp: ts });
  }

  getClientCount(): number {
    return this.clients.size;
  }
}
