#!/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 = "product-analytics";
const MCP_VERSION = "1.0.0";

// ============================================
// API CLIENTS
// ============================================

// Mixpanel Client - Query API uses basic auth, Ingestion uses project token
class MixpanelClient {
  private projectId: string;
  private serviceAccountSecret: string;
  private queryBaseUrl = "https://mixpanel.com/api/2.0";
  private ingestionBaseUrl = "https://api.mixpanel.com";

  constructor(projectId: string, serviceAccountSecret: string) {
    this.projectId = projectId;
    this.serviceAccountSecret = serviceAccountSecret;
  }

  async request(endpoint: string, options: RequestInit = {}, useIngestionUrl = false) {
    const baseUrl = useIngestionUrl ? this.ingestionBaseUrl : this.queryBaseUrl;
    const url = `${baseUrl}${endpoint}`;
    
    // Basic auth: project_id:secret (base64)
    const authString = Buffer.from(`${this.projectId}:${this.serviceAccountSecret}`).toString('base64');
    
    const response = await fetch(url, {
      ...options,
      headers: {
        "Authorization": `Basic ${authString}`,
        "Content-Type": "application/json",
        "Accept": "application/json",
        ...options.headers,
      },
    });

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

    return response.json();
  }

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

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

// Amplitude Client - Uses API Key + Secret Key
class AmplitudeClient {
  private apiKey: string;
  private secretKey: string;
  private baseUrl = "https://amplitude.com/api/2";

  constructor(apiKey: string, secretKey: string) {
    this.apiKey = apiKey;
    this.secretKey = secretKey;
  }

  async request(endpoint: string, options: RequestInit = {}) {
    const url = `${this.baseUrl}${endpoint}`;
    
    // Basic auth: api_key:secret_key (base64)
    const authString = Buffer.from(`${this.apiKey}:${this.secretKey}`).toString('base64');
    
    const response = await fetch(url, {
      ...options,
      headers: {
        "Authorization": `Basic ${authString}`,
        "Content-Type": "application/json",
        "Accept": "application/json",
        ...options.headers,
      },
    });

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

    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),
    });
  }
}

// PostHog Client - Uses Personal API Key
class PostHogClient {
  private apiKey: string;
  private projectId: string;
  private baseUrl = "https://app.posthog.com/api";

  constructor(apiKey: string, projectId: string) {
    this.apiKey = apiKey;
    this.projectId = projectId;
  }

  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(`PostHog API error: ${response.status} ${response.statusText} - ${text}`);
    }

    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),
    });
  }
}

// ============================================
// TOOL DEFINITIONS
// ============================================
const tools = [
  // ========== MIXPANEL TOOLS ==========
  {
    name: "query_mixpanel_segmentation",
    description: "Query Mixpanel segmentation data for event analysis over time",
    inputSchema: {
      type: "object" as const,
      properties: {
        event: { type: "string", description: "Event name to analyze" },
        type: { type: "string", description: "Analysis type: general, unique, average" },
        unit: { type: "string", description: "Time unit: minute, hour, day, week, month" },
        from_date: { type: "string", description: "Start date (YYYY-MM-DD)" },
        to_date: { type: "string", description: "End date (YYYY-MM-DD)" },
        where: { type: "string", description: "Event property filter (JSON string)" },
        on: { type: "string", description: "Property to segment by" },
        limit: { type: "number", description: "Limit number of segments returned" },
      },
      required: ["event", "from_date", "to_date"],
    },
  },
  {
    name: "query_mixpanel_funnels",
    description: "Query Mixpanel funnel conversion data",
    inputSchema: {
      type: "object" as const,
      properties: {
        events: { 
          type: "array", 
          description: "Array of event names defining funnel steps",
          items: { type: "string" }
        },
        unit: { type: "string", description: "Time unit: day, week, month" },
        from_date: { type: "string", description: "Start date (YYYY-MM-DD)" },
        to_date: { type: "string", description: "End date (YYYY-MM-DD)" },
        funnel_window_days: { type: "number", description: "Window for conversion (days)" },
        on: { type: "string", description: "Property to segment by" },
        where: { type: "string", description: "Event property filter (JSON string)" },
      },
      required: ["events", "from_date", "to_date"],
    },
  },
  {
    name: "query_mixpanel_retention",
    description: "Query Mixpanel retention/cohort analysis",
    inputSchema: {
      type: "object" as const,
      properties: {
        from_date: { type: "string", description: "Start date (YYYY-MM-DD)" },
        to_date: { type: "string", description: "End date (YYYY-MM-DD)" },
        retention_type: { type: "string", description: "Type: birth or compounded" },
        born_event: { type: "string", description: "Event defining cohort entry" },
        event: { type: "string", description: "Event defining retention" },
        born_where: { type: "string", description: "Filter for born event (JSON)" },
        where: { type: "string", description: "Filter for retention event (JSON)" },
        interval: { type: "number", description: "Retention interval in days" },
        interval_count: { type: "number", description: "Number of intervals to track" },
        unit: { type: "string", description: "Time unit: day, week, month" },
        on: { type: "string", description: "Property to segment by" },
      },
      required: ["from_date", "to_date"],
    },
  },
  {
    name: "get_mixpanel_user_profile",
    description: "Get Mixpanel user profile by distinct_id",
    inputSchema: {
      type: "object" as const,
      properties: {
        distinct_id: { type: "string", description: "User distinct ID" },
      },
      required: ["distinct_id"],
    },
  },
  {
    name: "query_mixpanel_jql",
    description: "Run custom Mixpanel JQL (JavaScript Query Language) query",
    inputSchema: {
      type: "object" as const,
      properties: {
        script: { type: "string", description: "JQL script to execute" },
        params: { type: "object", description: "Query parameters" },
      },
      required: ["script"],
    },
  },
  {
    name: "export_mixpanel_events",
    description: "Export raw Mixpanel events for a date range",
    inputSchema: {
      type: "object" as const,
      properties: {
        from_date: { type: "string", description: "Start date (YYYY-MM-DD)" },
        to_date: { type: "string", description: "End date (YYYY-MM-DD)" },
        event: { type: "array", items: { type: "string" }, description: "Filter by event names" },
        where: { type: "string", description: "Event property filter (JSON string)" },
        limit: { type: "number", description: "Limit number of events (max 1000 per request)" },
      },
      required: ["from_date", "to_date"],
    },
  },
  {
    name: "list_mixpanel_events",
    description: "List all event names in Mixpanel project",
    inputSchema: {
      type: "object" as const,
      properties: {
        type: { type: "string", description: "Event type: general or all" },
      },
    },
  },

  // ========== AMPLITUDE TOOLS ==========
  {
    name: "query_amplitude_segmentation",
    description: "Query Amplitude event segmentation data",
    inputSchema: {
      type: "object" as const,
      properties: {
        e: { type: "object", description: "Event definition (JSON object with event_type, filters)" },
        start: { type: "string", description: "Start date (YYYYMMDD)" },
        end: { type: "string", description: "End date (YYYYMMDD)" },
        m: { type: "string", description: "Metrics: uniques, totals, average, etc." },
        i: { type: "number", description: "Interval: 1 (hour), 7 (day), 30 (week), -1 (all)" },
        s: { type: "array", items: { type: "object" }, description: "Segments to group by" },
        g: { type: "string", description: "Group by property" },
        limit: { type: "number", description: "Limit results" },
      },
      required: ["e", "start", "end"],
    },
  },
  {
    name: "get_amplitude_user_activity",
    description: "Get user activity and event stream for specific user",
    inputSchema: {
      type: "object" as const,
      properties: {
        user: { type: "string", description: "User ID or Amplitude ID" },
        offset: { type: "number", description: "Pagination offset" },
        limit: { type: "number", description: "Number of events to return" },
      },
      required: ["user"],
    },
  },
  {
    name: "list_amplitude_cohorts",
    description: "List all cohorts in Amplitude project",
    inputSchema: {
      type: "object" as const,
      properties: {},
    },
  },
  {
    name: "get_amplitude_cohort",
    description: "Get details and users for specific Amplitude cohort",
    inputSchema: {
      type: "object" as const,
      properties: {
        cohort_id: { type: "string", description: "Cohort ID" },
      },
      required: ["cohort_id"],
    },
  },
  {
    name: "query_amplitude_charts",
    description: "Query Amplitude saved chart by ID",
    inputSchema: {
      type: "object" as const,
      properties: {
        chart_id: { type: "string", description: "Chart ID" },
        start: { type: "string", description: "Start date (YYYYMMDD)" },
        end: { type: "string", description: "End date (YYYYMMDD)" },
      },
      required: ["chart_id"],
    },
  },
  {
    name: "get_amplitude_taxonomy",
    description: "Get Amplitude event taxonomy (all events and properties)",
    inputSchema: {
      type: "object" as const,
      properties: {
        event_type: { type: "string", description: "Filter by specific event type" },
      },
    },
  },

  // ========== POSTHOG TOOLS ==========
  {
    name: "query_posthog_events",
    description: "Query PostHog events with filters",
    inputSchema: {
      type: "object" as const,
      properties: {
        event: { type: "string", description: "Event name to filter" },
        after: { type: "string", description: "ISO timestamp for events after this time" },
        before: { type: "string", description: "ISO timestamp for events before this time" },
        person_id: { type: "string", description: "Filter by person ID" },
        properties: { type: "object", description: "Event property filters" },
        limit: { type: "number", description: "Limit number of results (default 100)" },
        offset: { type: "number", description: "Pagination offset" },
      },
    },
  },
  {
    name: "get_posthog_person",
    description: "Get PostHog person profile by ID or distinct_id",
    inputSchema: {
      type: "object" as const,
      properties: {
        id: { type: "string", description: "Person ID or distinct_id" },
      },
      required: ["id"],
    },
  },
  {
    name: "list_posthog_feature_flags",
    description: "List all feature flags in PostHog project",
    inputSchema: {
      type: "object" as const,
      properties: {
        limit: { type: "number", description: "Limit results" },
        offset: { type: "number", description: "Pagination offset" },
      },
    },
  },
  {
    name: "query_posthog_insights",
    description: "Query PostHog insights (trends, funnels, retention)",
    inputSchema: {
      type: "object" as const,
      properties: {
        insight: { type: "string", description: "Insight type: TRENDS, FUNNELS, RETENTION, PATHS" },
        events: { type: "array", items: { type: "object" }, description: "Events to analyze" },
        date_from: { type: "string", description: "Start date or relative (e.g., -7d)" },
        date_to: { type: "string", description: "End date or relative" },
        interval: { type: "string", description: "Interval: hour, day, week, month" },
        properties: { type: "object", description: "Filter properties" },
        breakdown: { type: "string", description: "Property to breakdown by" },
        display: { type: "string", description: "Display type for trends" },
      },
      required: ["insight"],
    },
  },
  {
    name: "query_posthog_hogql",
    description: "Run custom PostHog HogQL query",
    inputSchema: {
      type: "object" as const,
      properties: {
        query: { type: "string", description: "HogQL query string (SQL-like syntax)" },
        values: { type: "object", description: "Query parameter values" },
      },
      required: ["query"],
    },
  },
  {
    name: "list_posthog_annotations",
    description: "List annotations in PostHog project",
    inputSchema: {
      type: "object" as const,
      properties: {
        scope: { type: "string", description: "Scope: project or organization" },
      },
    },
  },
];

// ============================================
// TOOL HANDLERS
// ============================================
async function handleMixpanelTool(client: MixpanelClient, name: string, args: any) {
  switch (name) {
    case "query_mixpanel_segmentation": {
      const params = new URLSearchParams();
      params.append("event", args.event);
      params.append("from_date", args.from_date);
      params.append("to_date", args.to_date);
      if (args.type) params.append("type", args.type);
      if (args.unit) params.append("unit", args.unit);
      if (args.where) params.append("where", args.where);
      if (args.on) params.append("on", args.on);
      if (args.limit) params.append("limit", String(args.limit));
      return await client.get(`/segmentation?${params.toString()}`);
    }

    case "query_mixpanel_funnels": {
      const params = new URLSearchParams();
      params.append("events", JSON.stringify(args.events));
      params.append("from_date", args.from_date);
      params.append("to_date", args.to_date);
      if (args.unit) params.append("unit", args.unit);
      if (args.funnel_window_days) params.append("funnel_window_days", String(args.funnel_window_days));
      if (args.on) params.append("on", args.on);
      if (args.where) params.append("where", args.where);
      return await client.get(`/funnels?${params.toString()}`);
    }

    case "query_mixpanel_retention": {
      const params = new URLSearchParams();
      params.append("from_date", args.from_date);
      params.append("to_date", args.to_date);
      if (args.retention_type) params.append("retention_type", args.retention_type);
      if (args.born_event) params.append("born_event", args.born_event);
      if (args.event) params.append("event", args.event);
      if (args.born_where) params.append("born_where", args.born_where);
      if (args.where) params.append("where", args.where);
      if (args.interval) params.append("interval", String(args.interval));
      if (args.interval_count) params.append("interval_count", String(args.interval_count));
      if (args.unit) params.append("unit", args.unit);
      if (args.on) params.append("on", args.on);
      return await client.get(`/retention?${params.toString()}`);
    }

    case "get_mixpanel_user_profile": {
      return await client.get(`/engage?distinct_id=${encodeURIComponent(args.distinct_id)}`);
    }

    case "query_mixpanel_jql": {
      const payload: any = { script: args.script };
      if (args.params) payload.params = args.params;
      return await client.post("/jql", payload);
    }

    case "export_mixpanel_events": {
      const params = new URLSearchParams();
      params.append("from_date", args.from_date);
      params.append("to_date", args.to_date);
      if (args.event) params.append("event", JSON.stringify(args.event));
      if (args.where) params.append("where", args.where);
      if (args.limit) params.append("limit", String(args.limit));
      return await client.get(`/export?${params.toString()}`);
    }

    case "list_mixpanel_events": {
      const params = new URLSearchParams();
      if (args.type) params.append("type", args.type);
      return await client.get(`/events/names?${params.toString()}`);
    }

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

async function handleAmplitudeTool(client: AmplitudeClient, name: string, args: any) {
  switch (name) {
    case "query_amplitude_segmentation": {
      const params = new URLSearchParams();
      params.append("e", JSON.stringify(args.e));
      params.append("start", args.start);
      params.append("end", args.end);
      if (args.m) params.append("m", args.m);
      if (args.i !== undefined) params.append("i", String(args.i));
      if (args.s) params.append("s", JSON.stringify(args.s));
      if (args.g) params.append("g", args.g);
      if (args.limit) params.append("limit", String(args.limit));
      return await client.get(`/events/segmentation?${params.toString()}`);
    }

    case "get_amplitude_user_activity": {
      const params = new URLSearchParams();
      params.append("user", args.user);
      if (args.offset) params.append("offset", String(args.offset));
      if (args.limit) params.append("limit", String(args.limit));
      return await client.get(`/useractivity?${params.toString()}`);
    }

    case "list_amplitude_cohorts": {
      return await client.get("/cohorts");
    }

    case "get_amplitude_cohort": {
      return await client.get(`/cohorts/${args.cohort_id}`);
    }

    case "query_amplitude_charts": {
      const params = new URLSearchParams();
      if (args.start) params.append("start", args.start);
      if (args.end) params.append("end", args.end);
      return await client.get(`/chart/${args.chart_id}/query?${params.toString()}`);
    }

    case "get_amplitude_taxonomy": {
      const params = new URLSearchParams();
      if (args.event_type) params.append("event_type", args.event_type);
      return await client.get(`/taxonomy/event?${params.toString()}`);
    }

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

async function handlePostHogTool(client: PostHogClient, name: string, args: any) {
  switch (name) {
    case "query_posthog_events": {
      const params = new URLSearchParams();
      if (args.event) params.append("event", args.event);
      if (args.after) params.append("after", args.after);
      if (args.before) params.append("before", args.before);
      if (args.person_id) params.append("person_id", args.person_id);
      if (args.properties) params.append("properties", JSON.stringify(args.properties));
      if (args.limit) params.append("limit", String(args.limit));
      if (args.offset) params.append("offset", String(args.offset));
      return await client.get(`/projects/${client['projectId']}/events?${params.toString()}`);
    }

    case "get_posthog_person": {
      return await client.get(`/projects/${client['projectId']}/persons/${encodeURIComponent(args.id)}`);
    }

    case "list_posthog_feature_flags": {
      const params = new URLSearchParams();
      if (args.limit) params.append("limit", String(args.limit));
      if (args.offset) params.append("offset", String(args.offset));
      return await client.get(`/projects/${client['projectId']}/feature_flags?${params.toString()}`);
    }

    case "query_posthog_insights": {
      const payload: any = {
        insight: args.insight,
      };
      if (args.events) payload.events = args.events;
      if (args.date_from) payload.date_from = args.date_from;
      if (args.date_to) payload.date_to = args.date_to;
      if (args.interval) payload.interval = args.interval;
      if (args.properties) payload.properties = args.properties;
      if (args.breakdown) payload.breakdown = args.breakdown;
      if (args.display) payload.display = args.display;
      return await client.post(`/projects/${client['projectId']}/insights`, payload);
    }

    case "query_posthog_hogql": {
      const payload: any = { query: args.query };
      if (args.values) payload.values = args.values;
      return await client.post(`/projects/${client['projectId']}/query`, payload);
    }

    case "list_posthog_annotations": {
      const params = new URLSearchParams();
      if (args.scope) params.append("scope", args.scope);
      return await client.get(`/projects/${client['projectId']}/annotations?${params.toString()}`);
    }

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

// ============================================
// SERVER SETUP
// ============================================
async function main() {
  // Check for required environment variables
  const mixpanelProjectId = process.env.MIXPANEL_PROJECT_ID;
  const mixpanelSecret = process.env.MIXPANEL_SERVICE_ACCOUNT_SECRET;
  const amplitudeApiKey = process.env.AMPLITUDE_API_KEY;
  const amplitudeSecretKey = process.env.AMPLITUDE_SECRET_KEY;
  const posthogApiKey = process.env.POSTHOG_API_KEY;
  const posthogProjectId = process.env.POSTHOG_PROJECT_ID;

  // At least one platform must be configured
  const hasMixpanel = !!(mixpanelProjectId && mixpanelSecret);
  const hasAmplitude = !!(amplitudeApiKey && amplitudeSecretKey);
  const hasPostHog = !!(posthogApiKey && posthogProjectId);

  if (!hasMixpanel && !hasAmplitude && !hasPostHog) {
    console.error("Error: At least one analytics platform must be configured.");
    console.error("Mixpanel requires: MIXPANEL_PROJECT_ID, MIXPANEL_SERVICE_ACCOUNT_SECRET");
    console.error("Amplitude requires: AMPLITUDE_API_KEY, AMPLITUDE_SECRET_KEY");
    console.error("PostHog requires: POSTHOG_API_KEY, POSTHOG_PROJECT_ID");
    process.exit(1);
  }

  // Initialize configured clients
  const mixpanelClient = hasMixpanel ? new MixpanelClient(mixpanelProjectId!, mixpanelSecret!) : null;
  const amplitudeClient = hasAmplitude ? new AmplitudeClient(amplitudeApiKey!, amplitudeSecretKey!) : null;
  const posthogClient = hasPostHog ? new PostHogClient(posthogApiKey!, posthogProjectId!) : null;

  // Filter tools based on available clients
  const availableTools = tools.filter(tool => {
    if (tool.name.includes("mixpanel")) return hasMixpanel;
    if (tool.name.includes("amplitude")) return hasAmplitude;
    if (tool.name.includes("posthog")) return hasPostHog;
    return false;
  });

  console.error(`Product Analytics MCP Server initialized with:`);
  if (hasMixpanel) console.error("  ✓ Mixpanel");
  if (hasAmplitude) console.error("  ✓ Amplitude");
  if (hasPostHog) console.error("  ✓ PostHog");
  console.error(`  Total tools available: ${availableTools.length}`);

  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: any;

      // Route to appropriate client
      if (name.includes("mixpanel")) {
        if (!mixpanelClient) throw new Error("Mixpanel not configured");
        result = await handleMixpanelTool(mixpanelClient, name, args || {});
      } else if (name.includes("amplitude")) {
        if (!amplitudeClient) throw new Error("Amplitude not configured");
        result = await handleAmplitudeTool(amplitudeClient, name, args || {});
      } else if (name.includes("posthog")) {
        if (!posthogClient) throw new Error("PostHog not configured");
        result = await handlePostHogTool(posthogClient, 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);
  console.error(`${MCP_NAME} MCP server running on stdio`);
}

main().catch(console.error);
