frontend/src/lib/interactiveClient.ts

/** Typed fetch wrapper for the FastAPI `/interactive/*` router.
 *
 * Base URL precedence:
 *   1. explicit `options.baseUrl` from the caller
 *   2. `VITE_PHYSIX_API_URL` baked in at build time. An EXPLICIT empty
 *      string means "use same-origin / relative URLs" — that's how the
 *      HF Space build is configured (the SPA and the FastAPI server
 *      live behind the same host, so `/interactive/*` resolves to the
 *      Space itself).
 *   3. fall back to `http://localhost:8000` for `pnpm dev` against a
 *      separately-running uvicorn.
 */

import type { PhysiXAction, PhysiXObservation, TrajectorySample } from "@/types/physix";

const DEFAULT_BASE_URL = "http://localhost:8000";

function resolveBaseUrl(override?: string): string {
  if (override !== undefined) return override;
  const env = (import.meta as ImportMeta & { env?: Record<string, string> }).env;
  // `??` (not `||`/truthy) so an explicit empty string survives — the
  // production build sets VITE_PHYSIX_API_URL="" to mean same-origin.
  // If the variable wasn't defined at all (dev), `env?.X` is undefined
  // and we fall through to the localhost default.
  const fromEnv = env?.VITE_PHYSIX_API_URL;
  return fromEnv ?? DEFAULT_BASE_URL;
}

export interface SystemDescriptor {
  system_id: string;
  state_variables: string[];
}

export interface InteractiveStartResponse {
  session_id: string;
  observation: PhysiXObservation;
  system: SystemDescriptor;
  max_turns: number;
}

export interface LlmStepRequest {
  /** OpenAI-compatible /v1 base URL. Required by the server. */
  base_url: string;
  /** Provider-native model id (HF repo id, Ollama tag, OpenAI name, …). */
  model: string;
  /** Bearer token. Server falls back to env vars when omitted. */
  api_key?: string | undefined;
  temperature?: number | undefined;
  max_tokens?: number | undefined;
}

export interface LlmStepResponse {
  observation: PhysiXObservation;
  predicted_trajectory: TrajectorySample[];
  action: PhysiXAction;
  raw_completion: string;
  latency_s: number;
  model: string;
}

export interface SessionSummary {
  session_id: string;
  system_id: string;
  turn: number;
  max_turns: number;
  converged: boolean;
  done: boolean;
}

export interface LlmModelInfo {
  name: string;
  size_bytes: number | null;
  parameter_size: string | null;
  family: string | null;
}

export interface LlmModelsResponse {
  models: LlmModelInfo[];
  error: string | null;
}

/**
 * OpenEnv stateless contract. The official OpenEnv server constructs a
 * fresh env per request, so bare `/reset` and `/step` are stateless —
 * they're useful for inspection and `metadata`/`schema` introspection,
 * but you cannot run a multi-turn episode through them. For stateful
 * step-by-step interaction we use the `/interactive/*` session router
 * instead (see `directStep`, `getSummary`, `endSession`).
 */
export interface OpenEnvResetRequest {
  seed?: number | null;
  episode_id?: string | null;
  /** PhysiX-specific extension forwarded via `additionalProperties: true`. */
  system_id?: string | null;
}

export interface OpenEnvResetResponse {
  observation: PhysiXObservation;
  reward: number | null;
  done: boolean;
}

export interface OpenEnvMetadata {
  name: string;
  description: string;
  readme_content?: string | null;
  version?: string | null;
  author?: string | null;
  documentation_url?: string | null;
}

export interface OpenEnvSchemaResponse {
  action: Record<string, unknown>;
  observation: Record<string, unknown>;
  state: Record<string, unknown>;
}

/**
 * Stateful per-session direct step. Same shape as `LlmStepResponse`
 * minus the `raw_completion` / `model` / `latency_s` fields, which
 * don't apply when the user supplies the action directly.
 */
export interface DirectStepResponse {
  observation: PhysiXObservation;
  predicted_trajectory: TrajectorySample[];
}

export class InteractiveApiError extends Error {
  public readonly status: number;
  public readonly detail: string;

  public constructor(status: number, detail: string) {
    super(`[${status}] ${detail}`);
    this.name = "InteractiveApiError";
    this.status = status;
    this.detail = detail;
  }
}

export interface InteractiveClientOptions {
  baseUrl?: string;
  fetchImpl?: typeof fetch;
}

export class InteractiveClient {
  private readonly baseUrl: string;
  private readonly fetchImpl: typeof fetch;

  public constructor(options: InteractiveClientOptions = {}) {
    this.baseUrl = resolveBaseUrl(options.baseUrl).replace(/\/+$/, "");
    this.fetchImpl = options.fetchImpl ?? fetch.bind(globalThis);
  }

  public async listSystems(): Promise<SystemDescriptor[]> {
    return this.request<SystemDescriptor[]>("/interactive/systems", { method: "GET" });
  }

  public async listModels(): Promise<LlmModelsResponse> {
    return this.request<LlmModelsResponse>("/interactive/models", { method: "GET" });
  }

  public async startSession(payload: {
    system_id?: string | undefined;
    seed?: number | undefined;
    max_turns?: number | undefined;
  }): Promise<InteractiveStartResponse> {
    const compact = stripUndefined(payload);
    return this.request<InteractiveStartResponse>("/interactive/sessions", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(compact),
    });
  }

  public async llmStep(
    sessionId: string,
    request: LlmStepRequest,
  ): Promise<LlmStepResponse> {
    const compact = stripUndefined(request as unknown as Record<string, unknown>);
    return this.request<LlmStepResponse>(
      `/interactive/sessions/${encodeURIComponent(sessionId)}/llm-step`,
      {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(compact),
      },
    );
  }

  public async endSession(sessionId: string): Promise<void> {
    const response = await this.fetchImpl(
      `${this.baseUrl}/interactive/sessions/${encodeURIComponent(sessionId)}`,
      { method: "DELETE" },
    );
    if (!response.ok && response.status !== 404) {
      const detail = await safeReadDetail(response);
      throw new InteractiveApiError(response.status, detail);
    }
  }

  public async getSummary(sessionId: string): Promise<SessionSummary> {
    return this.request<SessionSummary>(
      `/interactive/sessions/${encodeURIComponent(sessionId)}`,
      { method: "GET" },
    );
  }

  // -------- OpenEnv read-only endpoints --------

  public async openEnvMetadata(): Promise<OpenEnvMetadata> {
    return this.request<OpenEnvMetadata>("/metadata", { method: "GET" });
  }

  public async openEnvSchema(): Promise<OpenEnvSchemaResponse> {
    return this.request<OpenEnvSchemaResponse>("/schema", { method: "GET" });
  }

  /**
   * Stateless `POST /reset` — useful for inspecting what the env
   * returns at episode start, but the OpenEnv framework constructs a
   * new env per request so a follow-up `POST /step` cannot succeed.
   * For multi-step episodes use `startSession` + `directStep`.
   */
  public async openEnvReset(
    payload: OpenEnvResetRequest,
  ): Promise<OpenEnvResetResponse> {
    const compact = stripUndefined(payload as Record<string, unknown>);
    return this.request<OpenEnvResetResponse>("/reset", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(compact),
    });
  }

  // -------- Stateful per-session direct step --------

  /** Apply a user-supplied action to a session created via `startSession`. */
  public async directStep(
    sessionId: string,
    action: PhysiXAction,
  ): Promise<DirectStepResponse> {
    return this.request<DirectStepResponse>(
      `/interactive/sessions/${encodeURIComponent(sessionId)}/step`,
      {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ action }),
      },
    );
  }

  private async request<T>(path: string, init: RequestInit): Promise<T> {
    const response = await this.fetchImpl(`${this.baseUrl}${path}`, init);
    if (!response.ok) {
      const detail = await safeReadDetail(response);
      throw new InteractiveApiError(response.status, detail);
    }
    return (await response.json()) as T;
  }
}

function stripUndefined<T extends Record<string, unknown>>(input: T): Partial<T> {
  const out: Partial<T> = {};
  for (const [key, value] of Object.entries(input) as [keyof T, T[keyof T]][]) {
    if (value !== undefined) {
      out[key] = value;
    }
  }
  return out;
}

async function safeReadDetail(response: Response): Promise<string> {
  try {
    const body = (await response.json()) as { detail?: unknown };
    if (typeof body.detail === "string") return body.detail;
    return JSON.stringify(body);
  } catch {
    return response.statusText || "Request failed";
  }
}