import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { z } from 'zod';
import fs from 'fs';
import { createApiClient } from '../lib/api.js';
import { draftManager } from '../lib/draft-manager.js';
import { requestContext } from '../context.js';
import { logger } from '../lib/logger.js';
import { config } from '../config.js';
import { oauth2Schemes, setToolSecuritySchemes, handleApiError } from '../lib/tool-security.js';
import { UI_TEMPLATE_URI, withWidgetData } from '../lib/ui.js';

const checkAuth = () => {
  const context = requestContext.getStore();
  if (!context?.token) {
    throw new Error("AUTHENTICATION_REQUIRED");
  }
  return context;
};

export const registerProofTools = (server: McpServer) => {
  server.registerTool(
    "list_proof_templates",
    {
      title: "List Proof Templates",
      description: "List available proof templates.",
      inputSchema: {
        raw: z.boolean().optional().default(false).describe("If true, returns the raw JSON list")
      },
      _meta: {
        "openai/outputTemplate": UI_TEMPLATE_URI,
        "openai/widgetAccessible": true
      }
    },
    async ({ raw }) => {
      try {
        const context = checkAuth();
        const client = createApiClient();
        const { data: response } = await client.get<any>('/api/v1/proof-models?per_page=200');

        const items: any[] = Array.isArray(response) ? response : (response.data || []);

        const normalizeSchema = (template: any) => {
          let schema: any = template.data_schema || template.data || {};
          if (typeof schema === 'string') {
            try {
              schema = JSON.parse(schema);
            } catch {
              schema = {};
            }
          }

          let customFields: any =
            schema.customFields ||
            schema.custom_fields ||
            template.custom_fields ||
            template.customFields ||
            [];

          if (!Array.isArray(customFields) && customFields && typeof customFields === 'object') {
            customFields = Object.entries(customFields).map(([key, value]) => ({
              key,
              ...(value as Record<string, unknown>)
            }));
          }

          return {
            ...schema,
            customFields
          };
        };

        // Prepare structured content for the UI
        const templates = items.map(t => ({
          id: t.id,
          name: t.name || t.title || 'Unnamed Template',
          description: t.description || 'No description',
          schema: normalizeSchema(t)
        }));

        const structuredContent = withWidgetData({ templates, screen: "create_proof" });

        let output = "";
        if (items.length === 0) {
          output = "No proof templates available.";
        } else {
          output = "### Available Proof Templates\n\n";
          items.forEach((t) => {
            const name = t.name || t.title || 'Unnamed Template';
            output += `- **${name}** (ID: 
${t.id}
)\n`;
          });
          output += "\nI've opened the ProovIT dashboard where you can see all templates and their details.";
        }

        return { 
          content: [{ type: "text" as const, text: raw ? JSON.stringify(items, null, 2) : output }],
          structuredContent,
          _meta: {
            apiUrl: config.publicUrl,
            sessionId: context.sessionId
          }
        };
      } catch (error: any) {
        return handleApiError(error);
      }
    }
  );

  setToolSecuritySchemes("list_proof_templates", oauth2Schemes());

  server.registerTool(
    "list_drafts",
    {
      title: "List Drafts",
      description: "List current local proof drafts.",
      inputSchema: {
        raw: z.boolean().optional().default(false).describe("If true, returns the raw JSON")
      },
      _meta: {
        "openai/outputTemplate": UI_TEMPLATE_URI,
        "openai/widgetAccessible": true
      }
    },
    async ({ raw }) => {
      try {
        const context = checkAuth();
        const drafts = await draftManager.listDrafts();
        
        const structuredContent = withWidgetData({ drafts, screen: "drafts" });

        if (drafts.length === 0) {
          return { 
            content: [{ type: "text" as const, text: "No drafts found." }],
            structuredContent
          };
        }

        let summary = "### Current Local Drafts\n\n";
        drafts.forEach(d => {
          summary += `- **${d.fields.title || 'Untitled'}** (ID: 
${d.id}
)\n`;
        });

        return { 
          content: [{ type: "text" as const, text: raw ? JSON.stringify(drafts, null, 2) : summary }],
          structuredContent,
          _meta: {
            apiUrl: config.publicUrl,
            sessionId: context.sessionId
          }
        };
      } catch (error: any) {
        return handleApiError(error);
      }
    }
  );

  setToolSecuritySchemes("list_drafts", oauth2Schemes());

  server.registerTool(
    "get_draft",
    {
      title: "Get Draft",
      description: "Get full details of the current active proof draft."
    },
    async () => {
      try {
        checkAuth();
        const draft = await draftManager.getDraft();
        const payload = draft ?? null;
        return { content: [{ type: "text" as const, text: JSON.stringify(payload, null, 2) }] };
      } catch (e: any) {
        return handleApiError(e);
      }
    }
  );

  setToolSecuritySchemes("get_draft", oauth2Schemes());

  server.registerTool(
    "create_draft",
    {
      title: "Create Draft",
      description: "Create or update the current local draft for a proof.",
      inputSchema: {
        templateId: z.string().describe("The ProovIT Template ID"),
        title: z.string().describe("The title of the proof"),
        description: z.string().optional(),
        categoryId: z.string().optional(),
        tagIds: z.array(z.string()).optional(),
        customFields: z.union([z.string(), z.record(z.string(), z.any())]).optional().describe("Custom fields values as an object or JSON string")
      }
    },
    async ({ templateId, title, ...otherFields }) => {
      try {
        checkAuth();
        const client = createApiClient();
        const { data: response } = await client.get<any>('/api/v1/proof-models?per_page=200');
        const items: any[] = Array.isArray(response) ? response : (response.data || []);
        const template = items.find(t => t.id === templateId);

        if (!template) {
          return { isError: true, content: [{ type: "text" as const, text: `Template with ID "${templateId}" not found.` }] };
        }

        let initialFields: any = { title, ...otherFields };
        if (otherFields.customFields && typeof otherFields.customFields === 'string') {
          try {
            initialFields.customFields = JSON.parse(otherFields.customFields);
          } catch (e) {
            return { isError: true, content: [{ type: "text" as const, text: "Invalid JSON in customFields parameter." }] };
          }
        }

        const draft = await draftManager.createDraft(templateId, initialFields);

        let msg = `Draft "${title}" initialized successfully.\n`;

        let schema = template.data_schema || template.data || {};
        if (typeof schema === 'string') {
          try { schema = JSON.parse(schema); } catch (e) { schema = {}; }
        }
        const fields = schema.customFields || template.custom_fields || template.customFields || [];

        const required = fields.filter((f: any) => f.required);
        if (required.length > 0) {
          msg += `\n**Required fields for this template:**\n`;
          required.forEach((f: any) => {
            msg += `- 
${f.key || f.name}
: ${f.label || f.title} (${f.type || 'text'})
`;
          });
        }

        return { content: [{ type: "text" as const, text: msg }] };
      } catch (error: any) {
        return handleApiError(error);
      }
    }
  );

  setToolSecuritySchemes("create_draft", oauth2Schemes());

  server.registerTool(
    "update_draft",
    {
      title: "Update Draft",
      description: "Update fields and metadata in the current active proof draft.",
      inputSchema: {
        title: z.string().optional(),
        description: z.string().optional(),
        categoryId: z.string().optional(),
        tagIds: z.array(z.string()).optional(),
        customFields: z.union([z.string(), z.record(z.string(), z.any())]).optional().describe("Custom fields values as an object or JSON string")
      }
    },
    async ({ ...updates }) => {
      try {
        checkAuth();
        let fieldsToUpdate: any = { ...updates };
        if (updates.customFields) {
          if (typeof updates.customFields === 'string') {
            try {
              fieldsToUpdate.customFields = JSON.parse(updates.customFields);
            } catch (e) {
              return { isError: true, content: [{ type: "text" as const, text: "Invalid JSON in customFields parameter." }] };
            }
          } else {
            fieldsToUpdate.customFields = updates.customFields;
          }
        }
        await draftManager.updateDraftFields(fieldsToUpdate);
        return {
          content: [{ type: "text" as const, text: `Draft updated. Use 
get_draft
 to see current state.` }]
        };
      } catch (e: any) {
        return handleApiError(e);
      }
    }
  );

  setToolSecuritySchemes("update_draft", oauth2Schemes());

  server.registerTool(
    "upload_file_to_draft",
    {
      title: "Upload File To Draft",
      description: "Upload a file (Base64) to the current active draft.",
      inputSchema: {
        filename: z.string(),
        base64Content: z.string(),
        mimeType: z.string().optional()
      }
    },
    async ({ filename, base64Content, mimeType }) => {
      try {
        checkAuth();
        await draftManager.addFile(filename, base64Content, mimeType);
        return {
          content: [{ type: "text" as const, text: `File "${filename}" successfully uploaded to the active draft.` }]
        };
      } catch (e: any) {
        return handleApiError(e);
      }
    }
  );

  setToolSecuritySchemes("upload_file_to_draft", oauth2Schemes());

  server.registerTool(
    "delete_draft",
    {
      title: "Delete Draft",
      description: "Delete the current active proof draft."
    },
    async () => {
      try {
        checkAuth();
        await draftManager.deleteDraft();
        return { content: [{ type: "text" as const, text: "Active draft has been deleted." }] };
      } catch (error: any) {
        return handleApiError(error);
      }
    }
  );

  setToolSecuritySchemes("delete_draft", oauth2Schemes());

  server.registerTool(
    "submit_proof",
    {
      title: "Submit Proof",
      description: "Submit the active draft to ProovIT API and generate the certificate."
    },
    async () => {
      try {
        checkAuth();
        const draft = await draftManager.getDraft();
        if (!draft.files || draft.files.length === 0) {
          return {
            isError: true,
            content: [{
              type: "text" as const,
              text: "At least one file is required before submitting a proof. Please upload a file to the draft."
            }]
          };
        }
        const client = createApiClient();

        const { data: templatesResponse } = await client.get<any>('/api/v1/proof-models?per_page=200');
        const templates: any[] = Array.isArray(templatesResponse) ? templatesResponse : (templatesResponse.data || []);
        const template = templates.find(t => t.id === draft.templateId);

        if (!template) {
          throw new Error(`Template 
${draft.templateId}
 not found for this draft.`);
        }

        let schema = template.data_schema || template.data || {};
        if (typeof schema === 'string') {
          try { schema = JSON.parse(schema); } catch (e) { schema = {}; }
        }
        const fields = schema.customFields || template.custom_fields || template.customFields || [];

        const requiredFields = fields.filter((f: any) => f.required);
        const providedFields = draft.fields.customFields || {};
        const missingFields = requiredFields.filter((f: any) => {
          const val = providedFields[f.key || f.name];
          return !val || String(val).trim() === "";
        });

        if (missingFields.length > 0) {
          const fieldList = missingFields.map((f: any) => `
${f.key || f.name}
 (${f.label || f.title})`).join(', ');
          return {
            isError: true,
            content: [{
              type: "text" as const,
              text: `❌ Cannot submit: Missing required custom fields: ${fieldList}. Please use 
update_draft
 to provide them.`
            }]
          };
        }

        const initPayload = {
          name: draft.fields.title,
          description: draft.fields.description,
          metadata: {
            category_id: draft.fields.categoryId,
            tags_id: draft.fields.tagIds || [],
            custom_fields: providedFields
          },
          proof_template_id: draft.templateId
        };

        const initResponse = await client.post('/api/v1/proofs/init', initPayload);
        const proofId = initResponse.data.id || initResponse.data.data?.id;

        if (!proofId) {
          throw new Error("API did not return a proof ID.");
        }

        if (draft.files.length > 0) {
          const formData = new FormData();
          for (const file of draft.files) {
            const filePath = await draftManager.getFilePath(file.filename);
            try {
              const fileBuffer = await fs.promises.readFile(filePath);
              const blob = new Blob([fileBuffer], { type: file.mimeType });
              formData.append('files[]', blob, file.filename);
            } catch (fileError: any) {
              logger.error(`Error reading file ${file.filename}`, fileError);
              throw new Error(`File ${file.filename} not found in draft.`);
            }
          }
          await client.post(`/api/v1/proofs/${proofId}/files`, formData);
        }

        const certResponse = await client.post(`/api/v1/proofs/${proofId}/certificate-link`);
        const certificateUrl = certResponse.data.url || certResponse.data.data?.url;

        const detailResponse = await client.get<any>(`/api/v1/proofs/${proofId}?include=files,template`);
        const proofDetail = detailResponse.data?.data ?? detailResponse.data;
        const isTerminated = Boolean(proofDetail?.state?.terminated || proofDetail?.terminated);
        const isSubmitted = Boolean(proofDetail?.state?.submitted);
        const certificateFromDetail =
          proofDetail?.metadata?.certificate_url ||
          proofDetail?.links?.certificate_pdf;

        try {
          await draftManager.deleteDraft();
        } catch (err) {
          logger.error(`Error deleting draft`, err);
        }

        if (isTerminated && isSubmitted) {
          const finalCertificateUrl = certificateUrl || certificateFromDetail;
          return {
            content: [{
              type: "text" as const,
              text: `✅ Proof submitted successfully!\n\n**ID:** ${proofId}\n**Certificate:** ${finalCertificateUrl}`
            }]
          };
        }

        return {
          content: [{
            type: "text" as const,
            text: `✅ Proof submitted successfully!\n\n**ID:** ${proofId}\n\nThe proof is being finalized. You will receive an email when it is ready.`
          }]
        };
      } catch (error: any) {
        return handleApiError(error);
      }
    }
  );

  setToolSecuritySchemes("submit_proof", oauth2Schemes());

  server.registerTool(
    "get_proof_status",
    {
      title: "Get Proof Status",
      description: "Get the current status and details of a proof by ID.",
      inputSchema: {
        proofId: z.string().describe("The proof ID"),
        raw: z.boolean().optional().default(false).describe("If true, returns the raw JSON")
      }
    },
    async ({ proofId, raw }) => {
      try {
        checkAuth();
        const client = createApiClient();
        const { data: response } = await client.get<any>(`/api/v1/proofs/${proofId}?include=files,template`);
        const payload = response?.data ?? response ?? null;

        if (raw) {
          return { content: [{ type: "text" as const, text: JSON.stringify(payload, null, 2) }] };
        }

        const status = payload?.status || payload?.state?.status || 'unknown';
        const statusLabel = payload?.status_label || '';
        const filesCount = payload?.files_count ?? (Array.isArray(payload?.files) ? payload.files.length : 0);
        const terminated = Boolean(payload?.state?.terminated || payload?.terminated);
        const submitted = Boolean(payload?.state?.submitted);

        const summary = [
          `Status: ${status}${statusLabel ? ` (${statusLabel})` : ''}`,
          `Files: ${filesCount}`,
          `Submitted: ${submitted ? 'yes' : 'no'}`,
          `Terminated: ${terminated ? 'yes' : 'no'}`
        ].join('\n');

        return { content: [{ type: "text" as const, text: summary }] };
      } catch (error: any) {
        return handleApiError(error);
      }
    }
  );

  setToolSecuritySchemes("get_proof_status", oauth2Schemes());
};