forge-codegen 0.0.1-alpha

use std::path::PathBuf;

use super::Error;

/// Generates the TypeScript client.
pub struct ClientGenerator {
    #[allow(dead_code)]
    output_dir: PathBuf,
}

impl ClientGenerator {
    /// Create a new client generator.
    pub fn new(output_dir: impl Into<PathBuf>) -> Self {
        Self {
            output_dir: output_dir.into(),
        }
    }

    /// Generate the client code.
    pub fn generate(&self) -> Result<String, Error> {
        Ok(self.generate_client_content())
    }

    /// Generate the client content.
    fn generate_client_content(&self) -> String {
        r#"// Auto-generated by FORGE - DO NOT EDIT
import type { ForgeError } from './types';

// Client configuration
export interface ForgeClientConfig {
  url: string;
  getToken?: () => string | null | Promise<string | null>;
  onAuthError?: (error: ForgeError) => void;
  timeout?: number;
  retries?: number;
}

// WebSocket connection state
export type ConnectionState = 'connecting' | 'connected' | 'reconnecting' | 'disconnected';

// RPC request
interface RpcRequest {
  function: string;
  args: unknown;
  requestId?: string;
}

// RPC response
interface RpcResponse<T = unknown> {
  success: boolean;
  data?: T;
  error?: ForgeError;
  requestId?: string;
}

// The main FORGE client
export class ForgeClient {
  private config: ForgeClientConfig;
  private ws: WebSocket | null = null;
  private connectionState: ConnectionState = 'disconnected';
  private reconnectAttempts = 0;
  private maxReconnectAttempts = 10;
  private reconnectDelay = 1000;
  private subscriptions = new Map<string, (data: unknown) => void>();
  private pendingRequests = new Map<string, {
    resolve: (value: unknown) => void;
    reject: (error: Error) => void;
  }>();

  constructor(config: ForgeClientConfig) {
    this.config = {
      timeout: 30000,
      retries: 3,
      ...config,
    };
  }

  // Get the current connection state
  getConnectionState(): ConnectionState {
    return this.connectionState;
  }

  // Connect to the WebSocket server
  async connect(): Promise<void> {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      return;
    }

    return new Promise((resolve, reject) => {
      const wsUrl = this.config.url.replace(/^http/, 'ws') + '/ws';
      this.ws = new WebSocket(wsUrl);
      this.connectionState = 'connecting';

      this.ws.onopen = () => {
        this.connectionState = 'connected';
        this.reconnectAttempts = 0;
        resolve();
      };

      this.ws.onerror = (event) => {
        reject(new Error('WebSocket connection failed'));
      };

      this.ws.onclose = () => {
        this.connectionState = 'disconnected';
        this.handleDisconnect();
      };

      this.ws.onmessage = (event) => {
        this.handleMessage(event.data);
      };
    });
  }

  // Disconnect from the server
  disconnect(): void {
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
    this.connectionState = 'disconnected';
    this.subscriptions.clear();
    this.pendingRequests.clear();
  }

  // Call a remote function
  async call<T>(functionName: string, args: unknown): Promise<T> {
    const token = await this.getToken();

    const response = await fetch(`${this.config.url}/rpc/${functionName}`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        ...(token ? { 'Authorization': `Bearer ${token}` } : {}),
      },
      body: JSON.stringify({ args }),
    });

    const result: RpcResponse<T> = await response.json();

    if (!result.success || result.error) {
      const error = result.error || { code: 'UNKNOWN', message: 'Unknown error' };
      if (error.code === 'UNAUTHORIZED' && this.config.onAuthError) {
        this.config.onAuthError(error);
      }
      throw new ForgeClientError(error.code, error.message, error.details);
    }

    return result.data as T;
  }

  // Subscribe to a query
  subscribe<T>(
    functionName: string,
    args: unknown,
    callback: (data: T) => void
  ): () => void {
    const subscriptionId = this.generateId();

    // Store the callback
    this.subscriptions.set(subscriptionId, callback as (data: unknown) => void);

    // Send subscription request
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify({
        type: 'subscribe',
        id: subscriptionId,
        function: functionName,
        args,
      }));
    }

    // Return unsubscribe function
    return () => {
      this.subscriptions.delete(subscriptionId);
      if (this.ws && this.ws.readyState === WebSocket.OPEN) {
        this.ws.send(JSON.stringify({
          type: 'unsubscribe',
          id: subscriptionId,
        }));
      }
    };
  }

  // Get the auth token
  private async getToken(): Promise<string | null> {
    if (!this.config.getToken) {
      return null;
    }
    return this.config.getToken();
  }

  // Handle WebSocket messages
  private handleMessage(data: string): void {
    try {
      const message = JSON.parse(data);

      switch (message.type) {
        case 'data':
        case 'delta': {
          const callback = this.subscriptions.get(message.subscriptionId);
          if (callback) {
            callback(message.data);
          }
          break;
        }
        case 'response': {
          const pending = this.pendingRequests.get(message.requestId);
          if (pending) {
            if (message.success) {
              pending.resolve(message.data);
            } else {
              pending.reject(new ForgeClientError(
                message.error?.code || 'UNKNOWN',
                message.error?.message || 'Unknown error',
                message.error?.details
              ));
            }
            this.pendingRequests.delete(message.requestId);
          }
          break;
        }
        case 'error': {
          console.error('FORGE error:', message.error);
          break;
        }
      }
    } catch (e) {
      console.error('Failed to parse WebSocket message:', e);
    }
  }

  // Handle disconnection
  private handleDisconnect(): void {
    if (this.reconnectAttempts >= this.maxReconnectAttempts) {
      return;
    }

    this.connectionState = 'reconnecting';
    this.reconnectAttempts++;

    const delay = Math.min(
      this.reconnectDelay * Math.pow(2, this.reconnectAttempts - 1),
      30000
    );

    setTimeout(() => {
      this.connect().catch(() => {
        // Will retry on next disconnect
      });
    }, delay);
  }

  // Generate a unique ID
  private generateId(): string {
    return Math.random().toString(36).substring(2, 15);
  }
}

// Custom error class
export class ForgeClientError extends Error {
  code: string;
  details?: Record<string, unknown>;

  constructor(code: string, message: string, details?: Record<string, unknown>) {
    super(message);
    this.name = 'ForgeClientError';
    this.code = code;
    this.details = details;
  }
}

// Create a new client instance
export function createForgeClient(config: ForgeClientConfig): ForgeClient {
  return new ForgeClient(config);
}
"#
        .to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_client_generator_creation() {
        let gen = ClientGenerator::new("/tmp/forge");
        assert_eq!(gen.output_dir, PathBuf::from("/tmp/forge"));
    }

    #[test]
    fn test_generate_client_content() {
        let gen = ClientGenerator::new("/tmp/forge");
        let content = gen.generate_client_content();
        assert!(content.contains("ForgeClient"));
        assert!(content.contains("ForgeClientConfig"));
        assert!(content.contains("ConnectionState"));
        assert!(content.contains("connect()"));
        assert!(content.contains("disconnect()"));
        assert!(content.contains("call<T>"));
        assert!(content.contains("subscribe<T>"));
    }
}