import Fastify, { type FastifyInstance } from "fastify";
import { ToolDef } from "./types";
import McpSession from "./mcpConnection";
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types";
import { pingTool, testAddTool, testRemoveTool, listToolsTool, invokeToolTool } from "./mcpTools";

export class Server {
    private readonly app: FastifyInstance;
    private readonly port: number;
    private readonly sessions: McpSession[] = [];
    public readonly tools: Record<string, ToolDef> = {};

    constructor(port: number = Number(process.env.PORT ?? 3000)) {
        this.port = port;
        this.app = Fastify({
            trustProxy: true,
            logger: {
                serializers: {
                    req(request) {
                        return {
                            method: request.method,
                            url: request.url,
                            host: request.headers.host,
                            remoteAddress: request.ip,
                            headers: request.headers, // <-- full header set
                        };
                    },
                },
            },

            // logger: {
            //     level: process.env.LOG_LEVEL ?? "info",
            //     ...(process.env.NODE_ENV !== "production"
            //         ? {
            //             transport: {
            //                 target: "pino-pretty",
            //                 options: { colorize: true, translateTime: "HH:MM:ss.l", ignore: "pid,hostname" },
            //             },
            //         }
            //         : {}),
            // },
        });
        this.registerLogging();
        this.registerRoutes();
    }

    private registerLogging(): void {
        // onResponse is skipped for hijacked replies (the /mcp routes hijack to
        // hand the raw response to the MCP SDK), so listen on the raw response's
        // "finish" event instead — that fires regardless of hijack status.
        this.app.addHook("onRequest", async (request, reply) => {
            const start = process.hrtime.bigint();
            request.log.info({ method: request.method, url: request.url, ip: request.ip }, "request received");
            reply.raw.on("finish", () => {
                const ms = Number(process.hrtime.bigint() - start) / 1e6;
                request.log.info(
                    {
                        method: request.method,
                        url: request.url,
                        statusCode: reply.raw.statusCode,
                        ms: Number(ms.toFixed(1)),
                    },
                    "request completed",
                );
            });
        });
    }

    public getTool(name: string): ToolDef {
        return this.tools[name];
    }

    public registerTools(...tools: ToolDef[]): void {

        for (const tool of tools) {
            if (!this.getTool(tool.title))
                this.tools[tool.title] = tool;
        }

        this.emitToolListChange();
    }

    public unregisterTools(...names: string[]): void {
        for (const name of names)
            delete (this.tools[name]);

        this.emitToolListChange();
    }

    private emitToolListChange() {
        for (const session of Object.values(this.sessions))
            session.handleToolListChange(Object.values(this.tools));
    }

    public getSession(sessionId: string): McpSession | undefined {
        return this.sessions.find((s) => s.sessionId === sessionId);
    }

    // Stateful Streamable HTTP transport, keyed by session id. Fastify's
    // request/reply are hijacked so the MCP SDK can write directly to the
    // underlying Node req/res it expects.
    private registerRoutes(): void {
        this.app.post("/mcp", async (request, reply) => {

            let session: McpSession | undefined;

            const sessionId = request.headers["mcp-session-id"] as string | undefined;
            if (!sessionId) {

                console.log("[mcp] typeof body:", typeof request.body,
                    "isInit:", isInitializeRequest(request.body),
                    "session:", request.headers["mcp-session-id"]);

                if (!isInitializeRequest(request.body))
                    throw new Error('Only initialize request is allowed to have no mcp-session-id header');
                session = new McpSession();
                session.handleToolListChange(Object.values(this.tools), false);
                await session.init();
                this.sessions.push(session);
            }
            else {
                session = this.getSession(sessionId);
                if (!session)
                    return reply.status(404).send("Not found");
            }

            session.handleRequest(request, reply);
        });

        // GET = open the SSE stream the server pushes notifications through
        this.app.get("/mcp", async (request, reply) => {
            const sessionId = request.headers["mcp-session-id"] as string | undefined;
            if (!sessionId)
                throw new Error('mcp-session-id header is missing');

            const session = this.getSession(sessionId);
            if (!session)
                return reply.status(404).send("Not found");

            await session.handleRequest(request, reply);
        });

        // DELETE = client terminates the session
        this.app.delete("/mcp", async (request, reply) => {
            const sessionId = request.headers["mcp-session-id"] as string | undefined;
            if (!sessionId)
                throw new Error('mcp-session-id header is missing');

            const session = this.getSession(sessionId);
            if (!session)
                return reply.status(404).send("Not found");

            await session.handleRequest(request, reply);
        });

        this.app.get("/healthz", async () => ({ ok: true, sessions: Object.keys(this.sessions).length }));
    }

    async start(): Promise<void> {
        await this.app.listen({ port: this.port, host: "0.0.0.0" });
        this.app.log.info(`MCP ping server listening on http://0.0.0.0:${this.port}/mcp`);
    }
}

export const server = new Server();
server.registerTools(pingTool, listToolsTool, invokeToolTool, testAddTool, testRemoveTool);