/**
 * RemoteServer - Stream Hypen apps over WebSocket with Session Management
 *
 * Usage:
 * ```typescript
 * import { RemoteServer, app } from "@hypen-space/core";
 *
 * const counter = app
 *   .defineState({ count: 0 })
 *   .onAction("increment", ({ state }) => state.count++)
 *   .onDisconnect(async ({ state, session }) => {
 *     await redis.set(`session:${session.id}`, JSON.stringify(state));
 *   })
 *   .onReconnect(async ({ session, restore }) => {
 *     const saved = await redis.get(`session:${session.id}`);
 *     if (saved) restore(JSON.parse(saved));
 *   })
 *   .onExpire(async ({ session }) => {
 *     await redis.del(`session:${session.id}`);
 *   })
 *   .build();
 *
 * new RemoteServer()
 *   .module("Counter", counter)
 *   .ui(`Column { Text("Count: \${state.count}") }`)
 *   .session({ ttl: 3600, concurrent: "kick-old" })
 *   .listen(3000);
 * ```
 */

import { Engine } from "../engine.js";
import { HypenModuleInstance } from "../app.js";
import type { HypenModule } from "../app.js";
import type {
  RemoteMessage,
  RemoteClient,
  RemoteServerConfig,
  InitialTreeMessage,
  PatchMessage,
  DispatchActionMessage,
  HelloMessage,
  SessionAckMessage,
  SessionExpiredMessage,
  Session,
  SessionConfig,
} from "./types.js";
import type { Patch } from "../types.js";
import type { ServerWebSocket } from "bun";
import { SessionManager } from "./session.js";
import { frameworkLoggers } from "../logger.js";

const log = frameworkLoggers.remote;

interface ClientData {
  id: string;
  engine: Engine;
  moduleInstance: HypenModuleInstance<any>;
  revision: number;
  connectedAt: Date;
  /** Session ID for this client */
  sessionId: string;
  /** Whether we've received the hello message */
  helloReceived: boolean;
  /** Timeout for legacy clients that don't send hello */
  helloTimeout?: ReturnType<typeof setTimeout>;
}

/**
 * Builder pattern for hosting Hypen apps over WebSocket
 */
export class RemoteServer {
  private _module: HypenModule<any> | null = null;
  private _moduleName: string = "App";
  private _ui: string = "";
  private _config: RemoteServerConfig = {};
  private _sessionConfig: SessionConfig = {};
  private _onConnectionCallbacks: Array<(client: RemoteClient) => void> = [];
  private _onDisconnectionCallbacks: Array<(client: RemoteClient) => void> = [];
  private clients = new Map<ServerWebSocket<unknown>, ClientData>();
  private nextClientId = 1;
  private server: ReturnType<typeof Bun.serve> | null = null;
  private sessionManager: SessionManager | null = null;

  /**
   * Set the module for this app
   */
  module(name: string, module: HypenModule<any>): this {
    this._moduleName = name;
    this._module = module;
    return this;
  }

  /**
   * Set the UI DSL string
   */
  ui(dsl: string): this {
    this._ui = dsl;
    return this;
  }

  /**
   * Set server configuration
   */
  config(config: RemoteServerConfig): this {
    this._config = { ...this._config, ...config };
    return this;
  }

  /**
   * Configure session management
   */
  session(config: SessionConfig): this {
    this._sessionConfig = config;
    return this;
  }

  /**
   * Register connection callback
   */
  onConnection(callback: (client: RemoteClient) => void): this {
    this._onConnectionCallbacks.push(callback);
    return this;
  }

  /**
   * Register disconnection callback
   */
  onDisconnection(callback: (client: RemoteClient) => void): this {
    this._onDisconnectionCallbacks.push(callback);
    return this;
  }

  /**
   * Start the WebSocket server
   */
  listen(port?: number): this {
    if (!this._module) {
      throw new Error("Module not set. Call .module() before .listen()");
    }

    if (!this._ui) {
      throw new Error("UI not set. Call .ui() before .listen()");
    }

    // Initialize session manager
    this.sessionManager = new SessionManager(this._sessionConfig);

    const finalPort = port ?? this._config.port ?? 3000;
    const hostname = this._config.hostname ?? "0.0.0.0";

    this.server = Bun.serve({
      port: finalPort,
      hostname,
      websocket: {
        open: (ws) => this.handleOpen(ws),
        message: (ws, message) => this.handleMessage(ws, message),
        close: (ws) => this.handleClose(ws),
      },
      fetch: (req, server) => {
        const url = new URL(req.url);

        // Upgrade to WebSocket
        if (server.upgrade(req, { data: undefined })) {
          return; // Connection upgraded
        }

        // Health check endpoint
        if (url.pathname === "/health") {
          return new Response("OK", { status: 200 });
        }

        // Stats endpoint
        if (url.pathname === "/stats") {
          const stats = this.sessionManager?.getStats() ?? {
            activeSessions: 0,
            pendingSessions: 0,
            totalConnections: 0,
          };
          return new Response(JSON.stringify(stats), {
            headers: { "Content-Type": "application/json" },
          });
        }

        return new Response("Hypen Remote Server", { status: 200 });
      },
    });

    log.info(`Hypen app streaming on ws://${hostname}:${finalPort}`);

    return this;
  }

  /**
   * Stop the server
   */
  stop(): void {
    if (this.server) {
      this.server.stop();
      this.server = null;
    }
    if (this.sessionManager) {
      this.sessionManager.destroy();
      this.sessionManager = null;
    }
  }

  /**
   * Get the server URL
   */
  get url(): string | null {
    if (!this.server) return null;
    const hostname = this._config.hostname ?? "localhost";
    const port = this._config.port ?? 3000;
    return `ws://${hostname}:${port}`;
  }

  /**
   * Handle new WebSocket connection
   * Waits for hello message before fully initializing
   */
  private async handleOpen(ws: ServerWebSocket<unknown>) {
    try {
      const clientId = `client_${this.nextClientId++}`;
      const connectedAt = new Date();

      // Create engine instance for this client
      const engine = new Engine();
      await engine.init();

      // Create module instance
      const moduleInstance = new HypenModuleInstance(engine, this._module!);

      // Store client data (session will be set when hello is received)
      const clientData: ClientData = {
        id: clientId,
        engine,
        moduleInstance,
        revision: 0,
        connectedAt,
        sessionId: "",
        helloReceived: false,
      };

      this.clients.set(ws, clientData);

      // Set timeout for legacy clients that don't send hello
      clientData.helloTimeout = setTimeout(() => {
        if (!clientData.helloReceived) {
          // Legacy client - create new session automatically
          this.initializeSession(ws, clientData, undefined, undefined);
        }
      }, 1000); // 1 second grace period

    } catch (error) {
      log.error("Error handling WebSocket open:", error);
      ws.close(1011, "Internal server error");
    }
  }

  /**
   * Initialize session for a client (new or resumed)
   */
  private async initializeSession(
    ws: ServerWebSocket<unknown>,
    clientData: ClientData,
    requestedSessionId: string | undefined,
    props: Record<string, any> | undefined
  ): Promise<void> {
    if (clientData.helloReceived) return;
    clientData.helloReceived = true;

    // Clear hello timeout
    if (clientData.helloTimeout) {
      clearTimeout(clientData.helloTimeout);
      clientData.helloTimeout = undefined;
    }

    let session: Session;
    let isNew = true;
    let isRestored = false;
    let restoredState: unknown = null;

    if (requestedSessionId && this.sessionManager) {
      // Try to resume pending session
      const resumed = this.sessionManager.resumeSession(requestedSessionId);
      if (resumed) {
        session = resumed.session;
        restoredState = resumed.savedState;
        isNew = false;
        isRestored = true;

        // Call onReconnect hook
        await this.triggerReconnect(clientData, session, restoredState);
      } else {
        // Check for concurrent active session
        const activeSession = this.sessionManager.getActiveSession(requestedSessionId);
        if (activeSession) {
          const handled = await this.handleConcurrentConnection(
            ws,
            clientData,
            activeSession,
            props
          );
          if (!handled) return; // Connection was rejected
          session = activeSession;
          isNew = false;
        } else {
          // Session not found, create new
          session = this.sessionManager.createSession(props);
        }
      }
    } else if (this.sessionManager) {
      // New session
      session = this.sessionManager.createSession(props);
    } else {
      // No session manager (shouldn't happen, but fallback)
      session = {
        id: crypto.randomUUID(),
        ttl: 3600,
        createdAt: new Date(),
        lastConnectedAt: new Date(),
        props,
      };
    }

    clientData.sessionId = session.id;
    this.sessionManager?.trackConnection(session.id, ws);

    // Send session acknowledgment
    const sessionAck: SessionAckMessage = {
      type: "sessionAck",
      sessionId: session.id,
      isNew,
      isRestored,
    };
    ws.send(JSON.stringify(sessionAck));

    // Set up render callback
    this.setupRenderCallback(ws, clientData);

    // Render initial tree
    const initialPatches: Patch[] = [];
    clientData.engine.setRenderCallback((patches) => {
      initialPatches.push(...patches);
    });
    clientData.engine.renderSource(this._ui);

    // Now set up streaming render callback
    this.setupRenderCallback(ws, clientData);

    // Send initial tree
    const initialMessage: InitialTreeMessage = {
      type: "initialTree",
      module: this._moduleName,
      state: clientData.moduleInstance.getState(),
      patches: initialPatches,
      revision: 0,
    };
    ws.send(JSON.stringify(initialMessage));

    // Notify connection callbacks
    const client: RemoteClient = {
      id: clientData.id,
      socket: ws,
      connectedAt: clientData.connectedAt,
    };
    this._onConnectionCallbacks.forEach((cb) => cb(client));
  }

  /**
   * Set up the render callback for streaming patches
   */
  private setupRenderCallback(
    ws: ServerWebSocket<unknown>,
    clientData: ClientData
  ): void {
    clientData.engine.setRenderCallback((patches) => {
      const data = this.clients.get(ws);
      if (!data) return;

      data.revision++;

      const message: PatchMessage = {
        type: "patch",
        module: this._moduleName,
        patches,
        revision: data.revision,
      };

      ws.send(JSON.stringify(message));

      // If allow-multiple, broadcast to other connections on same session
      if (
        this.sessionManager?.getConcurrentPolicy() === "allow-multiple" &&
        data.sessionId
      ) {
        const connections = this.sessionManager.getConnections(data.sessionId);
        if (connections) {
          for (const conn of connections) {
            if (conn !== ws) {
              (conn as ServerWebSocket<unknown>).send(JSON.stringify(message));
            }
          }
        }
      }
    });
  }

  /**
   * Handle concurrent connection based on policy
   * Returns true if connection is allowed, false if rejected
   */
  private async handleConcurrentConnection(
    ws: ServerWebSocket<unknown>,
    clientData: ClientData,
    existingSession: Session,
    props: Record<string, any> | undefined
  ): Promise<boolean> {
    const policy = this.sessionManager?.getConcurrentPolicy() ?? "kick-old";

    switch (policy) {
      case "kick-old": {
        // Kick existing connections
        const existingConnections = this.sessionManager?.getConnections(
          existingSession.id
        );
        if (existingConnections) {
          for (const conn of existingConnections) {
            const oldWs = conn as ServerWebSocket<unknown>;
            const expiredMsg: SessionExpiredMessage = {
              type: "sessionExpired",
              sessionId: existingSession.id,
              reason: "kicked",
            };
            oldWs.send(JSON.stringify(expiredMsg));
            oldWs.close(1000, "Session taken by new connection");
          }
        }
        return true;
      }

      case "reject-new": {
        // Reject this connection
        const expiredMsg: SessionExpiredMessage = {
          type: "sessionExpired",
          sessionId: existingSession.id,
          reason: "kicked",
        };
        ws.send(JSON.stringify(expiredMsg));
        ws.close(1000, "Session already active");
        return false;
      }

      case "allow-multiple": {
        // Allow both connections
        return true;
      }

      default:
        return true;
    }
  }

  /**
   * Trigger onReconnect hook
   */
  private async triggerReconnect(
    clientData: ClientData,
    session: Session,
    savedState: unknown
  ): Promise<void> {
    const handler = this._module?.handlers.onReconnect;
    if (!handler) return;

    let restored = false;
    const restore = (state: unknown) => {
      restored = true;
      clientData.moduleInstance.updateState(state as any);
    };

    await handler({ session, restore });
  }

  /**
   * Handle incoming WebSocket message
   */
  private handleMessage(
    ws: ServerWebSocket<unknown>,
    message: string | Buffer
  ) {
    try {
      const clientData = this.clients.get(ws);
      if (!clientData) return;

      const msg = JSON.parse(message.toString()) as RemoteMessage;

      switch (msg.type) {
        case "hello": {
          const helloMsg = msg as HelloMessage;
          this.initializeSession(
            ws,
            clientData,
            helloMsg.sessionId,
            helloMsg.props
          );
          break;
        }

        case "dispatchAction": {
          const actionMsg = msg as DispatchActionMessage;
          clientData.engine.dispatchAction(actionMsg.action, actionMsg.payload);
          break;
        }

        default:
          // Unknown message type
          break;
      }
    } catch (error) {
      log.error("Error handling WebSocket message:", error);
    }
  }

  /**
   * Handle WebSocket close - suspend session instead of destroying
   */
  private async handleClose(ws: ServerWebSocket<unknown>) {
    const clientData = this.clients.get(ws);
    if (!clientData) return;

    // Clear hello timeout if still pending
    if (clientData.helloTimeout) {
      clearTimeout(clientData.helloTimeout);
    }

    // Get current state for session save
    const currentState = clientData.moduleInstance.getState();

    // Trigger onDisconnect hook
    if (clientData.sessionId && this._module?.handlers.onDisconnect) {
      const session = this.sessionManager?.getActiveSession(clientData.sessionId);
      if (session) {
        await this._module.handlers.onDisconnect({
          state: currentState,
          session,
        });
      }
    }

    // Suspend session (don't destroy immediately)
    if (clientData.sessionId && this.sessionManager) {
      this.sessionManager.untrackConnection(clientData.sessionId, ws);

      // Only suspend if no other connections for this session
      if (this.sessionManager.getConnectionCount(clientData.sessionId) === 0) {
        const session = this.sessionManager.getActiveSession(clientData.sessionId);
        if (session) {
          this.sessionManager.suspendSession(
            clientData.sessionId,
            currentState,
            async (expiredSession) => {
              // Trigger onExpire hook
              if (this._module?.handlers.onExpire) {
                await this._module.handlers.onExpire({ session: expiredSession });
              }
            }
          );
        }
      }
    }

    // Cleanup module instance
    await clientData.moduleInstance.destroy();
    this.clients.delete(ws);

    // Notify disconnection callbacks
    const client: RemoteClient = {
      id: clientData.id,
      socket: ws,
      connectedAt: clientData.connectedAt,
    };
    this._onDisconnectionCallbacks.forEach((cb) => cb(client));
  }

  /**
   * Get current client count
   */
  getClientCount(): number {
    return this.clients.size;
  }

  /**
   * Get session stats
   */
  getSessionStats(): {
    activeSessions: number;
    pendingSessions: number;
    totalConnections: number;
  } {
    return this.sessionManager?.getStats() ?? {
      activeSessions: 0,
      pendingSessions: 0,
      totalConnections: 0,
    };
  }

  /**
   * Broadcast a message to all connected clients
   */
  broadcast(message: RemoteMessage): void {
    const json = JSON.stringify(message);
    this.clients.forEach((_, ws) => {
      ws.send(json);
    });
  }
}

/**
 * Convenience function to create and start a RemoteServer
 */
export function serve(options: {
  module: HypenModule<any>;
  moduleName?: string;
  ui: string;
  port?: number;
  hostname?: string;
  session?: SessionConfig;
  onConnection?: (client: RemoteClient) => void;
  onDisconnection?: (client: RemoteClient) => void;
}): RemoteServer {
  const server = new RemoteServer()
    .module(options.moduleName ?? "App", options.module)
    .ui(options.ui);

  if (options.port || options.hostname) {
    server.config({
      port: options.port,
      hostname: options.hostname,
    });
  }

  if (options.session) {
    server.session(options.session);
  }

  if (options.onConnection) {
    server.onConnection(options.onConnection);
  }

  if (options.onDisconnection) {
    server.onDisconnection(options.onDisconnection);
  }

  return server.listen(options.port);
}
