#!/usr/bin/env node
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import {
  CallToolRequestSchema,
  ListToolsRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";

// ============================================
// CONFIGURATION
// ============================================
const MCP_NAME = "compliance-grc";
const MCP_VERSION = "1.0.0";
const VANTA_BASE_URL = "https://api.vanta.com";
const DRATA_BASE_URL = "https://public-api.drata.com";

// ============================================
// VANTA API CLIENT - OAuth2 client_credentials
// ============================================
class VantaClient {
  private clientId: string;
  private clientSecret: string;
  private baseUrl: string;
  private accessToken: string | null = null;
  private tokenExpiry: number = 0;

  constructor(clientId: string, clientSecret: string) {
    this.clientId = clientId;
    this.clientSecret = clientSecret;
    this.baseUrl = VANTA_BASE_URL;
  }

  async getAccessToken(): Promise<string> {
    // Check if we have a valid token
    if (this.accessToken && Date.now() < this.tokenExpiry) {
      return this.accessToken;
    }

    // Get new token using client_credentials flow
    const tokenUrl = `${this.baseUrl}/oauth/token`;
    const response = await fetch(tokenUrl, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify({
        grant_type: "client_credentials",
        client_id: this.clientId,
        client_secret: this.clientSecret,
      }),
    });

    if (!response.ok) {
      const text = await response.text();
      throw new Error(`Vanta OAuth error: ${response.status} ${response.statusText} - ${text}`);
    }

    const data = await response.json();
    this.accessToken = data.access_token as string;
    // Set expiry to 5 minutes before actual expiry for safety
    this.tokenExpiry = Date.now() + (data.expires_in - 300) * 1000;
    return this.accessToken!;
  }

  async request(endpoint: string, options: RequestInit = {}) {
    const token = await this.getAccessToken();
    const url = `${this.baseUrl}${endpoint}`;
    const response = await fetch(url, {
      ...options,
      headers: {
        Authorization: `Bearer ${token}`,
        "Content-Type": "application/json",
        "Accept": "application/json",
        ...options.headers,
      },
    });

    if (!response.ok) {
      const text = await response.text();
      throw new Error(`Vanta API error: ${response.status} ${response.statusText} - ${text}`);
    }

    if (response.status === 204) {
      return { success: true };
    }

    return response.json();
  }

  async get(endpoint: string) {
    return this.request(endpoint, { method: "GET" });
  }

  async post(endpoint: string, data: any) {
    return this.request(endpoint, {
      method: "POST",
      body: JSON.stringify(data),
    });
  }
}

// ============================================
// DRATA API CLIENT - Bearer token
// ============================================
class DrataClient {
  private apiKey: string;
  private baseUrl: string;

  constructor(apiKey: string) {
    this.apiKey = apiKey;
    this.baseUrl = DRATA_BASE_URL;
  }

  async request(endpoint: string, options: RequestInit = {}) {
    const url = `${this.baseUrl}${endpoint}`;
    const response = await fetch(url, {
      ...options,
      headers: {
        Authorization: `Bearer ${this.apiKey}`,
        "Content-Type": "application/json",
        "Accept": "application/json",
        ...options.headers,
      },
    });

    if (!response.ok) {
      const text = await response.text();
      throw new Error(`Drata API error: ${response.status} ${response.statusText} - ${text}`);
    }

    if (response.status === 204) {
      return { success: true };
    }

    return response.json();
  }

  async get(endpoint: string) {
    return this.request(endpoint, { method: "GET" });
  }

  async post(endpoint: string, data: any) {
    return this.request(endpoint, {
      method: "POST",
      body: JSON.stringify(data),
    });
  }

  async put(endpoint: string, data: any) {
    return this.request(endpoint, {
      method: "PUT",
      body: JSON.stringify(data),
    });
  }
}

// ============================================
// TOOL DEFINITIONS
// ============================================
const tools = [
  // ========== VANTA TOOLS ==========
  {
    name: "list_vanta_controls",
    description: "List all compliance controls in Vanta (SOC2, HIPAA, GDPR, etc.)",
    inputSchema: {
      type: "object" as const,
      properties: {
        framework: { type: "string", description: "Filter by framework (soc2, hipaa, gdpr, iso27001)" },
        status: { type: "string", description: "Filter by status (passing, failing, not_applicable)" },
        limit: { type: "number", description: "Number of results to return" },
        offset: { type: "number", description: "Pagination offset" },
      },
    },
  },
  {
    name: "get_vanta_control",
    description: "Get detailed information about a specific Vanta control",
    inputSchema: {
      type: "object" as const,
      properties: {
        controlId: { type: "string", description: "Control ID" },
      },
      required: ["controlId"],
    },
  },
  {
    name: "list_vanta_tests",
    description: "List compliance tests and their results in Vanta",
    inputSchema: {
      type: "object" as const,
      properties: {
        controlId: { type: "string", description: "Filter by control ID" },
        status: { type: "string", description: "Filter by status (passing, failing, pending)" },
        limit: { type: "number", description: "Number of results to return" },
      },
    },
  },
  {
    name: "list_vanta_vulnerabilities",
    description: "List security vulnerabilities detected by Vanta",
    inputSchema: {
      type: "object" as const,
      properties: {
        severity: { type: "string", description: "Filter by severity (critical, high, medium, low)" },
        status: { type: "string", description: "Filter by status (open, resolved, ignored)" },
        source: { type: "string", description: "Filter by source (aws, github, gcp, etc.)" },
        limit: { type: "number", description: "Number of results to return" },
      },
    },
  },
  {
    name: "list_vanta_evidence",
    description: "List compliance evidence collected by Vanta",
    inputSchema: {
      type: "object" as const,
      properties: {
        controlId: { type: "string", description: "Filter by control ID" },
        type: { type: "string", description: "Filter by evidence type (screenshot, document, integration)" },
        startDate: { type: "string", description: "Filter by start date (YYYY-MM-DD)" },
        endDate: { type: "string", description: "Filter by end date (YYYY-MM-DD)" },
        limit: { type: "number", description: "Number of results to return" },
      },
    },
  },
  {
    name: "list_vanta_users",
    description: "List users and their access levels in Vanta",
    inputSchema: {
      type: "object" as const,
      properties: {
        role: { type: "string", description: "Filter by role (admin, member, viewer)" },
        status: { type: "string", description: "Filter by status (active, inactive)" },
        limit: { type: "number", description: "Number of results to return" },
      },
    },
  },
  {
    name: "list_vanta_integrations",
    description: "List active integrations in Vanta (AWS, GitHub, Google Workspace, etc.)",
    inputSchema: {
      type: "object" as const,
      properties: {
        type: { type: "string", description: "Filter by integration type (aws, github, gsuite, okta, etc.)" },
        status: { type: "string", description: "Filter by status (active, inactive, error)" },
      },
    },
  },
  {
    name: "get_vanta_compliance_status",
    description: "Get overall compliance status and readiness summary from Vanta",
    inputSchema: {
      type: "object" as const,
      properties: {
        framework: { type: "string", description: "Framework to check (soc2, hipaa, gdpr, iso27001)" },
      },
    },
  },

  // ========== DRATA TOOLS ==========
  {
    name: "list_drata_controls",
    description: "List all compliance controls in Drata",
    inputSchema: {
      type: "object" as const,
      properties: {
        framework: { type: "string", description: "Filter by framework (soc2, hipaa, gdpr, iso27001)" },
        status: { type: "string", description: "Filter by status (active, inactive, satisfied, unsatisfied)" },
        page: { type: "number", description: "Page number for pagination" },
        perPage: { type: "number", description: "Results per page (default 50, max 100)" },
      },
    },
  },
  {
    name: "get_drata_control",
    description: "Get detailed information about a specific Drata control",
    inputSchema: {
      type: "object" as const,
      properties: {
        controlId: { type: "string", description: "Control ID" },
      },
      required: ["controlId"],
    },
  },
  {
    name: "list_drata_personnel",
    description: "List personnel/employees tracked in Drata",
    inputSchema: {
      type: "object" as const,
      properties: {
        status: { type: "string", description: "Filter by status (active, inactive, offboarded)" },
        department: { type: "string", description: "Filter by department" },
        role: { type: "string", description: "Filter by role" },
        page: { type: "number", description: "Page number for pagination" },
        perPage: { type: "number", description: "Results per page" },
      },
    },
  },
  {
    name: "list_drata_vendors",
    description: "List third-party vendors tracked in Drata",
    inputSchema: {
      type: "object" as const,
      properties: {
        riskLevel: { type: "string", description: "Filter by risk level (high, medium, low)" },
        status: { type: "string", description: "Filter by status (active, inactive)" },
        page: { type: "number", description: "Page number for pagination" },
        perPage: { type: "number", description: "Results per page" },
      },
    },
  },
  {
    name: "list_drata_risks",
    description: "List identified risks in Drata",
    inputSchema: {
      type: "object" as const,
      properties: {
        severity: { type: "string", description: "Filter by severity (critical, high, medium, low)" },
        status: { type: "string", description: "Filter by status (open, in_progress, resolved, accepted)" },
        category: { type: "string", description: "Filter by category (security, privacy, operational)" },
        page: { type: "number", description: "Page number for pagination" },
        perPage: { type: "number", description: "Results per page" },
      },
    },
  },
  {
    name: "list_drata_assets",
    description: "List IT assets tracked in Drata (servers, databases, applications)",
    inputSchema: {
      type: "object" as const,
      properties: {
        type: { type: "string", description: "Filter by asset type (server, database, application, device)" },
        status: { type: "string", description: "Filter by status (active, inactive)" },
        environment: { type: "string", description: "Filter by environment (production, staging, development)" },
        page: { type: "number", description: "Page number for pagination" },
        perPage: { type: "number", description: "Results per page" },
      },
    },
  },
  {
    name: "list_drata_frameworks",
    description: "List compliance frameworks configured in Drata",
    inputSchema: {
      type: "object" as const,
      properties: {
        status: { type: "string", description: "Filter by status (active, inactive)" },
      },
    },
  },
  {
    name: "list_drata_evidence",
    description: "List compliance evidence collected by Drata",
    inputSchema: {
      type: "object" as const,
      properties: {
        controlId: { type: "string", description: "Filter by control ID" },
        type: { type: "string", description: "Filter by evidence type" },
        startDate: { type: "string", description: "Filter by start date (YYYY-MM-DD)" },
        endDate: { type: "string", description: "Filter by end date (YYYY-MM-DD)" },
        page: { type: "number", description: "Page number for pagination" },
        perPage: { type: "number", description: "Results per page" },
      },
    },
  },
  {
    name: "get_drata_compliance_summary",
    description: "Get overall compliance summary and readiness from Drata",
    inputSchema: {
      type: "object" as const,
      properties: {
        framework: { type: "string", description: "Framework to check (soc2, hipaa, gdpr, iso27001)" },
      },
    },
  },
];

// ============================================
// TOOL HANDLERS
// ============================================
async function handleVantaTool(client: VantaClient, name: string, args: any) {
  switch (name) {
    case "list_vanta_controls": {
      const params = new URLSearchParams();
      if (args.framework) params.append("framework", args.framework);
      if (args.status) params.append("status", args.status);
      if (args.limit) params.append("limit", String(args.limit));
      if (args.offset) params.append("offset", String(args.offset));
      const query = params.toString();
      return await client.get(`/v1/controls${query ? `?${query}` : ""}`);
    }

    case "get_vanta_control": {
      return await client.get(`/v1/controls/${args.controlId}`);
    }

    case "list_vanta_tests": {
      const params = new URLSearchParams();
      if (args.controlId) params.append("control_id", args.controlId);
      if (args.status) params.append("status", args.status);
      if (args.limit) params.append("limit", String(args.limit));
      const query = params.toString();
      return await client.get(`/v1/tests${query ? `?${query}` : ""}`);
    }

    case "list_vanta_vulnerabilities": {
      const params = new URLSearchParams();
      if (args.severity) params.append("severity", args.severity);
      if (args.status) params.append("status", args.status);
      if (args.source) params.append("source", args.source);
      if (args.limit) params.append("limit", String(args.limit));
      const query = params.toString();
      return await client.get(`/v1/vulnerabilities${query ? `?${query}` : ""}`);
    }

    case "list_vanta_evidence": {
      const params = new URLSearchParams();
      if (args.controlId) params.append("control_id", args.controlId);
      if (args.type) params.append("type", args.type);
      if (args.startDate) params.append("start_date", args.startDate);
      if (args.endDate) params.append("end_date", args.endDate);
      if (args.limit) params.append("limit", String(args.limit));
      const query = params.toString();
      return await client.get(`/v1/evidence${query ? `?${query}` : ""}`);
    }

    case "list_vanta_users": {
      const params = new URLSearchParams();
      if (args.role) params.append("role", args.role);
      if (args.status) params.append("status", args.status);
      if (args.limit) params.append("limit", String(args.limit));
      const query = params.toString();
      return await client.get(`/v1/users${query ? `?${query}` : ""}`);
    }

    case "list_vanta_integrations": {
      const params = new URLSearchParams();
      if (args.type) params.append("type", args.type);
      if (args.status) params.append("status", args.status);
      const query = params.toString();
      return await client.get(`/v1/integrations${query ? `?${query}` : ""}`);
    }

    case "get_vanta_compliance_status": {
      const params = new URLSearchParams();
      if (args.framework) params.append("framework", args.framework);
      const query = params.toString();
      return await client.get(`/v1/compliance/status${query ? `?${query}` : ""}`);
    }

    default:
      throw new Error(`Unknown Vanta tool: ${name}`);
  }
}

async function handleDrataTool(client: DrataClient, name: string, args: any) {
  switch (name) {
    case "list_drata_controls": {
      const params = new URLSearchParams();
      if (args.framework) params.append("framework", args.framework);
      if (args.status) params.append("status", args.status);
      if (args.page) params.append("page", String(args.page));
      if (args.perPage) params.append("per_page", String(args.perPage));
      const query = params.toString();
      return await client.get(`/v1/controls${query ? `?${query}` : ""}`);
    }

    case "get_drata_control": {
      return await client.get(`/v1/controls/${args.controlId}`);
    }

    case "list_drata_personnel": {
      const params = new URLSearchParams();
      if (args.status) params.append("status", args.status);
      if (args.department) params.append("department", args.department);
      if (args.role) params.append("role", args.role);
      if (args.page) params.append("page", String(args.page));
      if (args.perPage) params.append("per_page", String(args.perPage));
      const query = params.toString();
      return await client.get(`/v1/personnel${query ? `?${query}` : ""}`);
    }

    case "list_drata_vendors": {
      const params = new URLSearchParams();
      if (args.riskLevel) params.append("risk_level", args.riskLevel);
      if (args.status) params.append("status", args.status);
      if (args.page) params.append("page", String(args.page));
      if (args.perPage) params.append("per_page", String(args.perPage));
      const query = params.toString();
      return await client.get(`/v1/vendors${query ? `?${query}` : ""}`);
    }

    case "list_drata_risks": {
      const params = new URLSearchParams();
      if (args.severity) params.append("severity", args.severity);
      if (args.status) params.append("status", args.status);
      if (args.category) params.append("category", args.category);
      if (args.page) params.append("page", String(args.page));
      if (args.perPage) params.append("per_page", String(args.perPage));
      const query = params.toString();
      return await client.get(`/v1/risks${query ? `?${query}` : ""}`);
    }

    case "list_drata_assets": {
      const params = new URLSearchParams();
      if (args.type) params.append("type", args.type);
      if (args.status) params.append("status", args.status);
      if (args.environment) params.append("environment", args.environment);
      if (args.page) params.append("page", String(args.page));
      if (args.perPage) params.append("per_page", String(args.perPage));
      const query = params.toString();
      return await client.get(`/v1/assets${query ? `?${query}` : ""}`);
    }

    case "list_drata_frameworks": {
      const params = new URLSearchParams();
      if (args.status) params.append("status", args.status);
      const query = params.toString();
      return await client.get(`/v1/frameworks${query ? `?${query}` : ""}`);
    }

    case "list_drata_evidence": {
      const params = new URLSearchParams();
      if (args.controlId) params.append("control_id", args.controlId);
      if (args.type) params.append("type", args.type);
      if (args.startDate) params.append("start_date", args.startDate);
      if (args.endDate) params.append("end_date", args.endDate);
      if (args.page) params.append("page", String(args.page));
      if (args.perPage) params.append("per_page", String(args.perPage));
      const query = params.toString();
      return await client.get(`/v1/evidence${query ? `?${query}` : ""}`);
    }

    case "get_drata_compliance_summary": {
      const params = new URLSearchParams();
      if (args.framework) params.append("framework", args.framework);
      const query = params.toString();
      return await client.get(`/v1/compliance/summary${query ? `?${query}` : ""}`);
    }

    default:
      throw new Error(`Unknown Drata tool: ${name}`);
  }
}

// ============================================
// SERVER SETUP
// ============================================
async function main() {
  // Check Vanta credentials
  const vantaClientId = process.env.VANTA_CLIENT_ID;
  const vantaClientSecret = process.env.VANTA_CLIENT_SECRET;
  const vantaEnabled = Boolean(vantaClientId && vantaClientSecret);

  // Check Drata credentials
  const drataApiKey = process.env.DRATA_API_KEY;
  const drataEnabled = Boolean(drataApiKey);

  // At least one service must be configured
  if (!vantaEnabled && !drataEnabled) {
    console.error("Error: At least one service must be configured:");
    console.error("  - Vanta: Set VANTA_CLIENT_ID and VANTA_CLIENT_SECRET");
    console.error("  - Drata: Set DRATA_API_KEY");
    process.exit(1);
  }

  // Initialize clients
  const vantaClient = vantaEnabled ? new VantaClient(vantaClientId!, vantaClientSecret!) : null;
  const drataClient = drataEnabled ? new DrataClient(drataApiKey!) : null;

  // Filter tools based on available services
  const availableTools = tools.filter(tool => {
    if (tool.name.startsWith("list_vanta_") || tool.name.startsWith("get_vanta_")) {
      return vantaEnabled;
    }
    if (tool.name.startsWith("list_drata_") || tool.name.startsWith("get_drata_")) {
      return drataEnabled;
    }
    return false;
  });

  const server = new Server(
    { name: `${MCP_NAME}-mcp`, version: MCP_VERSION },
    { capabilities: { tools: {} } }
  );

  server.setRequestHandler(ListToolsRequestSchema, async () => ({
    tools: availableTools,
  }));

  server.setRequestHandler(CallToolRequestSchema, async (request) => {
    const { name, arguments: args } = request.params;
    
    try {
      let result;

      // Route to appropriate handler
      if (name.startsWith("list_vanta_") || name.startsWith("get_vanta_")) {
        if (!vantaClient) {
          throw new Error("Vanta is not configured. Set VANTA_CLIENT_ID and VANTA_CLIENT_SECRET.");
        }
        result = await handleVantaTool(vantaClient, name, args || {});
      } else if (name.startsWith("list_drata_") || name.startsWith("get_drata_")) {
        if (!drataClient) {
          throw new Error("Drata is not configured. Set DRATA_API_KEY.");
        }
        result = await handleDrataTool(drataClient, name, args || {});
      } else {
        throw new Error(`Unknown tool: ${name}`);
      }

      return {
        content: [{ type: "text", text: JSON.stringify(result, null, 2) }],
      };
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      return {
        content: [{ type: "text", text: `Error: ${message}` }],
        isError: true,
      };
    }
  });

  const transport = new StdioServerTransport();
  await server.connect(transport);
  
  const services = [];
  if (vantaEnabled) services.push("Vanta");
  if (drataEnabled) services.push("Drata");
  console.error(`${MCP_NAME} MCP server running on stdio (${services.join(" + ")})`);
}

main().catch(console.error);
