#!/usr/bin/env node
/**
 * MCP image-gen server (thin custom wrapper around OpenRouter).
 *
 * Tools:
 *   - generate_image(prompt, model, aspect)  -> saves PNG to IMAGE_MCP_OUTPUT_DIR
 *   - list_image_models()                    -> 35 image-capable models (1h cache)
 *
 * Auth: Bearer token in Authorization header (token from env MCP_TOKEN).
 * Sourceful Riverflow models branch to /api/v1/images instead of /chat/completions.
 *
 * Design notes (from mcp-image-generation skill, validated 2026-07):
 *   - Stateful StreamableHTTPServerTransport (sessionIdGenerator: crypto.randomUUID)
 *   - try/catch around fetch AND transport.handleRequest
 *   - AbortSignal.timeout(120_000) on upstream
 *   - process.stderr.write on every request/response (SDK swallows errors)
 *   - MODEL_CACHE 1h TTL
 */

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import {
  ListToolsRequestSchema,
  CallToolRequestSchema,
} from '@modelcontextprotocol/sdk/types.js';
import http from 'node:http';
import crypto from 'node:crypto';
import fs from 'node:fs/promises';
import path from 'node:path';

const API_KEY = process.env.OPENROUTER_API_KEY;
const MCP_TOKEN = process.env.MCP_TOKEN;

if (!API_KEY) {
  console.error('FATAL: OPENROUTER_API_KEY env not set');
  process.exit(1);
}
if (!MCP_TOKEN) {
  console.error('FATAL: MCP_TOKEN env not set (required to protect the endpoint)');
  process.exit(1);
}

const PORT = Number(process.env.PORT || 8935);
const OUTPUT_DIR = process.env.IMAGE_MCP_OUTPUT_DIR || '/app/generated-images';

// ---------- OpenRouter model catalog (cached 1h) ----------
let MODEL_CACHE = null;
let MODEL_CACHE_TIME = 0;
const MODEL_TTL_MS = 3600_000;

async function getImageModels() {
  if (MODEL_CACHE && Date.now() - MODEL_CACHE_TIME < MODEL_TTL_MS) {
    return MODEL_CACHE;
  }
  const r = await fetch(
    'https://openrouter.ai/api/v1/models?output_modalities=image',
    { headers: { Authorization: `Bearer ${API_KEY}` } },
  );
  if (!r.ok) throw new Error(`OpenRouter /models ${r.status}`);
  const data = await r.json();
  MODEL_CACHE = data.data.map((m) => m.id);
  MODEL_CACHE_TIME = Date.now();
  process.stderr.write(`[cache] refreshed: ${MODEL_CACHE.length} image models\n`);
  return MODEL_CACHE;
}

function isRiverflow(modelId) {
  return /^sourceful\/riverflow/i.test(modelId);
}

function aspectToOpenRouter(ar) {
  // OpenRouter accepts 1:1 / 16:9 / 9:16 etc. via image_config.aspect_ratio (vendor-specific, optional).
  // To stay safe we just include the keyword in the prompt — let the model decide when the
  // vendor-specific field is unsupported.
  if (ar === 'square') return '1:1';
  if (ar === 'landscape') return '16:9';
  if (ar === 'portrait') return '9:16';
  return ar;
}

// ---------- API call router ----------
async function callOpenRouter({ prompt, model, aspect }) {
  // Riverflow models: OpenRouter exposes /api/v1/images for them (chat endpoint returns 404
  // even though models appear in the ?output_modalities=image list — verified 2026-07-05).
  if (isRiverflow(model)) {
    return await callRiverflowImages({ prompt, model, aspect });
  }
  return await callChatCompletions({ prompt, model, aspect });
}

async function callChatCompletions({ prompt, model, aspect }) {
  const ar = aspectToOpenRouter(aspect);
  const augmentedPrompt = ar ? `${prompt} (aspect ratio ${ar})` : prompt;
  const resp = await fetch('https://openrouter.ai/api/v1/chat/completions', {
    method: 'POST',
    headers: {
      Authorization: `Bearer ${API_KEY}`,
      'Content-Type': 'application/json',
    },
    body: JSON.stringify({
      model,
      messages: [{ role: 'user', content: augmentedPrompt }],
      modalities: ['image', 'text'],
    }),
    signal: AbortSignal.timeout(120_000),
  });
  if (!resp.ok) {
    const body = await resp.text();
    throw new Error(`OpenRouter chat ${resp.status}: ${body.slice(0, 600)}`);
  }
  const data = await resp.json();
  const images = data.choices?.[0]?.message?.images;
  if (!images?.length) {
    const text = data.choices?.[0]?.message?.content;
    throw new Error(
      `No images returned. Text content: ${JSON.stringify(text).slice(0, 300)}`,
    );
  }
  // Pick first image_url entry that has a data: URL
  const imgEntry = images.find(
    (e) => e?.image_url?.url?.startsWith('data:image/'),
  );
  if (!imgEntry) {
    throw new Error('images[] has no data: URLs (unsupported shape)');
  }
  return {
    dataUrl: imgEntry.image_url.url,
    provider: 'chat_completions',
  };
}

async function callRiverflowImages({ prompt, model, aspect }) {
  // OpenRouter /api/v1/images accepts a Shape similar to OpenAI Images API.
  // Token is sent as Bearer (same as chat completions).
  const sizeMap = {
    square: '1024x1024',
    landscape: '1536x1024',
    portrait: '1024x1536',
  };
  const resp = await fetch('https://openrouter.ai/api/v1/images', {
    method: 'POST',
    headers: {
      Authorization: `Bearer ${API_KEY}`,
      'Content-Type': 'application/json',
    },
    body: JSON.stringify({
      model,
      prompt,
      size: sizeMap[aspect] || '1024x1024',
      n: 1,
    }),
    signal: AbortSignal.timeout(120_000),
  });
  if (!resp.ok) {
    const body = await resp.text();
    throw new Error(`OpenRouter images ${resp.status}: ${body.slice(0, 600)}`);
  }
  const data = await resp.json();
  // OpenAI-style response: { data: [{ url, b64_json }] }
  const item = data?.data?.[0];
  if (!item) throw new Error('Riverflow returned empty data[]');
  // Prefer b64_json; fall back to URL fetch
  if (item.b64_json) {
    const mime = 'image/png';
    return { dataUrl: `data:${mime};base64,${item.b64_json}`, provider: 'images_ep' };
  }
  if (item.url) {
    const fetched = await fetch(item.url, { signal: AbortSignal.timeout(60_000) });
    if (!fetched.ok) throw new Error(`Riverflow URL fetch ${fetched.status}`);
    const buf = Buffer.from(await fetched.arrayBuffer());
    return {
      dataUrl: `data:image/png;base64,${buf.toString('base64')}`,
      provider: 'images_ep',
    };
  }
  throw new Error('Riverflow response has no url or b64_json');
}

// ---------- MCP server ----------
const server = new Server(
  { name: 'mcp-img-gen', version: '1.0.0' },
  { capabilities: { tools: {} } },
);

server.setRequestHandler(ListToolsRequestSchema, async () => {
  const models = await getImageModels();
  return {
    tools: [
      {
        name: 'generate_image',
        description:
          'Generate an image via OpenRouter. Aspect: square|landscape|portrait. ' +
          'Returns {model,filepath,bytes,provider,latency_ms}.',
        inputSchema: {
          type: 'object',
          properties: {
            prompt: { type: 'string', description: 'Text prompt' },
            model: {
              type: 'string',
              enum: models,
              default: 'google/gemini-2.5-flash-image',
            },
            aspect: {
              type: 'string',
              enum: ['square', 'landscape', 'portrait'],
              default: 'square',
            },
          },
          required: ['prompt'],
        },
      },
      {
        name: 'list_image_models',
        description: 'List 35 image-capable models on OpenRouter (cached 1h).',
        inputSchema: { type: 'object', properties: {} },
      },
    ],
  };
});

server.setRequestHandler(CallToolRequestSchema, async (req) => {
  const { name, arguments: args } = req.params;
  try {
    if (name === 'list_image_models') {
      const list = (await getImageModels()).join('\n');
      return {
        content: [{ type: 'text', text: `# OpenRouter image models (${MODEL_CACHE.length})\n\n${list}` }],
      };
    }
    if (name === 'generate_image') {
      const prompt = String(args?.prompt || '').trim();
      if (!prompt) return errResp('prompt required');
      const model = args?.model || 'google/gemini-2.5-flash-image';
      const aspect = args?.aspect || 'square';

      process.stderr.write(`[gen] req model=${model} aspect=${aspect} prompt="${prompt.slice(0, 80)}"\n`);
      const t0 = Date.now();
      const { dataUrl, provider } = await callOpenRouter({ prompt, model, aspect });
      const base64 = dataUrl.replace(/^data:image\/\w+;base64,/, '');
      const bytes = Buffer.from(base64, 'base64');
      await fs.mkdir(OUTPUT_DIR, { recursive: true });
      const safeModel = model.replace(/[/:]/g, '_');
      const filename = `${safeModel}-${Date.now()}.png`;
      const filepath = path.join(OUTPUT_DIR, filename);
      await fs.writeFile(filepath, bytes);
      process.stderr.write(`[gen] OK provider=${provider} ${Date.now() - t0}ms ${filepath} ${bytes.length}b\n`);
      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify(
              {
                model,
                provider,
                aspect,
                filepath,
                filename,
                bytes: bytes.length,
                latency_ms: Date.now() - t0,
              },
              null,
              2,
            ),
          },
        ],
      };
    }
    return errResp(`Unknown tool: ${name}`);
  } catch (e) {
    process.stderr.write(`[gen-ERR] ${e.stack || e}\n`);
    return {
      content: [{ type: 'text', text: `Error: ${e.message}` }],
      isError: true,
    };
  }
});

function errResp(text) {
  return { content: [{ type: 'text', text }], isError: true };
}

// ---------- HTTP transport w/ Bearer auth ----------
const transport = new StreamableHTTPServerTransport({
  sessionIdGenerator: () => crypto.randomUUID(),
});
await server.connect(transport);

const httpServer = http.createServer(async (req, res) => {
  // Auth gate (Skips the public /health endpoint).
  if (!req.url.startsWith('/health')) {
    const auth = req.headers.authorization || '';
    const expected = `Bearer ${MCP_TOKEN}`;
    if (auth !== expected) {
      process.stderr.write(`[auth] denied ${req.method} ${req.url}\n`);
      res.statusCode = 401;
      res.setHeader('www-authenticate', 'Bearer');
      res.setHeader('content-type', 'application/json');
      res.end(JSON.stringify({ error: 'invalid or missing token' }));
      return;
    }
  }

  // Lightweight /health for Coolify / Traefik probes
  if (req.url === '/health') {
    res.statusCode = 200;
    res.setHeader('content-type', 'application/json');
    res.end(JSON.stringify({ ok: true, version: '1.0.0', uptime: process.uptime() }));
    return;
  }

  process.stderr.write(`[HTTP] ${req.method} ${req.url} -> incoming\n`);
  try {
    await transport.handleRequest(req, res);
    process.stderr.write(`[HTTP] ${req.method} ${req.url} -> ${res.statusCode}\n`);
  } catch (e) {
    process.stderr.write(`[HTTP-ERR] ${e.stack || e}\n`);
    if (!res.headersSent) {
      res.statusCode = 500;
      res.setHeader('content-type', 'application/json');
      res.end(JSON.stringify({ error: e.message }));
    }
  }
});

httpServer.listen(PORT, '0.0.0.0', () => {
  process.stderr.write(`mcp-img-gen listening on http://0.0.0.0:${PORT}\n`);
});