import Fastify, { type FastifyInstance } from "fastify";
import { TextResult, ToolDef } from "./types";
import McpSession from "./mcpConnection";
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
import * as mcpTools from "./mcpTools";
import { logged } from "./utils/logged";
import cors from '@fastify/cors';

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

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

    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)) {
                // console.log(`Registering new tool ${tool.title}`);
                const toolLogged = logged(tool);
                this.tools[tool.title] = {
                    ...toolLogged,
                    handler: async (args) => {
                        try {
                            const res = await toolLogged.handler(args);
                            if (typeof res === 'string')
                                return { content: [{ type: "text", text: res, }] };
                            else
                                return res;
                        }
                        catch (err) {
                            console.error(err);
                            return {
                                isError: true,
                                content: [{ type: "text", text: (err as Error).message ?? String(err) }],
                            };
                        }
                    },
                };
            }
        }

        this.emitToolListChange();
    }

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

        this.emitToolListChange();
    }

    private emitToolListChange() {
        const toolList = Object.values(this.tools);
        console.log(`Advertising updated tool list: ${toolList.map((tool) => tool.title)}`);

        for (const session of Object.values(this.sessions))
            session.handleToolListChange(toolList);
    }

    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) => {
            reply.raw.setHeader('Access-Control-Allow-Origin', request.headers.origin ?? '*');
            reply.raw.setHeader('Access-Control-Allow-Credentials', 'true');
            reply.raw.setHeader('Vary', 'Origin');

            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.register(cors, {
            origin: true,
            credentials: true,
        });

        await this.registerLogging();
        await this.registerRoutes();

        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(...Object.values(mcpTools));