import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
// import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js"; // Corrected class name casing
import { z } from "zod";
import { BoardState, Team } from "../shared/types";
import { CodenamesGameHandler } from "../browser/handler";
import { Browser, chromium } from 'playwright'; // Added Browser, chromium
// import * as http from 'http'; // No longer using raw http server
// import * as url from 'url'; // No longer using raw url parsing
import express, { Request, Response } from 'express'; // Import express
// import { JSONRPCMessageSchema } from "@modelcontextprotocol/sdk/types.js"; // Import schema if needed for debug

// --- Global State ---
// let browserHandler: CodenamesGameHandler | null = null; // Replaced by handlers object
let mcpServer: McpServer | null = null;
const transports: { [sessionId: string]: SSEServerTransport } = {};
let handlers: { [key in Team.Red | Team.Blue]?: CodenamesGameHandler } = {}; // Store handlers for Red and Blue
let browser: Browser | null = null; // Store the main browser instance

async function initializeBrowserAndMcp() {
    try {
        console.log("Launching browser...");
        browser = await chromium.launch();
        console.log("Creating browser contexts...");
        const redContext = await browser.newContext();
        const blueContext = await browser.newContext();

        console.log("Initializing Red Codenames handler...");
        handlers[Team.Red] = await CodenamesGameHandler.createInContext(redContext);
        await handlers[Team.Red]!.joinRoom(); // Use non-null assertion
        await handlers[Team.Red]!.joinAsSpymaster(Team.Red);
        console.log("Red handler initialized and joined room/role.");

        console.log("Initializing Blue Codenames handler...");
        handlers[Team.Blue] = await CodenamesGameHandler.createInContext(blueContext);
        await handlers[Team.Blue]!.joinRoom(); // Use non-null assertion
        await handlers[Team.Blue]!.joinAsSpymaster(Team.Blue);
        console.log("Blue handler initialized and joined room/role.");

        console.log("All handlers initialized.");

        // Now define MCP server and tools *after* async setup is done
        mcpServer = new McpServer(
            { name: "codenames-mcp-server", version: "1.0.0" },
            { capabilities: {} }
        );

        mcpServer.tool(
            "getBoardState",
            { team: z.enum([Team.Red, Team.Blue]) },
            async ({ team }: { team: Team.Red | Team.Blue }) => {
                const handler = handlers[team];
                if (!handler) throw new Error(`Browser handler for team (${team}) not initialized`);
                console.log(`[getBoardState] Called by ${team} team.`);
                const board: BoardState = await handler.getBoardState();
                return {
                    content: [
                        { type: "text", text: JSON.stringify(board) }
                    ],
                };
            }
        );

        mcpServer.tool(
            "submitClue",
            { team: z.enum([Team.Red, Team.Blue]), word: z.string(), count: z.number() },
            async ({ team, word, count }: { team: Team.Red | Team.Blue; word: string; count: number }) => {
                const handler = handlers[team];
                if (!handler) throw new Error(`Browser handler for team (${team}) not initialized`);

                console.log(`[submitClue] Called by ${team} team. Submitting clue: ${word} (${count})`);
                await handler.submitClue({ word, count });

                // No turn toggling here!
                return {
                    content: [
                        { type: "text", text: `Clue submitted by ${team}: ${word} (${count}).` }
                    ],
                    success: true
                };
            }
        );

    } catch (initError) {
        console.error("Error during async browser/MCP initialization:", initError);
        process.exit(1); // Exit if essential setup fails
    }
}

// Main execution flow
async function main() {
    try {
        // Perform async initialization first
        await initializeBrowserAndMcp();

        // --- Express App Setup ---
        const app = express();
        const port = 4000;

        // --- SSE Connection Route (GET /sse) ---
        app.get("/sse", async (req: Request, res: Response): Promise<void> => {
            const logPrefixBase = `[SSE Connect ${req.ip}]`;
            console.log(`${logPrefixBase} SSE connection request received`);

            // Check if MCP server is ready
            if (!mcpServer) {
                console.error(`${logPrefixBase} Error: MCP Server not ready for SSE connection`);
                res.status(503).send("Server not ready"); // Send the response
                return; // Then return void
            }

            // Check for existing session ID (optional, for potential reconnection logic)
            const existingSessionId = req.query.sessionId as string;
            if (existingSessionId && transports[existingSessionId]) {
                console.log(`${logPrefixBase} Reconnecting session: ${existingSessionId}`);
                // Potentially handle reconnection logic here if needed
            }

            // Generate a unique session ID
            const sessionId = crypto.randomUUID();
            const logPrefix = `[${sessionId}]`;
            console.log(`${logPrefix} Generated new session ID`);

            try {
                // Create a new SSE transport
                const transport = new SSEServerTransport(`/messages`, res);
                // Use the transport's own sessionId as the key for our map
                const transportSessionId = transport.sessionId;
                transports[transportSessionId] = transport;
                console.log(`${logPrefix} SSE Transport created. Storing under session ID: ${transportSessionId}`); // Log the key being used

                // Handle client disconnect
                req.on('close', () => {
                    console.log(`[${transportSessionId}] Client disconnected`); // Log with the correct ID
                    transport.close(); // Close the transport on the server side
                    delete transports[transportSessionId]; // Clean up the transport reference using the correct key
                    console.log(`[${transportSessionId}] Transport closed and cleaned up.`); // Log with the correct ID
                });

                // Establish connection with MCP server
                console.log(`[${transportSessionId}] Connecting transport to MCP server...`); // Log with the correct ID
                await mcpServer.connect(transport);
                console.log(`[${transportSessionId}] Transport connected to MCP server.`); // Log with the correct ID

            } catch (error) {
                // Cannot access transportSessionId here if transport creation failed
                console.error(`${logPrefix} Error setting up SSE connection:`, error);
                // We don't know the transportSessionId if creation failed, so cannot clean up transports[transportSessionId]
                // If creation succeeded but connect failed, the disconnect handler should eventually clean it up.
                if (!res.headersSent) {
                    res.status(500).send('Failed to establish SSE connection');
                }
            }
        });

        // --- Message Posting Route (POST /messages) ---
        app.post(
            "/messages",
            express.json({ limit: "5mb" }),
            async (req: Request, res: Response): Promise<void> => {
                const sessionId = req.query.sessionId as string;
                if (!sessionId) {
                    console.error("POST /messages error: Missing sessionId");
                    res.status(400).send("Missing sessionId"); return;
                }
                const logPrefix = `[${sessionId}]`;
                console.log(`${logPrefix} POST /messages received`);
                const transport = transports[sessionId];
                const messageBody = req.body;
                if (!transport) {
                    console.error(`${logPrefix} POST /messages error: No transport`);
                    res.status(404).send("No session"); return;
                }
                if (!messageBody) {
                    console.error(`${logPrefix} POST /messages error: No body`);
                    res.status(400).send("No body"); return;
                }
                console.log(`${logPrefix} --> Received body:`, JSON.stringify(messageBody, null, 2));
                try {
                    console.log(`${logPrefix} DEBUG: Calling transport.handlePostMessage...`);
                    await transport.handlePostMessage(req, res, messageBody);
                    console.log(`${logPrefix} DEBUG: transport.handlePostMessage completed.`);
                } catch (error) {
                    console.error(`${logPrefix} ERROR in POST /messages handler:`, error);
                    if (!res.headersSent) {
                        res.status(500).send("Server Error");
                    }
                }
            }
        );

        // --- Start Express Server ---
        app.listen(port, '0.0.0.0', () => {
            console.log(`Express server listening on port ${port}`);
            console.log(`MCP SSE endpoint: http://localhost:${port}/sse`);
            console.log(`MCP Messages endpoint: POST http://localhost:${port}/messages`);
        });

        // Add an indefinite wait to keep the main function/script alive
        console.log("Server setup complete. Entering indefinite wait...");
        // Add cleanup logic on process exit
        process.on('SIGINT', async () => {
            console.log("\nReceived SIGINT. Cleaning up resources...");
            if (handlers[Team.Red]) {
                await handlers[Team.Red]?._cleanupResources();
            }
            if (handlers[Team.Blue]) {
                await handlers[Team.Blue]?._cleanupResources();
            }
            if (browser) {
                await browser.close();
                console.log("Browser closed.");
            }
            process.exit(0);
        });

        await new Promise(() => {}); // This promise never resolves
    } catch (initError) {
        console.error("Error during server initialization:", initError);
        process.exit(1);
    }
}

// Call the async main function
main(); // No catch needed here as errors are handled inside main or by Node