import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
import express from 'express';
import crypto from 'crypto';
import {
    CallToolRequestSchema,
    ListToolsRequestSchema,
    Tool,
} from "@modelcontextprotocol/sdk/types.js";
import axios from "axios";
import https from "https";

const UMAMI_BASE_URL = process.env.UMAMI_BASE_URL || "https://umami.infra.mintel.me";
const UMAMI_USERNAME = process.env.UMAMI_USERNAME;
const UMAMI_PASSWORD = process.env.UMAMI_PASSWORD;
const UMAMI_API_KEY = process.env.UMAMI_API_KEY; // optional if using API key auth

const httpsAgent = new https.Agent({
    rejectUnauthorized: false,
});

if (!UMAMI_USERNAME && !UMAMI_API_KEY) {
    console.error("Warning: Neither UMAMI_USERNAME/PASSWORD nor UMAMI_API_KEY is set.");
}

// Token cache to avoid logging in on every request
let cachedToken: string | null = null;

async function getAuthHeaders(): Promise<Record<string, string>> {
    if (UMAMI_API_KEY) {
        return { 'x-umami-api-key': UMAMI_API_KEY };
    }
    if (!cachedToken) {
        const res = await axios.post(`${UMAMI_BASE_URL}/api/auth/login`, {
            username: UMAMI_USERNAME,
            password: UMAMI_PASSWORD,
        }, { httpsAgent });
        cachedToken = res.data.token;
    }
    return { Authorization: `Bearer ${cachedToken}` };
}


// --- Tool Definitions ---
const LIST_WEBSITES_TOOL: Tool = {
    name: "umami_list_websites",
    description: "List all websites tracked in Umami",
    inputSchema: { type: "object", properties: {} },
};

const GET_WEBSITE_STATS_TOOL: Tool = {
    name: "umami_get_website_stats",
    description: "Get summary statistics for a website for a time range",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
            start_at: { type: "number", description: "Start timestamp in ms (e.g., Date.now() - 7 days)" },
            end_at: { type: "number", description: "End timestamp in ms (default: now)" },
        },
        required: ["website_id", "start_at"],
    },
};

const GET_PAGE_VIEWS_TOOL: Tool = {
    name: "umami_get_pageviews",
    description: "Get pageview/session time series for a website",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
            start_at: { type: "number", description: "Start timestamp in ms" },
            end_at: { type: "number", description: "End timestamp in ms (default: now)" },
            unit: { type: "string", description: "Time unit: 'hour', 'day', 'month' (default: day)" },
            timezone: { type: "string", description: "Timezone (default: Europe/Berlin)" },
        },
        required: ["website_id", "start_at"],
    },
};

const GET_TOP_PAGES_TOOL: Tool = {
    name: "umami_get_top_pages",
    description: "Get the most visited pages/URLs for a website",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
            start_at: { type: "number", description: "Start timestamp in ms" },
            end_at: { type: "number", description: "End timestamp in ms" },
            limit: { type: "number", description: "Number of results (default: 20)" },
        },
        required: ["website_id", "start_at"],
    },
};

const GET_TOP_REFERRERS_TOOL: Tool = {
    name: "umami_get_top_referrers",
    description: "Get the top traffic referrers for a website",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
            start_at: { type: "number", description: "Start timestamp in ms" },
            end_at: { type: "number", description: "End timestamp in ms" },
            limit: { type: "number", description: "Number of results (default: 10)" },
        },
        required: ["website_id", "start_at"],
    },
};

const GET_COUNTRY_STATS_TOOL: Tool = {
    name: "umami_get_country_stats",
    description: "Get visitor breakdown by country",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
            start_at: { type: "number", description: "Start timestamp in ms" },
            end_at: { type: "number", description: "End timestamp in ms" },
        },
        required: ["website_id", "start_at"],
    },
};

const GET_ACTIVE_VISITORS_TOOL: Tool = {
    name: "umami_get_active_visitors",
    description: "Get the number of visitors currently active on a website (last 5 minutes)",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
        },
        required: ["website_id"],
    },
};

// --- Server Setup ---
const server = new Server(
    { name: "umami-mcp", version: "1.0.0" },
    { capabilities: { tools: {} } }
);

server.setRequestHandler(ListToolsRequestSchema, async () => ({
    tools: [
        LIST_WEBSITES_TOOL,
        GET_WEBSITE_STATS_TOOL,
        GET_PAGE_VIEWS_TOOL,
        GET_TOP_PAGES_TOOL,
        GET_TOP_REFERRERS_TOOL,
        GET_COUNTRY_STATS_TOOL,
        GET_ACTIVE_VISITORS_TOOL,
    ],
}));

server.setRequestHandler(CallToolRequestSchema, async (request) => {
    const headers = await getAuthHeaders();
    const api = axios.create({ baseURL: `${UMAMI_BASE_URL}/api`, headers, httpsAgent });

    const now = Date.now();

    if (request.params.name === "umami_list_websites") {
        try {
            const res = await api.get('/websites');
            const sites = (res.data.data || res.data || []).map((s: any) => ({
                id: s.id, name: s.name, domain: s.domain
            }));
            return { content: [{ type: "text", text: JSON.stringify(sites, null, 2) }] };
        } catch (e: any) {
            return { isError: true, content: [{ type: "text", text: `Error: ${e.message}` }] };
        }
    }

    if (request.params.name === "umami_get_website_stats") {
        const { website_id, start_at, end_at = now } = request.params.arguments as any;
        try {
            const res = await api.get(`/websites/${website_id}/stats`, { params: { startAt: start_at, endAt: end_at } });
            return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
        } catch (e: any) {
            return { isError: true, content: [{ type: "text", text: `Error: ${e.message}` }] };
        }
    }

    if (request.params.name === "umami_get_pageviews") {
        const { website_id, start_at, end_at = now, unit = 'day', timezone = 'Europe/Berlin' } = request.params.arguments as any;
        try {
            const res = await api.get(`/websites/${website_id}/pageviews`, {
                params: { startAt: start_at, endAt: end_at, unit, timezone }
            });
            return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
        } catch (e: any) {
            return { isError: true, content: [{ type: "text", text: `Error: ${e.message}` }] };
        }
    }

    if (request.params.name === "umami_get_top_pages") {
        const { website_id, start_at, end_at = now, limit = 20 } = request.params.arguments as any;
        try {
            const res = await api.get(`/websites/${website_id}/metrics`, {
                params: { startAt: start_at, endAt: end_at, type: 'url', limit }
            });
            return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
        } catch (e: any) {
            return { isError: true, content: [{ type: "text", text: `Error: ${e.message}` }] };
        }
    }

    if (request.params.name === "umami_get_top_referrers") {
        const { website_id, start_at, end_at = now, limit = 10 } = request.params.arguments as any;
        try {
            const res = await api.get(`/websites/${website_id}/metrics`, {
                params: { startAt: start_at, endAt: end_at, type: 'referrer', limit }
            });
            return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
        } catch (e: any) {
            return { isError: true, content: [{ type: "text", text: `Error: ${e.message}` }] };
        }
    }

    if (request.params.name === "umami_get_country_stats") {
        const { website_id, start_at, end_at = now } = request.params.arguments as any;
        try {
            const res = await api.get(`/websites/${website_id}/metrics`, {
                params: { startAt: start_at, endAt: end_at, type: 'country' }
            });
            return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
        } catch (e: any) {
            return { isError: true, content: [{ type: "text", text: `Error: ${e.message}` }] };
        }
    }

    if (request.params.name === "umami_get_active_visitors") {
        const { website_id } = request.params.arguments as any;
        try {
            const res = await api.get(`/websites/${website_id}/active`);
            return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
        } catch (e: any) {
            return { isError: true, content: [{ type: "text", text: `Error: ${e.message}` }] };
        }
    }

    throw new Error(`Unknown tool: ${request.params.name}`);
});

// --- Express / SSE Server ---
async function run() {
    const isStdio = process.argv.includes('--stdio');

    if (isStdio) {
        const { StdioServerTransport } = await import('@modelcontextprotocol/sdk/server/stdio.js');
        const transport = new StdioServerTransport();
        await server.connect(transport);
        console.error('Umami MCP server is running on stdio');
    } else {
        const app = express();
        const transports = new Map<string, SSEServerTransport>();

        app.use((req, _res, next) => {
            console.error(`${req.method} ${req.url}`);
            next();
        });

        app.get('/sse', async (req, res) => {
            const sessionId = crypto.randomUUID();
            console.error(`New SSE connection: ${sessionId}`);
            const transport = new SSEServerTransport(`/message/${sessionId}`, res);
            transports.set(sessionId, transport);

            req.on('close', () => {
                console.error(`SSE connection closed: ${sessionId}`);
                transports.delete(sessionId);
            });

            await server.connect(transport);
        });

        app.post('/message/:sessionId', async (req, res) => {
            const { sessionId } = req.params;
            const transport = transports.get(sessionId as string);

            if (!transport) {
                console.error(`No transport found for session: ${sessionId}`);
                res.status(400).send('No active SSE connection for this session');
                return;
            }
            await transport.handlePostMessage(req, res);
        });

        const PORT = process.env.UMAMI_MCP_PORT || 3003;
        app.listen(PORT, () => {
            console.error(`Umami MCP server running on http://localhost:${PORT}/sse`);
        });
    }
}

run().catch((err) => {
    console.error("Fatal error:", err);
    process.exit(1);
});