const { Server } = require("@modelcontextprotocol/sdk/server/index.js");
const { SSEServerTransport } = require("@modelcontextprotocol/sdk/server/sse.js");
const { 
  ListToolsRequestSchema, 
  CallToolRequestSchema 
} = require("@modelcontextprotocol/sdk/types.js");
const express = require("express");
const axios = require("axios");

const app = express();
const API_BASE = "https://jules.googleapis.com/v1alpha";

app.use(express.json());

/**
 * We use a Map to keep track of multiple active transports.
 * This allows multiple users/clients to connect simultaneously without crashing.
 */
const transports = new Map();

/**
 * Helper for Google API Requests using the user-provided API Key
 */
const julesRequest = async (method, endpoint, apiKey, data = null) => {
  if (!apiKey) {
    throw new Error("Missing Jules API Key. Please provide 'apiKey' in the tool arguments.");
  }
  try {
    const response = await axios({
      method,
      url: `${API_BASE}/${endpoint}`,
      headers: {
        "X-Goog-Api-Key": apiKey,
        "Content-Type": "application/json",
      },
      data,
    });
    return response.data;
  } catch (error) {
    throw new Error(error.response?.data?.error?.message || error.message);
  }
};

/**
 * Creates a fresh server instance with all tool definitions.
 * This prevents the "Already connected to a transport" error.
 */
function createJulesServer() {
  const server = new Server(
    {
      name: "jules-mcp-user-auth",
      version: "1.2.0",
    },
    {
      capabilities: {
        tools: {},
      },
    }
  );

  // 1. Define the List Tools Handler
  server.setRequestHandler(ListToolsRequestSchema, async () => {
    return {
      tools: [
        {
          name: "list_sources",
          description: "List available GitHub repositories connected to Jules.",
          inputSchema: {
            type: "object",
            properties: {
              apiKey: { type: "string", description: "Your Jules API Key" }
            },
            required: ["apiKey"]
          },
        },
        {
          name: "create_session",
          description: "Start a new coding task with Jules.",
          inputSchema: {
            type: "object",
            properties: {
              apiKey: { type: "string", description: "Your Jules API Key" },
              prompt: { type: "string", description: "What do you want Jules to do?" },
              source: { type: "string", description: "Source name (e.g., sources/github/user/repo)" },
              title: { type: "string", description: "Title for the session" },
              branch: { type: "string", description: "Starting branch (default: main)", default: "main" },
              autoPR: { type: "boolean", description: "Automatically create a PR?", default: false }
            },
            required: ["apiKey", "prompt", "source", "title"],
          },
        },
        {
          name: "approve_plan",
          description: "Approve a plan generated by Jules for a specific session.",
          inputSchema: {
            type: "object",
            properties: {
              apiKey: { type: "string", description: "Your Jules API Key" },
              sessionId: { type: "string", description: "The ID of the session" }
            },
            required: ["apiKey", "sessionId"],
          },
        },
        {
          name: "get_activities",
          description: "List activities and progress for a session.",
          inputSchema: {
            type: "object",
            properties: {
              apiKey: { type: "string", description: "Your Jules API Key" },
              sessionId: { type: "string", description: "The ID of the session" }
            },
            required: ["apiKey", "sessionId"],
          },
        },
        {
          name: "send_message",
          description: "Send a follow-up message to Jules during a session.",
          inputSchema: {
            type: "object",
            properties: {
              apiKey: { type: "string", description: "Your Jules API Key" },
              sessionId: { type: "string", description: "The ID of the session" },
              prompt: { type: "string", description: "Your message to the agent" }
            },
            required: ["apiKey", "sessionId", "prompt"],
          },
        }
      ],
    };
  });

  // 2. Define the Call Tool Handler
  server.setRequestHandler(CallToolRequestSchema, async (request) => {
    const { name, arguments: args } = request.params;
    const userApiKey = args.apiKey;

    try {
      switch (name) {
        case "list_sources":
          const sources = await julesRequest("GET", "sources", userApiKey);
          return { content: [{ type: "text", text: JSON.stringify(sources, null, 2) }] };

        case "create_session":
          const sessionData = {
            prompt: args.prompt,
            title: args.title,
            sourceContext: {
              source: args.source,
              githubRepoContext: { startingBranch: args.branch || "main" }
            },
            automationMode: args.autoPR ? "AUTO_CREATE_PR" : undefined
          };
          const newSession = await julesRequest("POST", "sessions", userApiKey, sessionData);
          return { content: [{ type: "text", text: JSON.stringify(newSession, null, 2) }] };

        case "approve_plan":
          const approval = await julesRequest("POST", `sessions/${args.sessionId}:approvePlan`, userApiKey);
          return { content: [{ type: "text", text: JSON.stringify(approval, null, 2) }] };

        case "get_activities":
          const activities = await julesRequest("GET", `sessions/${args.sessionId}/activities?pageSize=50`, userApiKey);
          return { content: [{ type: "text", text: JSON.stringify(activities, null, 2) }] };

        case "send_message":
          const message = await julesRequest("POST", `sessions/${args.sessionId}:sendMessage`, userApiKey, { prompt: args.prompt });
          return { content: [{ type: "text", text: JSON.stringify(message, null, 2) }] };

        default:
          throw new Error(`Unknown tool: ${name}`);
      }
    } catch (err) {
      return { isError: true, content: [{ type: "text", text: err.message }] };
    }
  });

  return server;
}

app.get("/sse", async (req, res) => {
  console.log("New SSE connection attempt...");
  
  // Create a new server and transport for EVERY connection
  const server = createJulesServer();
  const transport = new SSEServerTransport("/messages", res);
  
  // Store the transport so the POST route can find it by sessionId
  transports.set(transport.sessionId, transport);
  
  // Cleanup when connection closes
  res.on("close", () => {
    console.log(`Closing transport for session: ${transport.sessionId}`);
    transports.delete(transport.sessionId);
  });

  await server.connect(transport);
});

app.post("/messages", async (req, res) => {
  const sessionId = req.query.sessionId;
  const transport = transports.get(sessionId);

  if (!transport) {
    return res.status(404).send("Session not found");
  }

  await transport.handlePostMessage(req, res);
});

app.get("/", (req, res) => {
  const host = req.get('host');
  const protocol = req.get('x-forwarded-proto') || req.protocol;
  const fullUrl = `${protocol}://${host}/sse`;

  res.send(`
    <html>
      <head>
        <title>Jules Dynamic MCP Server</title>
        <style>
          body { font-family: -apple-system, sans-serif; line-height: 1.6; padding: 40px; color: #333; max-width: 800px; margin: auto; background: #f9f9f9; }
          .card { background: #fff; padding: 30px; border-radius: 12px; box-shadow: 0 4px 6px rgba(0,0,0,0.05); }
          code { background: #f0f0f0; padding: 4px 8px; border-radius: 4px; font-weight: bold; color: #d63384; }
          pre { background: #1e1e1e; color: #569cd6; padding: 20px; border-radius: 8px; overflow-x: auto; border: 1px solid #333; }
          .url { color: #4caf50; font-weight: bold; }
          .step { margin-bottom: 20px; padding-left: 20px; border-left: 4px solid #007bff; }
        </style>
      </head>
      <body>
        <div class="card">
          <h1>🚀 Jules API MCP Server</h1>
          <p>Your dynamic MCP server for Google Jules is ready!</p>
          
          <div class="step">
            <h3>1. Copy SSE URL</h3>
            <p>Add this URL to your MCP client (Claude Desktop or Cursor):</p>
            <pre class="url">${fullUrl}</pre>
          </div>

          <div class="step">
            <h3>2. Provide API Key</h3>
            <p>Every tool call requires your <code>apiKey</code>. You can obtain it from the <strong>Jules Web App Settings</strong>.</p>
          </div>

          <div class="step">
            <h3>3. Usage Tip</h3>
            <p>Tell the AI: <i>"Using the Jules MCP, list my sources. Here is my API key: [YOUR_KEY]"</i></p>
          </div>
        </div>
      </body>
    </html>
  `);
});

const PORT = process.env.PORT || 7860;
app.listen(PORT, () => {
  console.log(`Server listening on port ${PORT}`);
});