import type { GeneratorSpec, ToolSpec } from '@bmm/types';

function toZod(param: ToolSpec['inputSchema'][string]): string {
  const required = param.required !== false;
  let base: string;
  switch (param.type) {
    case 'string':
      base = 'z.string()';
      break;
    case 'number':
      base = 'z.number()';
      break;
    case 'boolean':
      base = 'z.boolean()';
      break;
    case 'array':
      base = 'z.array(z.any())';
      break;
    case 'object':
      base = 'z.record(z.string(), z.any())';
      break;
  }
  if (param.description) base += `.describe(${JSON.stringify(param.description)})`;
  if (!required) base += '.optional()';
  return base;
}

function renderTool(tool: ToolSpec): string {
  const entries = Object.entries(tool.inputSchema)
    .map(([k, v]) => `    ${JSON.stringify(k)}: ${toZod(v)}`)
    .join(',\n');
  const schemaShape = entries ? `{\n${entries}\n  }` : '{}';
  return `server.registerTool(
  ${JSON.stringify(tool.name)},
  {
    title: ${JSON.stringify(tool.name)},
    description: ${JSON.stringify(tool.description)},
    inputSchema: ${schemaShape},
  },
  async (args) => {
    try {
      ${tool.implementation}
    } catch (err) {
      const msg = err instanceof Error ? err.message : String(err);
      return { content: [{ type: 'text', text: 'Error: ' + msg }], isError: true };
    }
  },
);`;
}

export function renderServerCode(spec: GeneratorSpec): string {
  const toolBlocks = spec.tools.map(renderTool).join('\n\n');
  return `// AUTO-GENERATED. Do not edit by hand.
// Generated by BuildMyMCPServer.
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { z } from 'zod';
import Fastify from 'fastify';
import { createRemoteJWKSet, jwtVerify } from 'jose';
import { randomUUID } from 'node:crypto';

function stripTrailingSlash(value) {
  return value.replace(/\\/$/, '');
}

function protectedResourceMetadataUrl(resourceUrl) {
  const url = new URL(resourceUrl);
  const resourcePath = url.pathname === '/' ? '' : url.pathname;
  url.pathname = '/.well-known/oauth-protected-resource' + resourcePath;
  url.hash = '';
  return url.toString();
}

const PUBLIC_URL = stripTrailingSlash(process.env.PUBLIC_URL ?? 'http://localhost:3000');
const CONTROL_PLANE_URL = stripTrailingSlash(process.env.CONTROL_PLANE_URL ?? 'http://host.docker.internal:4000');
const OAUTH_ISSUER = stripTrailingSlash(process.env.OAUTH_ISSUER ?? CONTROL_PLANE_URL + '/oauth');
const MCP_RESOURCE_URL = PUBLIC_URL + '/mcp';
const PROTECTED_RESOURCE_METADATA_URL = protectedResourceMetadataUrl(MCP_RESOURCE_URL);
const EXPECTED_AUDIENCES = Array.from(new Set([MCP_RESOURCE_URL, PUBLIC_URL]));
const PORT = Number.parseInt(process.env.PORT ?? '3000', 10);

const server = new McpServer(
  { name: ${JSON.stringify(spec.name)}, version: '1.0.0' },
  { capabilities: { tools: {}, resources: {}, prompts: {} } },
);

${toolBlocks}

const app = Fastify({ logger: { level: 'info' } });

app.get('/health', async () => ({ ok: true }));

const protectedResourceMetadata = async () => ({
  resource: MCP_RESOURCE_URL,
  authorization_servers: [OAUTH_ISSUER],
  bearer_methods_supported: ['header'],
  scopes_supported: ${JSON.stringify(spec.scopes)},
});

app.get('/.well-known/oauth-protected-resource', protectedResourceMetadata);
app.get('/.well-known/oauth-protected-resource/*', protectedResourceMetadata);

app.get('/.well-known/oauth-authorization-server', async () => {
  const r = await fetch(CONTROL_PLANE_URL + '/.well-known/oauth-authorization-server/oauth');
  return await r.json();
});

const JWKS = createRemoteJWKSet(new URL(CONTROL_PLANE_URL + '/oauth/jwks'));

const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => randomUUID() });

app.all('/mcp', async (request, reply) => {
  const auth = request.headers.authorization;
  if (!auth || !auth.startsWith('Bearer ')) {
    return reply
      .code(401)
      .header('WWW-Authenticate', \`Bearer resource_metadata="\${PROTECTED_RESOURCE_METADATA_URL}"\`)
      .send({ error: 'unauthorized' });
  }
  const token = auth.slice(7);
  try {
    const { payload } = await jwtVerify(token, JWKS, {
      issuer: OAUTH_ISSUER,
      audience: EXPECTED_AUDIENCES,
    });
    const audiences = Array.isArray(payload.aud) ? payload.aud : payload.aud ? [payload.aud] : [];
    if (!audiences.some((aud) => EXPECTED_AUDIENCES.includes(aud))) {
      return reply.code(403).send({ error: 'invalid_audience' });
    }
  } catch (e) {
    request.log.warn({ err: e }, 'token verify failed');
    return reply.code(401).send({ error: 'invalid_token' });
  }
  await transport.handleRequest(request.raw, reply.raw, request.body);
});

await server.connect(transport);
await app.listen({ port: PORT, host: '0.0.0.0' });
app.log.info('mcp server up on :' + PORT);

// suppress unused-z warning when there are no tools
void z;
`;
}