import { WebSocketServer, WebSocket } from 'ws';
import type { Server } from 'node:http';
import { nanoid } from 'nanoid';
import type { AuthManager } from './auth.js';
import type { SessionManager } from '../bridge/session-manager.js';
import type { ScreenPoller } from '../bridge/screen-poller.js';
import type { InputHandler } from '../bridge/input-handler.js';
import type { CmuxClient } from '../bridge/cmux-client.js';
import type { ClientMessage, ServerMessage } from '../protocol/messages.js';

interface ClientInfo {
  id: string;
  ws: WebSocket;
  lastPong: number;
}

export class WsServer {
  private wss: WebSocketServer;
  private clients = new Map<string, ClientInfo>();
  private auth: AuthManager;
  private sessions: SessionManager;
  private poller: ScreenPoller;
  private input: InputHandler;
  private cmux: CmuxClient;
  private verbose: boolean;
  private heartbeatInterval: ReturnType<typeof setInterval>;

  constructor(opts: {
    server: Server;
    auth: AuthManager;
    sessions: SessionManager;
    poller: ScreenPoller;
    input: InputHandler;
    cmux: CmuxClient;
    verbose?: boolean;
  }) {
    this.auth = opts.auth;
    this.sessions = opts.sessions;
    this.poller = opts.poller;
    this.input = opts.input;
    this.cmux = opts.cmux;
    this.verbose = opts.verbose ?? false;

    this.wss = new WebSocketServer({
      server: opts.server,
      perMessageDeflate: true,
      maxPayload: 1024 * 1024, // 1MB limit
    });

    this.wss.on('connection', (ws) => this.handleConnection(ws));
    this.heartbeatInterval = setInterval(() => this.checkHeartbeats(), 5000);

    // Listen for screen updates from poller
    this.poller.on('update', (update) => this.broadcastUpdate(update));
    this.poller.on('browser-update', (update) => this.broadcastBrowserUpdate(update));
  }

  private handleConnection(ws: WebSocket): void {
    const clientId = nanoid(12);
    this.clients.set(clientId, { id: clientId, ws, lastPong: Date.now() });
    this.sessions.addClient(clientId);

    if (this.verbose) console.log(`[ws] Client connected: ${clientId}`);

    ws.on('message', (data) => {
      let msg: ClientMessage;
      try {
        msg = JSON.parse(data.toString()) as ClientMessage;
      } catch {
        this.sendTo(clientId, { type: 'error', message: 'Invalid message format' });
        return;
      }
      this.handleMessage(clientId, msg).catch((err) => {
        if (this.verbose) console.error(`[ws] Error handling ${msg.type}:`, err.message);
      });
    });

    ws.on('close', () => {
      if (this.verbose) console.log(`[ws] Client disconnected: ${clientId}`);
      this.sessions.removeClient(clientId);
      this.clients.delete(clientId);
    });

    ws.on('error', (err) => {
      console.error(`[ws] Client error ${clientId}:`, err.message);
    });
  }

  private async handleMessage(clientId: string, msg: ClientMessage): Promise<void> {
    // Heartbeat pong — handle before auth check
    if (msg.type === 'pong') {
      const client = this.clients.get(clientId);
      if (client) client.lastPong = Date.now();
      return;
    }

    // Auth messages don't require prior authentication
    if (msg.type === 'auth') {
      const valid = this.auth.verifyToken(msg.token);
      if (valid) {
        this.sessions.authenticateClient(clientId);
        this.sendTo(clientId, { type: 'auth-ok' });
      } else {
        this.sendTo(clientId, { type: 'auth-fail', message: 'Invalid token' });
      }
      return;
    }

    // All other messages require authentication
    if (!this.sessions.isAuthenticated(clientId)) {
      this.sendTo(clientId, { type: 'auth-fail', message: 'Not authenticated' });
      return;
    }

    switch (msg.type) {
      case 'subscribe':
        this.sessions.subscribe(clientId, msg.workspace, msg.surface);
        if (this.verbose) console.log(`[ws] ${clientId} subscribed to ${msg.surface}`);
        // Send cached content immediately so client doesn't wait for next poll cycle
        {
          const lastLines = this.poller.getLastLines(msg.workspace, msg.surface);
          if (lastLines) {
            this.sendTo(clientId, {
              type: 'screen',
              surface: msg.surface,
              content: lastLines.join('\n'),
              lines: lastLines,
            });
          }
        }
        break;

      case 'unsubscribe':
        this.sessions.unsubscribe(clientId, msg.surface);
        break;

      case 'send-text':
        if (this.verbose) console.log(`[ws] ${clientId} send-text to ${msg.surface}`);
        this.poller.resetToFast(msg.surface);
        await this.input.handleText(msg.workspace, msg.surface, msg.text);
        if (msg.submit) {
          // Await-ordered submit: send_text has now fully completed (bytes are
          // in the PTY), so the Enter cannot race ahead of the bracketed paste
          // the way a client-side timer could. The short local delay covers the
          // TUI's async ingestion of the paste before the key event (send_key
          // uses ghostty key-injection, which is not PTY-ordered against bytes).
          await new Promise((resolve) => setTimeout(resolve, 200));
          await this.input
            .handleKey(msg.workspace, msg.surface, 'Enter')
            .catch((err) => {
              if (this.verbose) console.error(`[ws] submit Enter error:`, err);
            });
        }
        break;

      case 'send-key':
        this.poller.resetToFast(msg.surface);
        await this.input.handleKey(msg.workspace, msg.surface, msg.key).catch((err) => {
          if (this.verbose) console.error(`[ws] send-key error:`, err);
        });
        break;

      case 'list-workspaces': {
        const workspaces = await this.cmux.listWorkspaces();
        this.sendTo(clientId, { type: 'workspaces', workspaces });
        break;
      }

      case 'list-surfaces': {
        const surfaces = await this.cmux.listPaneSurfaces(msg.workspace);
        // Wrap in workspace info format
        const workspaces = await this.cmux.listWorkspaces();
        this.sendTo(clientId, { type: 'workspaces', workspaces });
        break;
      }

      case 'scroll-request': {
        // Use capture-pane for scrollback history (read-screen only shows current visible area)
        const content = await this.cmux.capturePaneScrollback(msg.workspace, msg.surface, msg.lines);
        const lines = content.split('\n');
        this.sendTo(clientId, {
          type: 'screen',
          surface: msg.surface,
          content,
          lines,
          scrollback: true,
        });
        break;
      }

      default:
        this.sendTo(clientId, { type: 'error', message: `Unknown message type` });
    }
  }

  private broadcastUpdate(update: {
    surface: string;
    workspace: string;
    full: boolean;
    lines: string[];
    patches?: Array<{ startLine: number; deleteCount: number; insertLines: string[] }>;
  }): void {
    const subscribedClients = this.sessions.getSubscribedClients(update.surface);

    for (const clientId of subscribedClients) {
      if (update.full || !update.patches) {
        this.sendTo(clientId, {
          type: 'screen',
          surface: update.surface,
          content: update.lines.join('\n'),
          lines: update.lines,
        });
      } else {
        this.sendTo(clientId, {
          type: 'screen-diff',
          surface: update.surface,
          patches: update.patches,
        });
      }
    }
  }

  private broadcastBrowserUpdate(update: { surface: string; data: string; mime: string }): void {
    for (const clientId of this.sessions.getSubscribedClients(update.surface)) {
      this.sendTo(clientId, {
        type: 'browser-screenshot',
        surface: update.surface,
        imageData: update.data,
        mime: update.mime,
      });
    }
  }

  private sendTo(clientId: string, msg: ServerMessage): void {
    const client = this.clients.get(clientId);
    if (client && client.ws.readyState === WebSocket.OPEN) {
      client.ws.send(JSON.stringify(msg));
    }
  }

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

  private checkHeartbeats(): void {
    const now = Date.now();
    for (const [clientId, client] of this.clients) {
      if (now - client.lastPong > 15000) {
        if (this.verbose) console.log(`[ws] Heartbeat timeout: ${clientId}`);
        client.ws.terminate();
      } else {
        this.sendTo(clientId, { type: 'ping' });
      }
    }
  }

  close(): void {
    clearInterval(this.heartbeatInterval);
    this.wss.close();
  }
}
