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 SERPBEAR_BASE_URL = process.env.SERPBEAR_BASE_URL || "https://serpbear.infra.mintel.me";
const SERPBEAR_API_KEY = process.env.SERPBEAR_API_KEY;

if (!SERPBEAR_API_KEY) {
    console.error("Warning: SERPBEAR_API_KEY is not set. API calls will fail.");
}

const serpbearClient = axios.create({
    baseURL: `${SERPBEAR_BASE_URL}/api`,
    headers: { 'x-api-key': SERPBEAR_API_KEY || '' },
    httpsAgent: new https.Agent({
        rejectUnauthorized: false,
    }),
    timeout: 10000
});

// Logging middleware for SerpBear client
serpbearClient.interceptors.request.use((config) => {
    console.error(`SerpBear Request: ${config.method?.toUpperCase()} ${config.url}`);
    return config;
}, (error) => {
    console.error(`SerpBear Request Error: ${error.message}`);
    return Promise.reject(error);
});

serpbearClient.interceptors.response.use((response) => {
    return response;
}, (error) => {
    const msg = error.response?.data?.error || error.response?.data?.message || error.message;
    console.error(`SerpBear API Error (${error.config?.url}): ${msg}`);
    return Promise.reject(error);
});


// --- Tool Definitions ---
const LIST_DOMAINS_TOOL: Tool = {
    name: "serpbear_list_domains",
    description: "List all domains/projects tracked in SerpBear",
    inputSchema: { type: "object", properties: {} },
};

const GET_KEYWORDS_TOOL: Tool = {
    name: "serpbear_get_keywords",
    description: "Get all tracked keywords for a domain, with their current ranking positions",
    inputSchema: {
        type: "object",
        properties: {
            domain_id: { type: "string", description: "Domain ID from serpbear_list_domains" },
        },
        required: ["domain_id"],
    },
};

const ADD_KEYWORDS_TOOL: Tool = {
    name: "serpbear_add_keywords",
    description: "Add new keywords to track for a domain",
    inputSchema: {
        type: "object",
        properties: {
            domain_id: { type: "string", description: "Domain ID" },
            keywords: {
                type: "array",
                items: { type: "string" },
                description: "List of keywords to add (e.g., ['Webentwickler Frankfurt', 'Next.js Agentur'])"
            },
            country: { type: "string", description: "Country code for SERP tracking (e.g., 'de', 'us'). Default: 'de'" },
            device: { type: "string", description: "Device type: 'desktop' or 'mobile'. Default: 'desktop'" },
        },
        required: ["domain_id", "keywords"],
    },
};

const DELETE_KEYWORDS_TOOL: Tool = {
    name: "serpbear_delete_keywords",
    description: "Remove keywords from tracking",
    inputSchema: {
        type: "object",
        properties: {
            keyword_ids: {
                type: "array",
                items: { type: "number" },
                description: "Array of keyword IDs to delete"
            },
        },
        required: ["keyword_ids"],
    },
};

const REFRESH_KEYWORDS_TOOL: Tool = {
    name: "serpbear_refresh_keywords",
    description: "Trigger an immediate SERP position refresh for specific keywords",
    inputSchema: {
        type: "object",
        properties: {
            keyword_ids: {
                type: "array",
                items: { type: "number" },
                description: "List of keyword IDs to refresh"
            },
        },
        required: ["keyword_ids"],
    },
};

const GET_KEYWORD_HISTORY_TOOL: Tool = {
    name: "serpbear_get_keyword_history",
    description: "Get the ranking history for a specific keyword over time",
    inputSchema: {
        type: "object",
        properties: {
            keyword_id: { type: "number", description: "Keyword ID from serpbear_get_keywords" },
        },
        required: ["keyword_id"],
    },
};

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

server.setRequestHandler(ListToolsRequestSchema, async () => ({
    tools: [
        LIST_DOMAINS_TOOL,
        GET_KEYWORDS_TOOL,
        ADD_KEYWORDS_TOOL,
        DELETE_KEYWORDS_TOOL,
        REFRESH_KEYWORDS_TOOL,
        GET_KEYWORD_HISTORY_TOOL,
    ],
}));

server.setRequestHandler(CallToolRequestSchema, async (request) => {
    try {
        if (request.params.name === "serpbear_list_domains") {
            const res = await serpbearClient.get('/domains');
            const domains = (res.data.domains || []).map((d: any) => ({
                id: d.id, domain: d.domain, keywords: d.keywordCount
            }));
            return { content: [{ type: "text", text: JSON.stringify(domains, null, 2) }] };
        }

        if (request.params.name === "serpbear_get_keywords") {
            const { domain_id } = request.params.arguments as any;
            const res = await serpbearClient.get('/keywords', { params: { domain: domain_id } });
            const keywords = (res.data.keywords || []).map((k: any) => ({
                id: k.id,
                keyword: k.keyword,
                position: k.position,
                lastUpdated: k.lastUpdated,
                country: k.country,
                device: k.device,
                change: k.position_change ?? null,
            }));
            return { content: [{ type: "text", text: JSON.stringify(keywords, null, 2) }] };
        }

        if (request.params.name === "serpbear_add_keywords") {
            const { domain_id, keywords, country = 'de', device = 'desktop' } = request.params.arguments as any;
            const res = await serpbearClient.post('/keywords', {
                domain: domain_id,
                keywords: keywords.map((kw: string) => ({ keyword: kw, country, device })),
            });
            return { content: [{ type: "text", text: `Added ${keywords.length} keywords. Result: ${JSON.stringify(res.data)}` }] };
        }

        if (request.params.name === "serpbear_delete_keywords") {
            const { keyword_ids } = request.params.arguments as any;
            await serpbearClient.delete('/keywords', { data: { ids: keyword_ids } });
            return { content: [{ type: "text", text: `Deleted ${keyword_ids.length} keywords.` }] };
        }

        if (request.params.name === "serpbear_refresh_keywords") {
            const { keyword_ids } = request.params.arguments as any;
            await serpbearClient.post('/keywords/refresh', { ids: keyword_ids });
            return { content: [{ type: "text", text: `Triggered refresh for ${keyword_ids.length} keywords.` }] };
        }

        if (request.params.name === "serpbear_get_keyword_history") {
            const { keyword_id } = request.params.arguments as any;
            const res = await serpbearClient.get(`/keywords/${keyword_id}/history`);
            return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
        }

        throw new Error(`Unknown tool: ${request.params.name}`);
    } catch (e: any) {
        const msg = e.response?.data?.error || e.response?.data?.message || e.message;
        return { isError: true, content: [{ type: "text", text: `SerpBear Service Error: ${msg}` }] };
    }
});

// --- 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('SerpBear MCP server is running on stdio');
    } else {
        const app = express();
        const transports = new Map<string, SSEServerTransport>();

        // Manual CORS middleware
        app.use((req, res, next) => {
            res.setHeader('Access-Control-Allow-Origin', '*');
            res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
            res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization, x-mcp-protocol-version');
            if (req.method === 'OPTIONS') {
                res.status(204).end();
                return;
            }
            next();
        });

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

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

            const heartbeatInterval = setInterval(() => {
                res.write(": heartbeat\n\n");
            }, 15000);

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

            await server.connect(transport);
            
            await new Promise((resolve) => {
                req.on('close', resolve);
            });
        });

        app.post('/message', express.json(), async (req, res) => {
            const sessionId = req.query.sessionId as string;
            const transport = transports.get(sessionId);

            console.error(`POST /message/${sessionId} - Active transports: ${transports.size}`);

            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 = Number(process.env.SERPBEAR_MCP_PORT) || 3004;
        app.listen(PORT, '0.0.0.0', () => {
            console.error(`SerpBear MCP server running on http://0.0.0.0:${PORT}/sse`);
        });
    }
}


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