import { Type } from "@sinclair/typebox";
import { definePluginEntry } from "openclaw/plugin-sdk/plugin-entry";
import { homedir } from "node:os";
import { createHash, randomUUID, randomBytes } from "node:crypto";
import fs from "node:fs/promises";
import path from "node:path";

type JsonRpcSuccess<T> = { jsonrpc: "2.0"; id?: string | number | null; result: T };
type JsonRpcError = { jsonrpc: "2.0"; id?: string | number | null; error: { code: number; message: string; data?: unknown } };
type JsonRpcResponse<T> = JsonRpcSuccess<T> | JsonRpcError;

type ToolCallResult = {
  content?: Array<{ type: string; text?: string; [k: string]: unknown }>;
  structuredContent?: unknown;
  isError?: boolean;
  [k: string]: unknown;
};

type PluginConfig = {
  endpoint?: string;
  callbackPort?: number;
  requestTimeoutMs?: number;
  transport?: "streamable-http";
};

type RuntimeApi = {
  registerTool: (...args: any[]) => void;
  getConfig?: <T = unknown>() => T;
  logger?: {
    info?: (...args: unknown[]) => void;
    warn?: (...args: unknown[]) => void;
    error?: (...args: unknown[]) => void;
  };
};

type OAuthCredential = {
  type: "oauth";
  provider: string;
  access: string;
  refresh?: string;
  expires?: number;
};

type RequiredPluginConfig = Required<Omit<PluginConfig, never>>;

type ProtectedResourceMetadata = {
  resource?: string;
  authorization_servers?: string[];
  bearer_methods_supported?: string[];
  scopes_supported?: string[];
  [k: string]: unknown;
};

type AuthorizationServerMetadata = {
  issuer?: string;
  authorization_endpoint?: string;
  token_endpoint?: string;
  registration_endpoint?: string;
  scopes_supported?: string[];
  code_challenge_methods_supported?: string[];
  token_endpoint_auth_methods_supported?: string[];
  [k: string]: unknown;
};

type ThunderCredential = OAuthCredential & { _thunderMeta?: { authorizationServer?: string; tokenEndpoint?: string; clientId?: string } };
type PendingAuth = { state: string; codeVerifier: string; clientId: string; issuer: string; tokenEndpoint: string; callbackUrl: string; createdAt: number };

const DEFAULT_ENDPOINT = "https://api.thundercompute.com:8443/mcp";
const DEFAULT_TIMEOUT_MS = 60_000;
const DEFAULT_CALLBACK_PORT = 1455;
const PROTOCOL_VERSION = "2025-03-26";
const PROVIDER_ID = "thunder-compute";
const APP_NAME = "OpenClaw Thunder Compute MCP";

function base64url(input: Buffer | string): string {
  const buf = Buffer.isBuffer(input) ? input : Buffer.from(input);
  return buf.toString("base64").replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/g, "");
}

function sha256Base64Url(input: string): string {
  return base64url(createHash("sha256").update(input).digest());
}

function getConfig(api: RuntimeApi): RequiredPluginConfig {
  const raw = (api.getConfig?.<PluginConfig>() ?? {}) as PluginConfig;
  return {
    endpoint: raw.endpoint || DEFAULT_ENDPOINT,
    callbackPort: raw.callbackPort || DEFAULT_CALLBACK_PORT,
    requestTimeoutMs: raw.requestTimeoutMs || DEFAULT_TIMEOUT_MS,
    transport: raw.transport || "streamable-http",
  };
}

function getPluginStateDir(): string {
  return path.join(homedir(), ".openclaw", "plugins", PROVIDER_ID);
}
function getCredentialPath(): string {
  return path.join(getPluginStateDir(), "auth.json");
}
function getPendingAuthPath(): string {
  return path.join(getPluginStateDir(), "pending-auth.json");
}

async function readJsonFile(filePath: string): Promise<any | null> {
  try {
    return JSON.parse(await fs.readFile(filePath, "utf8"));
  } catch {
    return null;
  }
}
async function writeJsonFile(filePath: string, value: any): Promise<void> {
  await fs.mkdir(path.dirname(filePath), { recursive: true });
  await fs.writeFile(filePath, JSON.stringify(value, null, 2));
}

async function loadStoredCredential(): Promise<ThunderCredential | null> {
  const cred = await readJsonFile(getCredentialPath());
  if (!cred || cred.type !== 'oauth') return null;
  return cred as ThunderCredential;
}
async function persistCredential(cred: ThunderCredential): Promise<void> {
  await writeJsonFile(getCredentialPath(), cred);
}
async function loadPendingAuth(): Promise<PendingAuth | null> {
  return await readJsonFile(getPendingAuthPath()) as PendingAuth | null;
}
async function persistPendingAuth(p: PendingAuth): Promise<void> {
  await writeJsonFile(getPendingAuthPath(), p);
}
async function clearPendingAuth(): Promise<void> {
  try { await fs.unlink(getPendingAuthPath()); } catch {}
}
async function clearCredential(): Promise<void> {
  try { await fs.unlink(getCredentialPath()); } catch {}
}

async function fetchJson<T>(url: string, init?: RequestInit): Promise<T> {
  const response = await fetch(url, init);
  const text = await response.text();
  if (!response.ok) {
    throw new Error(`HTTP ${response.status} from ${url}: ${text || response.statusText}`);
  }
  return JSON.parse(text) as T;
}

async function discoverProtectedResource(endpoint: string): Promise<ProtectedResourceMetadata> {
  const base = new URL(endpoint);
  const resourceUrl = new URL('/.well-known/oauth-protected-resource', `${base.protocol}//${base.host}`).toString();
  return fetchJson<ProtectedResourceMetadata>(resourceUrl);
}
function appendWellKnown(issuer: string, suffix: string): string {
  const normalized = issuer.replace(/\/+$/, '');
  return `${normalized}/.well-known/${suffix}`;
}
async function discoverAuthorizationServer(resource: ProtectedResourceMetadata): Promise<{ issuer: string; metadata: AuthorizationServerMetadata }> {
  const issuer = resource.authorization_servers?.[0];
  if (!issuer) throw new Error('Thunder Compute protected-resource metadata did not advertise an authorization server.');
  const candidates = [appendWellKnown(issuer, 'oauth-authorization-server'), appendWellKnown(issuer, 'openid-configuration')];
  let lastError: unknown = null;
  for (const url of candidates) {
    try {
      const metadata = await fetchJson<AuthorizationServerMetadata>(url);
      return { issuer, metadata };
    } catch (error) {
      lastError = error;
    }
  }
  throw lastError instanceof Error ? lastError : new Error('Unable to discover Thunder Compute authorization server metadata.');
}
async function maybeRegisterClient(metadata: AuthorizationServerMetadata, callbackUrl: string): Promise<string> {
  const hinted = (metadata as any).client_id || (metadata as any).default_client_id || (metadata as any).public_client_id;
  if (typeof hinted === 'string' && hinted) return hinted;
  if (!metadata.registration_endpoint) {
    throw new Error('Thunder Compute OAuth client_id is required but was not advertised.');
  }
  const response = await fetchJson<any>(metadata.registration_endpoint, {
    method: 'POST',
    headers: { 'content-type': 'application/json' },
    body: JSON.stringify({
      client_name: APP_NAME,
      redirect_uris: [callbackUrl],
      token_endpoint_auth_method: 'none',
      grant_types: ['authorization_code', 'refresh_token'],
      response_types: ['code'],
      application_type: 'native',
    }),
  });
  if (!response.client_id) throw new Error('Dynamic client registration succeeded but no client_id was returned.');
  return response.client_id;
}
async function exchangeCodeForToken(params: { tokenEndpoint: string; clientId: string; code: string; codeVerifier: string; redirectUri: string; }): Promise<any> {
  const body = new URLSearchParams({ grant_type: 'authorization_code', client_id: params.clientId, code: params.code, redirect_uri: params.redirectUri, code_verifier: params.codeVerifier });
  return fetchJson<any>(params.tokenEndpoint, { method: 'POST', headers: { 'content-type': 'application/x-www-form-urlencoded' }, body });
}
async function refreshAccessToken(params: { tokenEndpoint: string; clientId: string; refreshToken: string; }): Promise<any> {
  const body = new URLSearchParams({ grant_type: 'refresh_token', client_id: params.clientId, refresh_token: params.refreshToken });
  return fetchJson<any>(params.tokenEndpoint, { method: 'POST', headers: { 'content-type': 'application/x-www-form-urlencoded' }, body });
}
function tokenResponseToCredential(token: any, clientId: string, authorizationServer: string, tokenEndpoint: string): ThunderCredential {
  const expires = token.expires_in ? Date.now() + Number(token.expires_in) * 1000 : undefined;
  return { type: 'oauth', provider: PROVIDER_ID, access: token.access_token, refresh: token.refresh_token, expires, _thunderMeta: { authorizationServer, tokenEndpoint, clientId } };
}
async function ensureCredential(config: RequiredPluginConfig, logger?: RuntimeApi['logger']): Promise<ThunderCredential> {
  const current = await loadStoredCredential();
  if (!current) throw new Error('Thunder Compute OAuth not configured. Use tc_auth_begin, then tc_auth_complete.');
  const expires = current.expires ?? 0;
  if (!expires || expires > Date.now() + 30_000) return current;
  if (!current.refresh) throw new Error('Thunder Compute access token is expired and no refresh token is available. Re-run tc_auth_begin.');
  const tokenEndpoint = current._thunderMeta?.tokenEndpoint;
  const clientId = current._thunderMeta?.clientId;
  if (!tokenEndpoint || !clientId) throw new Error('Thunder Compute token refresh metadata is missing. Re-run tc_auth_begin.');
  logger?.info?.('Refreshing Thunder Compute OAuth access token.');
  const refreshed = await refreshAccessToken({ tokenEndpoint, clientId, refreshToken: current.refresh });
  const merged = tokenResponseToCredential({ ...refreshed, refresh_token: refreshed.refresh_token || current.refresh }, clientId, current._thunderMeta?.authorizationServer || '', tokenEndpoint);
  await persistCredential(merged);
  return merged;
}

class ThunderComputeMcpClient {
  private sessionId?: string;
  private initialized = false;
  constructor(private readonly config: RequiredPluginConfig, private readonly logger?: RuntimeApi['logger']) {}
  async listTools(): Promise<unknown> { await this.ensureInitialized(); const result = await this.rpcCall<{ tools?: unknown[] }>('tools/list', {}); return result.tools ?? result; }
  async callTool(name: string, args: Record<string, unknown>): Promise<ToolCallResult> { await this.ensureInitialized(); return this.rpcCall<ToolCallResult>('tools/call', { name, arguments: args }); }
  private async ensureInitialized(): Promise<void> {
    if (this.initialized) return;
    await this.rpcCall('initialize', { protocolVersion: PROTOCOL_VERSION, capabilities: {}, clientInfo: { name: '@thunder-compute/openclaw-plugin', version: '0.3.0' } });
    try { await this.rpcNotify('notifications/initialized', {}); } catch (error) { this.logger?.warn?.('Thunder Compute MCP initialized notification failed', error); }
    this.initialized = true;
  }
  private async rpcNotify(method: string, params: Record<string, unknown>): Promise<void> { await this.postJsonRpc({ jsonrpc: '2.0', method, params }); }
  private async rpcCall<T>(method: string, params: Record<string, unknown>): Promise<T> {
    const id = randomUUID();
    const response = (await this.postJsonRpc({ jsonrpc: '2.0', id, method, params })) as JsonRpcResponse<T>;
    if ('error' in response) {
      const details = response.error.data ? ` ${JSON.stringify(response.error.data)}` : '';
      throw new Error(`Thunder Compute MCP ${method} failed (${response.error.code}): ${response.error.message}${details}`);
    }
    return response.result;
  }
  private async postJsonRpc(body: Record<string, unknown>): Promise<unknown> {
    const controller = new AbortController();
    const timeout = setTimeout(() => controller.abort(), this.config.requestTimeoutMs);
    try {
      const credential = await ensureCredential(this.config, this.logger);
      const headers: Record<string, string> = { 'content-type': 'application/json', accept: 'application/json, text/event-stream', authorization: `Bearer ${credential.access}` };
      if (this.sessionId) headers['mcp-session-id'] = this.sessionId;
      const response = await fetch(this.config.endpoint, { method: 'POST', headers, body: JSON.stringify(body), signal: controller.signal });
      const nextSessionId = response.headers.get('mcp-session-id') ?? response.headers.get('x-mcp-session-id');
      if (nextSessionId) this.sessionId = nextSessionId;
      const text = await response.text();
      if (!response.ok) throw new Error(`HTTP ${response.status}: ${text || response.statusText}`);
      const trimmed = text.trim();
      if (!trimmed) return { ok: true };
      if (trimmed.startsWith('data:')) {
        const firstDataLine = trimmed.split(/\r?\n/).find((line) => line.startsWith('data:'))?.slice(5).trim();
        if (!firstDataLine) throw new Error(`Unexpected SSE response from Thunder Compute MCP: ${trimmed}`);
        return JSON.parse(firstDataLine);
      }
      return JSON.parse(trimmed);
    } finally { clearTimeout(timeout); }
  }
}

function createClient(api: RuntimeApi): ThunderComputeMcpClient { return new ThunderComputeMcpClient(getConfig(api), api.logger); }
function formatToolResult(result: ToolCallResult) { if (result?.content?.length) return result; return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }], structuredContent: result }; }
function withAuthHint(error: unknown) {
  const message = error instanceof Error ? error.message : String(error);
  if (message.includes('401') || /unauthorized|oauth/i.test(message)) return `${message}

Hint: run tc_auth_begin, open the returned URL in your browser, then run tc_auth_complete with the full redirected URL.`;
  return message;
}

const InstanceConfig = Type.Object({
  gpuType: Type.String({ description: 'Requested GPU type, e.g. A100, H100, RTX A6000, or a discovered SKU like a100xl_x1_prototyping.' }),
  template: Type.Optional(Type.String({ description: 'Requested template / image name.' })),
  region: Type.Optional(Type.String()),
  vcpus: Type.Optional(Type.Number({ minimum: 1 })),
  memoryGb: Type.Optional(Type.Number({ minimum: 1 })),
  storageGb: Type.Optional(Type.Number({ minimum: 1 })),
  sshKeyIds: Type.Optional(Type.Array(Type.String())),
  metadata: Type.Optional(Type.Record(Type.String(), Type.String())),
  raw: Type.Optional(Type.Any({ description: 'Escape hatch for MCP-server-specific arguments.' })),
});

function normalizeGpuType(input: string | undefined): { gpu_type?: string; num_gpus?: number; mode?: string } {
  if (!input) return {};
  const gpuType = input.trim();
  const sku = gpuType.toLowerCase().match(/^([a-z0-9]+)_x(\d+)_(prototyping|production)$/);
  if (sku) {
    return {
      gpu_type: sku[1],
      num_gpus: Number(sku[2]),
      mode: sku[3],
    };
  }

  const compact = gpuType.toLowerCase().replace(/[\s_-]+/g, '');
  const aliases: Record<string, string> = {
    a100: 'a100xl',
    a100xl: 'a100xl',
    h100: 'h100',
    a6000: 'a6000',
    rtxa6000: 'a6000',
    nvidiartxa6000: 'a6000',
    rtx4090: 'rtx4090',
    nvidiartx4090: 'rtx4090',
  };
  return { gpu_type: aliases[compact] ?? gpuType };
}

function buildCreateInstanceArgs(params: any): Record<string, unknown> {
  if (params.raw) return params.raw;
  const normalized = normalizeGpuType(params.gpuType);
  return {
    gpu_type: normalized.gpu_type,
    template: params.template,
    mode: normalized.mode,
    num_gpus: normalized.num_gpus,
    vcpus: params.vcpus,
    disk_size_gb: params.storageGb,
    ssh_key_name: params.sshKeyName,
    metadata: params.metadata,
  };
}

export default definePluginEntry({
  id: PROVIDER_ID,
  name: 'Thunder Compute',
  description: 'Native OpenClaw agent tools for the Thunder Compute MCP server with OAuth login.',
  register(api: RuntimeApi) {
    const config = getConfig(api);
    const buildExecute = (fn: (client: ThunderComputeMcpClient, params: any) => Promise<ToolCallResult>) => async (_id: string, params: any) => {
      const client = createClient(api);
      try { return formatToolResult(await fn(client, params)); } catch (error) { return { isError: true, content: [{ type: 'text', text: withAuthHint(error) }] }; }
    };

    api.registerTool({
      name: 'tc_auth_status',
      description: 'Show whether Thunder Compute OAuth is configured for this OpenClaw installation.',
      parameters: Type.Object({}),
      async execute() {
        try {
          const credential = await loadStoredCredential();
          if (!credential) return { content: [{ type: 'text', text: 'Thunder Compute OAuth is not configured. Run tc_auth_begin, then tc_auth_complete.' }] };
          return { content: [{ type: 'text', text: `Thunder Compute OAuth is configured. Expires: ${credential.expires ? new Date(credential.expires).toISOString() : 'unknown'}.` }], structuredContent: { expires: credential.expires ?? null } };
        } catch (error) { return { isError: true, content: [{ type: 'text', text: withAuthHint(error) }] }; }
      },
    });

    api.registerTool({
      name: 'tc_auth_begin',
      description: 'Begin Thunder Compute OAuth login. Returns a browser URL to open, then use tc_auth_complete with the redirected URL.',
      parameters: Type.Object({}),
      async execute() {
        try {
          const resource = await discoverProtectedResource(config.endpoint);
          const { issuer, metadata } = await discoverAuthorizationServer(resource);
          if (!metadata.authorization_endpoint || !metadata.token_endpoint) throw new Error('Thunder Compute authorization server metadata is missing authorization_endpoint or token_endpoint.');
          const callbackUrl = `http://127.0.0.1:${config.callbackPort}/auth/callback`;
          const clientId = await maybeRegisterClient(metadata, callbackUrl);
          const state = base64url(randomBytes(24));
          const codeVerifier = base64url(randomBytes(48));
          const codeChallenge = sha256Base64Url(codeVerifier);
          const authUrl = new URL(metadata.authorization_endpoint);
          authUrl.searchParams.set('response_type', 'code');
          authUrl.searchParams.set('client_id', clientId);
          authUrl.searchParams.set('redirect_uri', callbackUrl);
          authUrl.searchParams.set('code_challenge', codeChallenge);
          authUrl.searchParams.set('code_challenge_method', 'S256');
          authUrl.searchParams.set('state', state);
          authUrl.searchParams.set('resource', config.endpoint);
          authUrl.searchParams.set('scope', Array.isArray(resource.scopes_supported) && resource.scopes_supported.length ? resource.scopes_supported.join(' ') : 'openid offline_access');
          await persistPendingAuth({ state, codeVerifier, clientId, issuer, tokenEndpoint: metadata.token_endpoint, callbackUrl, createdAt: Date.now() });
          return { content: [{ type: 'text', text: `Open this URL in your browser to authenticate Thunder Compute:

${authUrl.toString()}

After approval, copy the full redirected URL from the browser address bar and call tc_auth_complete with that URL.` }], structuredContent: { authUrl: authUrl.toString(), callbackUrl } };
        } catch (error) { return { isError: true, content: [{ type: 'text', text: withAuthHint(error) }] }; }
      },
    });

    api.registerTool({
      name: 'tc_auth_complete',
      description: 'Complete Thunder Compute OAuth after tc_auth_begin by pasting the full redirected URL from the browser.',
      parameters: Type.Object({ redirectedUrl: Type.String({ description: 'The full redirected callback URL copied from the browser address bar.' }) }),
      async execute(_id: string, params: { redirectedUrl: string }) {
        try {
          const pending = await loadPendingAuth();
          if (!pending) throw new Error('No pending Thunder Compute auth flow found. Run tc_auth_begin first.');
          const redirected = new URL(params.redirectedUrl);
          const code = redirected.searchParams.get('code') || '';
          const state = redirected.searchParams.get('state') || '';
          if (!code) throw new Error('Redirect URL does not contain an authorization code.');
          if (state !== pending.state) throw new Error('OAuth state mismatch. Re-run tc_auth_begin and try again.');
          const token = await exchangeCodeForToken({ tokenEndpoint: pending.tokenEndpoint, clientId: pending.clientId, code, codeVerifier: pending.codeVerifier, redirectUri: pending.callbackUrl });
          const credential = tokenResponseToCredential(token, pending.clientId, pending.issuer, pending.tokenEndpoint);
          await persistCredential(credential);
          await clearPendingAuth();
          return { content: [{ type: 'text', text: 'Thunder Compute OAuth completed successfully. You can now use tc_list_available_tools and the other tc_* tools.' }], structuredContent: { expires: credential.expires ?? null } };
        } catch (error) { return { isError: true, content: [{ type: 'text', text: withAuthHint(error) }] }; }
      },
    });

    api.registerTool({
      name: 'tc_auth_clear',
      description: 'Clear stored Thunder Compute OAuth credentials and pending auth state.',
      parameters: Type.Object({}),
      async execute() {
        await clearPendingAuth();
        await clearCredential();
        return { content: [{ type: 'text', text: 'Cleared Thunder Compute OAuth state.' }] };
      },
    });

    api.registerTool({ name: 'tc_list_available_tools', description: 'List the tool names exposed by the Thunder Compute MCP server.', parameters: Type.Object({}), async execute() { const client = createClient(api); try { const tools = await client.listTools(); return { content: [{ type: 'text', text: JSON.stringify(tools, null, 2) }], structuredContent: tools }; } catch (error) { return { isError: true, content: [{ type: 'text', text: withAuthHint(error) }] }; } } });
    api.registerTool({ name: 'tc_get_specs', description: 'Get Thunder Compute GPU catalog / specifications.', parameters: Type.Object({ raw: Type.Optional(Type.Any()) }), execute: buildExecute((client, params) => client.callTool('get_specs', params.raw ?? {})) });
    api.registerTool({ name: 'tc_get_availability', description: 'Get live Thunder Compute GPU availability.', parameters: Type.Object({ raw: Type.Optional(Type.Any()) }), execute: buildExecute((client, params) => client.callTool('get_availability', params.raw ?? {})) });
    api.registerTool({ name: 'tc_get_pricing', description: 'Get current Thunder Compute pricing.', parameters: Type.Object({ raw: Type.Optional(Type.Any()) }), execute: buildExecute((client, params) => client.callTool('get_pricing', params.raw ?? {})) });
    api.registerTool({ name: 'tc_list_templates', description: 'List Thunder Compute instance templates / images.', parameters: Type.Object({ raw: Type.Optional(Type.Any()) }), execute: buildExecute((client, params) => client.callTool('list_templates', params.raw ?? {})) });
    api.registerTool({ name: 'tc_list_instances', description: 'List Thunder Compute instances and their current status.', parameters: Type.Object({ raw: Type.Optional(Type.Any()) }), execute: buildExecute((client, params) => client.callTool('list_instances', params.raw ?? {})) });
    api.registerTool({ name: 'tc_create_instance', description: 'Create a Thunder Compute instance.', parameters: InstanceConfig, execute: buildExecute((client, params) => client.callTool('create_instance', buildCreateInstanceArgs(params))) }, { optional: true });
    api.registerTool({ name: 'tc_run_command', description: 'Run a shell command on a Thunder Compute instance.', parameters: Type.Object({ instanceId: Type.String(), command: Type.String(), timeoutSeconds: Type.Optional(Type.Number({ minimum: 1 })), raw: Type.Optional(Type.Any()) }), execute: buildExecute((client, params) => { const args = params.raw ?? { instance_id: params.instanceId, command: params.command, timeout: params.timeoutSeconds }; return client.callTool('run_command', args); }) }, { optional: true });
    api.registerTool({ name: 'tc_delete_instance', description: 'Delete a Thunder Compute instance.', parameters: Type.Object({ instanceId: Type.String(), raw: Type.Optional(Type.Any()) }), execute: buildExecute((client, params) => client.callTool('delete_instance', params.raw ?? { instance_id: params.instanceId })) }, { optional: true });
    api.registerTool({ name: 'tc_create_snapshot', description: 'Create a snapshot of a Thunder Compute instance.', parameters: Type.Object({ instanceId: Type.String(), name: Type.Optional(Type.String()), raw: Type.Optional(Type.Any()) }), execute: buildExecute((client, params) => client.callTool('create_snapshot', params.raw ?? { instance_id: params.instanceId, name: params.name })) }, { optional: true });
    api.registerTool({ name: 'tc_get_upcoming_invoice', description: 'Get the current or upcoming Thunder Compute invoice estimate.', parameters: Type.Object({ raw: Type.Optional(Type.Any()) }), execute: buildExecute((client, params) => client.callTool('get_upcoming_invoice', params.raw ?? {})) });
  },
});
