src/lib/elements/compute/RemoteComputeManager.svelte.ts

import { RemoteCompute } from './RemoteCompute.svelte';
import type { Position3D } from '$lib/types/positionable.js';
import { generateName } from '$lib/utils/generateName.js';
import { positionManager } from '$lib/utils/positionManager.js';
import { 
  rootGet,
  healthCheckHealthGet,
  listSessionsSessionsGet,
  createSessionSessionsPost,
  startInferenceSessionsSessionIdStartPost,
  stopInferenceSessionsSessionIdStopPost,
  deleteSessionSessionsSessionIdDelete
} from '@robothub/inference-server-client';
import { settings } from '$lib/runes/settings.svelte';
import type { 
  CreateSessionRequest, 
  CreateSessionResponse
} from '@robothub/inference-server-client';

export type ModelType = 'act' | 'diffusion' | 'smolvla' | 'pi0' | 'groot' | 'custom';

export interface ModelTypeConfig {
  id: ModelType;
  label: string;
  icon: string;
  description: string;
  defaultPolicyPath: string;
  defaultCameraNames: string[];
  requiresLanguageInstruction?: boolean;
  enabled: boolean;
}

export const MODEL_TYPES: Record<ModelType, ModelTypeConfig> = {
  act: {
    id: 'act',
    label: 'ACT Model',
    icon: 'icon-[mdi--brain]',
    description: 'Action Chunking with Transformers',
    defaultPolicyPath: 'LaetusH/act_so101_beyond',
    defaultCameraNames: ['front'],
    enabled: true
  },
  diffusion: {
    id: 'diffusion',
    label: 'Diffusion Policy',
    icon: 'icon-[mdi--creation]',
    description: 'Diffusion-based robot control',
    defaultPolicyPath: 'diffusion_policy/default',
    defaultCameraNames: ['front', 'wrist'],
    enabled: true
  },
  smolvla: {
    id: 'smolvla',
    label: 'SmolVLA',
    icon: 'icon-[mdi--eye-outline]',
    description: 'Small Vision-Language-Action model',
    defaultPolicyPath: 'smolvla/latest',
    defaultCameraNames: ['front'],
    requiresLanguageInstruction: true,
    enabled: true
  },
  pi0: {
    id: 'pi0',
    label: 'Pi0',
    icon: 'icon-[mdi--pi]',
    description: 'Lightweight robotics model',
    defaultPolicyPath: 'pi0/base',
    defaultCameraNames: ['front'],
    enabled: true
  },
  groot: {
    id: 'groot',
    label: 'NVIDIA Groot',
    icon: 'icon-[mdi--robot-outline]',
    description: 'Humanoid robotics foundation model',
    defaultPolicyPath: 'nvidia/groot',
    defaultCameraNames: ['front', 'left', 'right'],
    requiresLanguageInstruction: true,
    enabled: false // Not yet implemented
  },
  custom: {
    id: 'custom',
    label: 'Custom Model',
    icon: 'icon-[mdi--cog]',
    description: 'Custom model configuration',
    defaultPolicyPath: '',
    defaultCameraNames: ['front'],
    enabled: true
  }
};

export interface AISessionConfig {
  sessionId: string;
  modelType: ModelType;
  policyPath: string;
  cameraNames: string[];
  transportServerUrl: string;
  workspaceId?: string;
  languageInstruction?: string;
}

export interface AISessionResponse {
  workspace_id: string;
  camera_room_ids: Record<string, string>;
  joint_input_room_id: string;
  joint_output_room_id: string;
}

export interface AISessionStatus {
  session_id: string;
  status: 'initializing' | 'ready' | 'running' | 'stopped';
  policy_path: string;
  camera_names: string[];
  workspace_id: string;
  rooms: {
    workspace_id: string;
    camera_room_ids: Record<string, string>;
    joint_input_room_id: string;
    joint_output_room_id: string;
  };
  stats: {
    inference_count: number;
    commands_sent: number;
    joints_received: number;
    images_received: Record<string, number>;
    errors: number;
    actions_in_queue: number;
  };
  inference_stats?: {
    inference_count: number;
    total_inference_time: number;
    average_inference_time: number;
    average_fps: number;
    is_loaded: boolean;
    device: string;
  };
  error_message?: string;
}

export class RemoteComputeManager {
  private _computes = $state<RemoteCompute[]>([]);

  constructor() {
    // No client initialization needed anymore
  }

  // Reactive getters
  get computes(): RemoteCompute[] {
    return this._computes;
  }

  get computeCount(): number {
    return this._computes.length;
  }

  get runningComputes(): RemoteCompute[] {
    return this._computes.filter(compute => compute.status === 'running');
  }

  /**
   * Get available model types
   */
  get availableModelTypes(): ModelTypeConfig[] {
    return Object.values(MODEL_TYPES).filter(model => model.enabled);
  }

  /**
   * Get model type configuration
   */
  getModelTypeConfig(modelType: ModelType): ModelTypeConfig | undefined {
    return MODEL_TYPES[modelType];
  }

  /**
   * Create a new AI compute instance with full configuration
   */
  async createComputeWithSession(
    config: AISessionConfig,
    computeId?: string,
    computeName?: string,
    position?: Position3D
  ): Promise<{ success: boolean; error?: string; compute?: RemoteCompute }> {
    const finalComputeId = computeId || generateName();
    
    // Check if compute already exists
    if (this._computes.find(c => c.id === finalComputeId)) {
      return { success: false, error: `Compute with ID ${finalComputeId} already exists` };
    }

    try {
      // Create compute instance
      const compute = new RemoteCompute(finalComputeId, computeName);
      compute.modelType = config.modelType;

      // Set position (from position manager if not provided)
      compute.position = position || positionManager.getNextPosition();

      // Add to reactive array
      this._computes.push(compute);

      // Create the session immediately
      const sessionResult = await this.createSession(compute.id, config);
      if (!sessionResult.success) {
        // Remove compute if session creation failed
        await this.removeCompute(compute.id);
        return { success: false, error: sessionResult.error };
      }

      console.log(`Created compute ${finalComputeId} with ${config.modelType} model at position (${compute.position.x.toFixed(1)}, ${compute.position.y.toFixed(1)}, ${compute.position.z.toFixed(1)}). Total computes: ${this._computes.length}`);

      return { success: true, compute };
    } catch (error) {
      console.error('Failed to create compute with session:', error);
      return { 
        success: false, 
        error: error instanceof Error ? error.message : String(error) 
      };
    }
  }

  /**
   * Create a new AI compute instance (legacy method)
   */
  createCompute(id?: string, name?: string, position?: Position3D): RemoteCompute {
    const computeId = id || generateName();
    
    // Check if compute already exists
    if (this._computes.find(c => c.id === computeId)) {
      throw new Error(`Compute with ID ${computeId} already exists`);
    }

    // Create compute instance
    const compute = new RemoteCompute(computeId, name);

    // Set position (from position manager if not provided)
    compute.position = position || positionManager.getNextPosition();

    // Add to reactive array
    this._computes.push(compute);

    console.log(`Created compute ${computeId} at position (${compute.position.x.toFixed(1)}, ${compute.position.y.toFixed(1)}, ${compute.position.z.toFixed(1)}). Total computes: ${this._computes.length}`);

    return compute;
  }

  /**
   * Get compute by ID
   */
  getCompute(id: string): RemoteCompute | undefined {
    return this._computes.find(c => c.id === id);
  }

  /**
   * Remove a compute instance
   */
  async removeCompute(id: string): Promise<void> {
    const computeIndex = this._computes.findIndex(c => c.id === id);
    if (computeIndex === -1) return;

    const compute = this._computes[computeIndex];

    // Clean up compute resources
    await this.stopSession(id);
    await this.deleteSession(id);

    // Remove from reactive array
    this._computes.splice(computeIndex, 1);

    console.log(`Removed compute ${id}. Remaining computes: ${this._computes.length}`);
  }

  /**
   * Create an Inference Session
   */
  async createSession(computeId: string, config: AISessionConfig): Promise<{ success: boolean; error?: string; data?: AISessionResponse }> {
    const compute = this.getCompute(computeId);
    if (!compute) {
      return { success: false, error: `Compute ${computeId} not found` };
    }

    try {
      const request: CreateSessionRequest = {
        session_id: config.sessionId,
        policy_path: config.policyPath,
        camera_names: config.cameraNames,
        transport_server_url: config.transportServerUrl,
        workspace_id: config.workspaceId || undefined,
        policy_type: config.modelType, // Use model type as policy type
        language_instruction: config.languageInstruction || undefined,
      };

      const response = await createSessionSessionsPost({
        body: request,
        baseUrl: settings.inferenceServerUrl
      });

      if (!response.data) {
        throw new Error('Failed to create session - no data returned');
      }

      const data: CreateSessionResponse = response.data;
      
      // Update compute with session info
      compute.sessionId = config.sessionId;
      compute.status = 'ready';
      compute.sessionConfig = config;
      compute.sessionData = {
        workspace_id: data.workspace_id,
        camera_room_ids: data.camera_room_ids,
        joint_input_room_id: data.joint_input_room_id,
        joint_output_room_id: data.joint_output_room_id
      };

      return { success: true, data: compute.sessionData };
    } catch (error) {
      console.error(`Failed to create session for compute ${computeId}:`, error);
      return { 
        success: false, 
        error: error instanceof Error ? error.message : String(error) 
      };
    }
  }

  /**
   * Start inference for a session
   */
  async startSession(computeId: string): Promise<{ success: boolean; error?: string }> {
    const compute = this.getCompute(computeId);
    if (!compute || !compute.sessionId) {
      return { success: false, error: 'No session to start' };
    }

    try {
      await startInferenceSessionsSessionIdStartPost({
        path: { session_id: compute.sessionId },
        baseUrl: settings.inferenceServerUrl
      });
      compute.status = 'running';
      return { success: true };
    } catch (error) {
      console.error(`Failed to start session for compute ${computeId}:`, error);
      return { 
        success: false, 
        error: error instanceof Error ? error.message : String(error) 
      };
    }
  }

  /**
   * Stop inference for a session
   */
  async stopSession(computeId: string): Promise<{ success: boolean; error?: string }> {
    const compute = this.getCompute(computeId);
    if (!compute || !compute.sessionId) {
      return { success: false, error: 'No session to stop' };
    }

    try {
      await stopInferenceSessionsSessionIdStopPost({
        path: { session_id: compute.sessionId },
        baseUrl: settings.inferenceServerUrl
      });
      compute.status = 'stopped';
      return { success: true };
    } catch (error) {
      console.error(`Failed to stop session for compute ${computeId}:`, error);
      return { 
        success: false, 
        error: error instanceof Error ? error.message : String(error) 
      };
    }
  }

  /**
   * Delete a session
   */
  async deleteSession(computeId: string): Promise<{ success: boolean; error?: string }> {
    const compute = this.getCompute(computeId);
    if (!compute || !compute.sessionId) {
      return { success: true }; // Already deleted
    }

    try {
      await deleteSessionSessionsSessionIdDelete({
        path: { session_id: compute.sessionId },
        baseUrl: settings.inferenceServerUrl
      });

      // Reset compute session info
      compute.sessionId = null;
      compute.status = 'disconnected';
      compute.sessionConfig = null;
      compute.sessionData = null;

      return { success: true };
    } catch (error) {
      console.error(`Failed to delete session for compute ${computeId}:`, error);
      return { 
        success: false, 
        error: error instanceof Error ? error.message : String(error) 
      };
    }
  }

  /**
   * Get session status
   */
  async getSessionStatus(computeId: string): Promise<{ success: boolean; data?: AISessionStatus; error?: string }> {
    const compute = this.getCompute(computeId);
    if (!compute || !compute.sessionId) {
      return { success: false, error: 'No session found' };
    }

    try {
      // Get all sessions and find the one we want
      const response = await listSessionsSessionsGet({
        baseUrl: settings.inferenceServerUrl
      });

      if (!response.data) {
        throw new Error('Failed to get sessions list');
      }

      const session = response.data.find(s => s.session_id === compute.sessionId);
      if (!session) {
        throw new Error(`Session ${compute.sessionId} not found`);
      }
      
      // Update compute status
      compute.status = session.status as 'initializing' | 'ready' | 'running' | 'stopped';
      
      // Convert to AISessionStatus format
      const sessionStatus: AISessionStatus = {
        session_id: session.session_id,
        status: session.status as 'initializing' | 'ready' | 'running' | 'stopped',
        policy_path: session.policy_path,
        camera_names: session.camera_names,
        workspace_id: session.workspace_id,
        rooms: session.rooms as any,
        stats: session.stats as any,
        inference_stats: session.inference_stats as any,
        error_message: session.error_message || undefined
      };
      
      return { success: true, data: sessionStatus };
    } catch (error) {
      console.error(`Failed to get session status for compute ${computeId}:`, error);
      return { 
        success: false, 
        error: error instanceof Error ? error.message : String(error) 
      };
    }
  }

  /**
   * Check AI server health
   */
  async checkServerHealth(): Promise<{ success: boolean; data?: any; error?: string }> {
    try {
      const healthResponse = await rootGet({
        baseUrl: settings.inferenceServerUrl
      });
      
      if (!healthResponse.data) {
        return { success: false, error: 'Server is not healthy' };
      }
      
      // Get detailed health info
      const detailedHealthResponse = await healthCheckHealthGet({
        baseUrl: settings.inferenceServerUrl
      });
      
      return { success: true, data: detailedHealthResponse.data };
    } catch (error) {
      console.error('Failed to check AI server health:', error);
      return { 
        success: false, 
        error: error instanceof Error ? error.message : String(error) 
      };
    }
  }

  /**
   * Clean up all computes
   */
  async destroy(): Promise<void> {
    const cleanupPromises = this._computes.map(async (compute) => {
      await this.stopSession(compute.id);
      await this.deleteSession(compute.id);
    });
    await Promise.allSettled(cleanupPromises);
    this._computes.length = 0;
  }
}

// Global compute manager instance
export const remoteComputeManager = new RemoteComputeManager();