import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import {
    CallToolRequestSchema,
    ListToolsRequestSchema,
    ToolSchema,
    ListResourcesRequestSchema,
    ReadResourceRequestSchema
} from "@modelcontextprotocol/sdk/types.js";
import { z } from "zod";
import { zodToJsonSchema } from "zod-to-json-schema";
import type { Props } from "../utils";
import { getThoughtSpotClient } from "../thoughtspot/thoughtspot-client";
import {
    type DataSource,
    fetchTMLAndCreateLiveboard,
    getAnswerForQuestion,
    getDataSources,
    getRelevantQuestions,
    getSessionInfo
} from "../thoughtspot/thoughtspot-service";
import { MixpanelTracker } from "../metrics/mixpanel/mixpanel";
import { Trackers, type Tracker, TrackEvent } from "../metrics";

const ToolInputSchema = ToolSchema.shape.inputSchema;
type ToolInput = z.infer<typeof ToolInputSchema>;

const PingSchema = z.object({});

const GetRelevantQuestionsSchema = z.object({
    query: z.string().describe("The query to get relevant data questions for, this could be a high level task or question the user is asking or hoping to get answered. Do minimal processing of the original question. You can even pass the complete raw query as it is, the system is smart to make sense of it as it has access to the entire schema. Do not add analytical hints or directions."),
    additionalContext: z.string()
        .describe("Additional context to add to the query, this might be older data returned for previous questions or any other relevant context that might help the system generate better questions.")
        .optional(),
    datasourceIds: z.array(z.string())
        .describe("The datasources to get questions for, this is the ids of the datasources to get data from. Each id is a GUID string.")
});

const GetAnswerSchema = z.object({
    question: z.string().describe("The question to get the answer for, these are generally the questions generated by the getRelevantQuestions tool."),
    datasourceId: z.string()
        .describe("The datasource to get the answer for, this is the id of the datasource to get data from")
});

const CreateLiveboardSchema = z.object({
    name: z.string().describe("The name of the liveboard to create"),
    answers: z.array(z.object({
        question: z.string(),
        session_identifier: z.string(),
        generation_number: z.number(),
    })).describe("The answers to create the liveboard from, these are the answers generated by the getAnswer tool."),
});

enum ToolName {
    Ping = "ping",
    GetRelevantQuestions = "getRelevantQuestions",
    GetAnswer = "getAnswer",
    CreateLiveboard = "createLiveboard",
}

interface Context {
    props: Props;
}

export class MCPServer extends Server {
    private trackers: Trackers = new Trackers();
    constructor(private ctx: Context) {
        super({
            name: "ThoughtSpot",
            version: "1.0.0",
        }, {
            capabilities: {
                tools: {},
                logging: {},
                completion: {},
                resources: {},
            }
        });
    }

    async init() {
        const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken);
        const sessionInfo = await getSessionInfo(client);
        const mixpanel = new MixpanelTracker(
            sessionInfo,
            this.ctx.props.clientName
        );
        this.addTracker(mixpanel);
        this.trackers.track(TrackEvent.Init);

        this.setRequestHandler(ListToolsRequestSchema, async () => {
            return {
                tools: [
                    {
                        name: ToolName.Ping,
                        description: "Simple ping tool to test connectivity and Auth",
                        inputSchema: zodToJsonSchema(PingSchema) as ToolInput,
                    },
                    {
                        name: ToolName.GetRelevantQuestions,
                        description: "Get relevant data questions from ThoughtSpot database",
                        inputSchema: zodToJsonSchema(GetRelevantQuestionsSchema) as ToolInput,
                    },
                    {
                        name: ToolName.GetAnswer,
                        description: "Get the answer to a question from ThoughtSpot database",
                        inputSchema: zodToJsonSchema(GetAnswerSchema) as ToolInput,
                    },
                    {
                        name: ToolName.CreateLiveboard,
                        description: "Create a liveboard from a list of answers",
                        inputSchema: zodToJsonSchema(CreateLiveboardSchema) as ToolInput,
                    }
                ]
            };
        });

        this.setRequestHandler(ListResourcesRequestSchema, async () => {
            const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken);
            const sources = await this.getDatasources();
            return {
                resources: sources.list.map((s) => ({
                    uri: `datasource:///${s.id}`,
                    name: s.name,
                    description: s.description,
                    mimeType: "text/plain"
                }))
            }
        });

        this.setRequestHandler(ReadResourceRequestSchema, async (request: z.infer<typeof ReadResourceRequestSchema>) => {
            const { uri } = request.params;
            const sourceId = uri.split("///").pop();
            if (!sourceId) {
                throw new Error("Invalid datasource uri");
            }
            const { map: sourceMap } = await this.getDatasources();
            const source = sourceMap.get(sourceId);
            if (!source) {
                throw new Error("Datasource not found");
            }
            return {
                contents: [{
                    uri: uri,
                    mimeType: "text/plain",
                    text: `
                    ${source.description}

                    The id of the datasource is ${sourceId}.

                    Use ThoughtSpot's getRelevantQuestions tool to get relevant questions for a query. And then use the getAnswer tool to get the answer for a question.
                    `,
                }],
            };
        });


        // Handle call tool request
        this.setRequestHandler(CallToolRequestSchema, async (request: z.infer<typeof CallToolRequestSchema>) => {
            const { name } = request.params;


            this.trackers.track(TrackEvent.CallTool, { toolName: name });

            switch (name) {
                case ToolName.Ping:
                    console.log("Received Ping request");
                    if (this.ctx.props.accessToken && this.ctx.props.instanceUrl) {
                        return {
                            content: [{ type: "text", text: "Pong" }],
                        };
                    }
                    return {
                        isError: true,
                        content: [{ type: "text", text: "ERROR: Not authenticated" }],
                    };

                case ToolName.GetRelevantQuestions: {
                    return this.callGetRelevantQuestions(request);
                }

                case ToolName.GetAnswer: {
                    return this.callGetAnswer(request);
                }

                case ToolName.CreateLiveboard: {
                    return this.callCreateLiveboard(request);
                }

                default:
                    throw new Error(`Unknown tool: ${name}`);
            }
        });
    }


    async callGetRelevantQuestions(request: z.infer<typeof CallToolRequestSchema>) {
        const { query, datasourceIds: sourceIds, additionalContext } = GetRelevantQuestionsSchema.parse(request.params.arguments);
        const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken);
        console.log("[DEBUG] Getting relevant questions for query: ", query, " and datasource: ", sourceIds);

        const relevantQuestions = await getRelevantQuestions(
            query,
            sourceIds!,
            additionalContext ?? "",
            client,
        );

        if (relevantQuestions.error) {
            return {
                isError: true,
                content: [{ type: "text", text: `ERROR: ${relevantQuestions.error.message}` }],
            };
        }

        if (relevantQuestions.questions.length === 0) {
            return {
                content: [{ type: "text", text: "No relevant questions found" }],
            };
        }

        return {
            content: relevantQuestions.questions.map(q => ({
                type: "text",
                text: `Question: ${q.question}\nDatasourceId: ${q.datasourceId}`,
            })),
        };
    }

    async callGetAnswer(request: z.infer<typeof CallToolRequestSchema>) {
        const { question, datasourceId: sourceId } = GetAnswerSchema.parse(request.params.arguments);
        const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken);
        const progressToken = request.params._meta?.progressToken;
        const progress = 0;
        console.log("[DEBUG] Getting answer for question: ", question, " and datasource: ", sourceId);

        const answer = await getAnswerForQuestion(question, sourceId, false, client);
        if (answer.error) {
            return {
                isError: true,
                content: [{ type: "text", text: `ERROR: ${answer.error.message}` }],
            };
        }

        return {
            content: [{
                type: "text",
                text: answer.data,
            }, {
                type: "text",
                text: `Question: ${question}\nSession Identifier: ${answer.session_identifier}\nGeneration Number: ${answer.generation_number} \n\nUse this information to create a liveboard with the createLiveboard tool, if the user asks.`,
            }],
        };
    }

    async callCreateLiveboard(request: z.infer<typeof CallToolRequestSchema>) {
        const { name, answers } = CreateLiveboardSchema.parse(request.params.arguments);
        const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken);
        const liveboard = await fetchTMLAndCreateLiveboard(name, answers, client);
        if (liveboard.error) {
            return {
                isError: true,
                content: [{ type: "text", text: `ERROR: ${liveboard.error.message}` }],
            };
        }
        return {
            content: [{
                type: "text",
                text: `Liveboard created successfully, you can view it at ${liveboard.url}
                
                Provide this url to the user as a link to view the liveboard in ThoughtSpot.`,
            }],
        };
    }

    private _sources: {
        list: DataSource[];
        map: Map<string, DataSource>;
    } | null = null;
    async getDatasources() {
        if (this._sources) {
            return this._sources;
        }

        const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken);
        const sources = await getDataSources(client);
        this._sources = {
            list: sources,
            map: new Map(sources.map(s => [s.id, s])),
        }
        return this._sources;
    }

    async addTracker(tracker: Tracker) {
        this.trackers.add(tracker);
    }
}