// ============================================================================
// @h33/mcp — Attested MCP Server wrapper
//
// Wraps any MCP Server so every tool call, resource read, and optionally
// prompt completion produces a cryptographic receipt. One line to integrate.
// ============================================================================

import { H33Agent, type H33AgentConfig } from '@h33/agent';
import { createHash } from 'crypto';

/**
 * Configuration for an attested MCP server.
 */
export interface AttestedMCPServerConfig {
  /** H33 Agent configuration (attested is forced to true) */
  agent: Omit<H33AgentConfig, 'attested'>;
  /** Whether to attest tool calls (default: true) */
  attestTools?: boolean;
  /** Whether to attest resource reads (default: true) */
  attestResources?: boolean;
  /** Whether to attest prompt completions (default: false) */
  attestPrompts?: boolean;
}

/**
 * Receipt metadata attached non-intrusively to MCP responses.
 */
export interface H33Receipt {
  receipt_id: string;
  node_hash: string;
  verification_url: string;
  replay_ref: string;
}

/**
 * Wraps an MCP Server with H33 attestation. Every tool call, resource
 * read, and optionally prompt completion produces a cryptographic receipt.
 *
 * @example
 * ```typescript
 * import { Server } from '@modelcontextprotocol/sdk/server/index.js';
 * import { attestMCPServer } from '@h33/mcp';
 *
 * const server = new Server(
 *   { name: 'my-server', version: '1.0.0' },
 *   { capabilities: { tools: {} } },
 * );
 *
 * // Wrap with attestation -- one line
 * const attested = attestMCPServer(server, {
 *   agent: {
 *     name: 'MCP Database Server',
 *     canonicalName: 'h33.agent.acme.mcp.database.prod.001',
 *     tenantId: 'acme-corp',
 *   },
 * });
 *
 * // Register tools as normal -- attestation is automatic
 * attested.setRequestHandler('tools/call', async (request) => {
 *   const result = await handleToolCall(request);
 *   return result; // receipt automatically attached
 * });
 * ```
 */
export function attestMCPServer(
  server: any,
  config: AttestedMCPServerConfig,
): any {
  const agent = new H33Agent({
    ...config.agent,
    attested: true,
  });

  const originalSetRequestHandler = server.setRequestHandler.bind(server);
  let started = false;

  async function ensureStarted(): Promise<void> {
    if (!started) {
      await agent.start();
      started = true;
    }
  }

  // Override setRequestHandler to intercept and wrap handlers
  server.setRequestHandler = (method: string, handler: Function) => {
    // ── Tool calls ────────────────────────────────────────────────────
    if (method === 'tools/call' && config.attestTools !== false) {
      originalSetRequestHandler(method, async (request: any) => {
        await ensureStarted();

        const toolName = `h33.tool.mcp.${request.params?.name || 'unknown'}.v1`;
        const inputHash = createHash('sha256')
          .update(JSON.stringify(request.params?.arguments || {}))
          .digest('hex');

        let result: any;
        let receipt: any;

        try {
          result = await handler(request);

          const outputHash = createHash('sha256')
            .update(JSON.stringify(result))
            .digest('hex');

          receipt = await agent.callTool(
            toolName,
            request.params?.arguments,
            result,
            'success',
          );
        } catch (err) {
          receipt = await agent.callTool(
            toolName,
            request.params?.arguments,
            undefined,
            'failed',
          );
          throw err;
        }

        // Attach receipt non-intrusively
        if (result && typeof result === 'object') {
          result._h33_receipt = {
            receipt_id: receipt.receipt_id,
            node_hash: receipt.node_hash,
            verification_url: receipt.verification_url,
            replay_ref: receipt.replay_ref,
          } satisfies Partial<H33Receipt>;
        }

        return result;
      });

    // ── Resource reads ────────────────────────────────────────────────
    } else if (method === 'resources/read' && config.attestResources !== false) {
      originalSetRequestHandler(method, async (request: any) => {
        await ensureStarted();

        const result = await handler(request);
        const uri = request.params?.uri || 'unknown';

        await agent.action(
          'resource_read',
          `Read resource: ${uri}`,
          createHash('sha256').update(uri).digest('hex'),
          result
            ? createHash('sha256').update(JSON.stringify(result)).digest('hex')
            : undefined,
        );

        return result;
      });

    // ── Prompt completions ────────────────────────────────────────────
    } else if (method === 'prompts/get' && config.attestPrompts === true) {
      originalSetRequestHandler(method, async (request: any) => {
        await ensureStarted();

        const result = await handler(request);
        const promptName = request.params?.name || 'unknown';

        await agent.action(
          'prompt_completion',
          `Prompt: ${promptName}`,
          createHash('sha256').update(promptName).digest('hex'),
          result
            ? createHash('sha256').update(JSON.stringify(result)).digest('hex')
            : undefined,
        );

        return result;
      });

    // ── Pass-through for everything else ──────────────────────────────
    } else {
      originalSetRequestHandler(method, handler);
    }
  };

  // ── Shutdown hook ─────────────────────────────────────────────────────────
  const originalClose = server.close?.bind(server);
  server.close = async () => {
    if (started) {
      await agent.stop();
      started = false;
    }
    if (originalClose) {
      await originalClose();
    }
  };

  // Expose agent for advanced usage
  server._h33Agent = agent;

  return server;
}

/**
 * Get the H33 agent from an attested MCP server.
 * Returns null if the server has not been wrapped.
 */
export function getAgent(server: any): H33Agent | null {
  return server._h33Agent || null;
}