import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { CallToolRequestSchema, ListToolsRequestSchema } from "@modelcontextprotocol/sdk/types.js";
import { JMAPClient } from "./jmap_client.js";
import * as dotenv from "dotenv";

dotenv.config();

const server = new Server(
    {
        name: "stalwart_mcp",
        version: "1.0.0",
    },
    {
        capabilities: {
            tools: {},
        },
    }
);

const JMAP_URL = process.env.STALWART_JMAP_URL || "https://mail.infra.mintel.me";
const JMAP_USER = process.env.STALWART_USER || "antigravity";
const JMAP_PASSWORD = process.env.STALWART_APP_PASSWORD || "90e21uh0eondwqa9i0owsdn";

const client = new JMAPClient(JMAP_URL, JMAP_USER, JMAP_PASSWORD);

server.setRequestHandler(ListToolsRequestSchema, async () => {
    return {
        tools: [
            {
                name: "stalwart_list_mailboxes",
                description: "List all mailboxes (folders) for the authenticated account.",
                inputSchema: {
                    type: "object",
                    properties: {},
                },
            },
            {
                name: "stalwart_search_emails",
                description: "Search for emails in a specific mailbox or across all mailboxes.",
                inputSchema: {
                    type: "object",
                    properties: {
                        mailboxId: {
                            type: "string",
                            description: "Optional mailbox ID to filter the search.",
                        },
                        query: {
                            type: "string",
                            description: "Search query (e.g., 'from:someone', 'subject:hello').",
                        },
                        limit: {
                            type: "number",
                            description: "Maximum number of emails to return (default 10).",
                        }
                    },
                },
            },
            {
                name: "stalwart_read_email",
                description: "Read the full contents of a specific email by its ID.",
                inputSchema: {
                    type: "object",
                    properties: {
                        emailId: {
                            type: "string",
                            description: "The ID of the email to read.",
                        },
                    },
                    required: ["emailId"],
                },
            },
            {
                name: "stalwart_send_email",
                description: "Send a new email.",
                inputSchema: {
                    type: "object",
                    properties: {
                        to: {
                            type: "string",
                            description: "Recipient email address.",
                        },
                        subject: {
                            type: "string",
                            description: "Subject of the email.",
                        },
                        textBody: {
                            type: "string",
                            description: "Plain text body of the email.",
                        },
                        htmlBody: {
                            type: "string",
                            description: "Optional HTML body of the email.",
                        }
                    },
                    required: ["to", "subject", "textBody"],
                },
            }
        ],
    };
});

server.setRequestHandler(CallToolRequestSchema, async (request) => {
    const { name, arguments: args } = request.params;

    try {
        if (name === "stalwart_list_mailboxes") {
            const response = await client.request([
                ["Mailbox/get", { accountId: client.getAccountId(), ids: null }, "0"]
            ]);
            
            const mailboxes = response.methodResponses[0][1].list;
            return {
                content: [{ type: "text", text: JSON.stringify(mailboxes, null, 2) }],
            };
        } 
        
        else if (name === "stalwart_search_emails") {
            const { mailboxId, query, limit = 10 } = args as any;
            
            // Note: A full search implementation in JMAP is slightly more complex, 
            // but we can query Email/query.
            const filter: any = {};
            if (mailboxId) filter.inMailbox = mailboxId;
            if (query) filter.text = query; // Simple text search

            const queryResponse = await client.request([
                ["Email/query", { 
                    accountId: client.getAccountId(), 
                    filter: Object.keys(filter).length > 0 ? filter : null,
                    limit 
                }, "0"],
                ["Email/get", {
                    accountId: client.getAccountId(),
                    "#ids": {
                        resultOf: "0",
                        name: "Email/query",
                        path: "/ids"
                    },
                    properties: ["id", "subject", "from", "to", "receivedAt", "preview"]
                }, "1"]
            ]);

            const emails = queryResponse.methodResponses[1][1].list;
            return {
                content: [{ type: "text", text: JSON.stringify(emails, null, 2) }],
            };
        }
        
        else if (name === "stalwart_read_email") {
            const { emailId } = args as any;
            const response = await client.request([
                ["Email/get", {
                    accountId: client.getAccountId(),
                    ids: [emailId],
                    properties: ["id", "subject", "from", "to", "receivedAt", "bodyValues", "textBody", "htmlBody"]
                }, "0"]
            ]);
            
            const email = response.methodResponses[0][1].list[0];
            if (!email) {
                return { content: [{ type: "text", text: `Email with ID ${emailId} not found.` }] };
            }
            
            return {
                content: [{ type: "text", text: JSON.stringify(email, null, 2) }],
            };
        }
        
        else if (name === "stalwart_send_email") {
            const { to, subject, textBody, htmlBody } = args as any;
            
            const accountId = client.getAccountId();
            const creationId = "draft-" + Date.now();
            
            const bodyParts: any[] = [
                { partId: "part_text", type: "text/plain" }
            ];
            
            const bodyValues: any = {
                "part_text": { value: textBody, isEncodingProblem: false, isTruncated: false }
            };
            
            if (htmlBody) {
                bodyParts.push({ partId: "part_html", type: "text/html" });
                bodyValues["part_html"] = { value: htmlBody, isEncodingProblem: false, isTruncated: false };
            }

            const response = await client.request([
                ["Email/set", {
                    accountId,
                    create: {
                        [creationId]: {
                            from: [{ email: JMAP_USER.includes("@") ? JMAP_USER : `${JMAP_USER}@mintel.me` }],
                            to: [{ email: to }],
                            subject: subject,
                            bodyValues: bodyValues,
                            textBody: [{ partId: "part_text" }],
                            ...(htmlBody ? { htmlBody: [{ partId: "part_html" }] } : {})
                        }
                    }
                }, "0"],
                ["EmailSubmission/set", {
                    accountId,
                    create: {
                        "send-1": {
                            emailId: `#${creationId}`
                        }
                    }
                }, "1"]
            ]);

            return {
                content: [{ type: "text", text: JSON.stringify(response, null, 2) }],
            };
        }

        throw new Error(`Unknown tool: ${name}`);
    } catch (error: any) {
        return {
            isError: true,
            content: [{ type: "text", text: `Error executing tool: ${error.message}\n${JSON.stringify(error.response?.data || {})}` }],
        };
    }
});

async function run() {
    // Initial connection test
    try {
        await client.connect();
        console.error("Successfully connected to Stalwart JMAP!");
    } catch (e) {
        console.error("Failed to connect to Stalwart on startup. Ensure credentials are correct.");
        process.exit(1);
    }

    const transport = new StdioServerTransport();
    await server.connect(transport);
    console.error("Stalwart MCP Server running on stdio");
}

run().catch(console.error);