#!/usr/bin/env node
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { z } from "zod";

const API = "https://inou.com";
const BRIDGE_VERSION = "1.6.1";
const accountId = process.env.INOU_API_TOKEN || process.argv.find(a => a.startsWith("--token="))?.slice(8);

if (!accountId) {
  console.error("Missing INOU_API_TOKEN environment variable");
  process.exit(1);
}

// Version check
let updateAvailable = null;
async function checkVersion() {
  try {
    const res = await fetch(`${API}/api/version`);
    if (res.ok) {
      const data = await res.json();
      if (data.latest_bridge_version && data.latest_bridge_version !== BRIDGE_VERSION) {
        updateAvailable = {
          current: BRIDGE_VERSION,
          latest: data.latest_bridge_version,
          message: data.update_message,
          url: data.download_url
        };
      }
    }
  } catch (e) {
    // Ignore version check errors
  }
}

// Token management (OAuth 2.0 refresh flow)
let accessToken = null;
let tokenExpires = 0;
let refreshTokenValue = accountId; // Initial refresh token from config

async function refreshToken() {
  const res = await fetch(`${API}/oauth/token`, {
    method: "POST",
    headers: { "Content-Type": "application/x-www-form-urlencoded" },
    body: new URLSearchParams({
      grant_type: "refresh_token",
      client_id: "inou-bridge",
      refresh_token: refreshTokenValue
    })
  });
  if (!res.ok) {
    const error = await res.text();
    throw new Error(`Token refresh failed: ${res.status} - ${error}`);
  }
  const data = await res.json();
  accessToken = data.access_token;
  refreshTokenValue = data.refresh_token; // OAuth rotates refresh tokens
  tokenExpires = Date.now() + (data.expires_in * 1000) - 60000; // Refresh 1 min early
}

async function getToken() {
  if (!accessToken || Date.now() > tokenExpires) {
    await refreshToken();
  }
  return accessToken;
}

// API helper with auto-refresh
async function api(path, params = {}) {
  const token = await getToken();
  const url = new URL(path, API);
  Object.entries(params).forEach(([k, v]) => {
    if (v !== undefined && v !== null && v !== "") url.searchParams.set(k, String(v));
  });

  let res = await fetch(url, { headers: { Authorization: `Bearer ${token}` } });

  // If unauthorized, try refreshing token once
  if (res.status === 401) {
    await refreshToken();
    const newToken = await getToken();
    res = await fetch(url, { headers: { Authorization: `Bearer ${newToken}` } });
  }

  if (!res.ok) throw new Error(`HTTP ${res.status}: ${await res.text()}`);
  return res;
}

// JSON response helper
async function json(path, params) {
  const res = await api(path, params);
  return { content: [{ type: "text", text: JSON.stringify(await res.json(), null, 2) }] };
}

// Image response helper
async function image(path, params, mimeType, label) {
  const res = await api(path, params);
  const buffer = await res.arrayBuffer();
  const base64 = Buffer.from(buffer).toString("base64");
  return {
    content: [
      { type: "image", data: base64, mimeType },
      { type: "text", text: `${label} (${Math.round(buffer.byteLength / 1024)}KB)` }
    ]
  };
}

// Create server
const server = new McpServer({ name: "inou-health", version: "1.6.0" });

// Tool annotations - all tools are read-only
const readOnly = { readOnlyHint: true };

// Tools
server.tool("list_dossiers", "List all patient dossiers accessible to this account.", {}, readOnly,
  async () => {
    const result = await json("/api/v1/dossiers");
    if (updateAvailable) {
      result.content.push({
        type: "text",
        text: `\n⚠️ UPDATE AVAILABLE: v${updateAvailable.latest}\n${updateAvailable.message}`
      });
    }
    return result;
  });

server.tool("list_studies", "List all imaging studies for a patient dossier.", {
  dossier: z.string().describe("Dossier ID (16-char hex)")
}, readOnly, async ({ dossier }) => json(`/api/v1/dossiers/${dossier}/entries`, { category: "imaging" }));

server.tool("list_series", "List series for a study. Filter by description (AX, T1, FLAIR, etc).", {
  dossier: z.string().describe("Dossier ID (16-char hex)"),
  study: z.string().describe("Study ID (16-char hex)"),
  filter: z.string().optional().describe("Filter by description")
}, readOnly, async ({ dossier, study, filter }) =>
  json(`/api/v1/dossiers/${dossier}/entries`, { parent: study, filter }));

server.tool("list_slices", "List slices for a series with position info.", {
  dossier: z.string().describe("Dossier ID (16-char hex)"),
  series: z.string().describe("Series ID (16-char hex)")
}, readOnly, async ({ dossier, series }) =>
  json(`/api/v1/dossiers/${dossier}/entries`, { parent: series }));

server.tool("fetch_image", "Fetch slice image as base64 PNG. Optionally set window/level.", {
  dossier: z.string().describe("Dossier ID (16-char hex)"),
  slice: z.string().describe("Slice ID (16-char hex)"),
  wc: z.number().optional().describe("Window center"),
  ww: z.number().optional().describe("Window width")
}, readOnly, async ({ dossier, slice, wc, ww }) =>
  image(`/image/${slice}`, { token: dossier, wc, ww }, "image/png", `Slice ${slice.slice(0, 8)}`));

server.tool("fetch_contact_sheet", "Fetch contact sheet (thumbnail grid) for NAVIGATION ONLY. Use to identify slices, then fetch at full resolution. NEVER diagnose from thumbnails.", {
  dossier: z.string().describe("Dossier ID (16-char hex)"),
  series: z.string().describe("Series ID (16-char hex)"),
  wc: z.number().optional().describe("Window center (e.g., -500 for lung)"),
  ww: z.number().optional().describe("Window width (e.g., 1500 for lung)")
}, readOnly, async ({ dossier, series, wc, ww }) =>
  image(`/contact-sheet.webp/${series}`, { token: dossier, wc, ww }, "image/webp", `Contact sheet ${series.slice(0, 8)}`));

server.tool("list_lab_tests", "List all lab test names for a patient dossier.", {
  dossier: z.string().describe("Dossier ID (16-char hex)")
}, readOnly, async ({ dossier }) => json("/api/labs/tests", { dossier }));

server.tool("get_lab_results", "Get lab results. Specify names, optional date range, optional latest flag.", {
  dossier: z.string().describe("Dossier ID (16-char hex)"),
  names: z.string().describe("Comma-separated test names"),
  from: z.string().optional().describe("Start date YYYY-MM-DD"),
  to: z.string().optional().describe("End date YYYY-MM-DD"),
  latest: z.boolean().optional().describe("Most recent only")
}, readOnly, async ({ dossier, names, from, to, latest }) =>
  json("/api/labs/results", { dossier, names, from, to, latest: latest ? "true" : undefined }));

server.tool("get_categories", "Get observation categories. Without type: top-level. With type=genome: genome categories.", {
  dossier: z.string().describe("Dossier ID (16-char hex)"),
  type: z.string().optional().describe("Observation type (e.g., genome)"),
  category: z.string().optional().describe("Get subcategories within this category")
}, readOnly, async ({ dossier, type, category }) => json("/api/categories", { dossier, type, category }));

server.tool("query_genome", "Query genome variants by gene(s), category, or rsids. Returns matches with magnitude, summary, and categories.", {
  dossier: z.string().describe("Dossier ID (16-char hex)"),
  gene: z.string().optional().describe("Gene name(s), comma-separated (e.g., MTHFR or MTHFR,COMT,CYP2D6)"),
  search: z.string().optional().describe("Search gene, subcategory, or summary"),
  category: z.string().optional().describe("Filter by category"),
  rsids: z.string().optional().describe("Comma-separated rsids"),
  min_magnitude: z.number().optional().describe("Minimum magnitude"),
  include_hidden: z.boolean().optional().describe("Include hidden categories and high-magnitude variants")
}, readOnly, async ({ dossier, gene, search, category, rsids, min_magnitude, include_hidden }) =>
  json("/api/genome", { dossier, gene, search, category, rsids, min_magnitude, include_hidden: include_hidden ? "true" : undefined }));

server.tool("get_version", "Get bridge and server version info.", {}, readOnly,
  async () => {
    let text = `Bridge: ${BRIDGE_VERSION} (Node.js)\nServer: inou.com`;
    if (updateAvailable) {
      text += `\n\n⚠️ UPDATE AVAILABLE: ${updateAvailable.latest}\n${updateAvailable.message}`;
    }
    return { content: [{ type: "text", text }] };
  });

// Run
await checkVersion(); // Check for updates on startup
const transport = new StdioServerTransport();
await server.connect(transport);