import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { z } from 'zod';
import { createApiClient } from '../lib/api.js';
import { requestContext } from '../context.js';
import { oauth2Schemes, setToolSecuritySchemes, handleApiError } from '../lib/tool-security.js';
import { UI_TEMPLATE_URI, withWidgetData } from '../lib/ui.js';
import { UserProfile, TokenBalance, ApiResponse } from '../types/api.js';
import { config } from '../config.js';

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

export const registerUserTools = (server: McpServer) => {
  server.registerTool(
    "whoami",
    {
      title: "Who Am I",
      description: "Get current user profile.",
      inputSchema: {
        raw: z.boolean().optional().default(false).describe("If true, returns the raw JSON from the API")
      },
      _meta: {
        "openai/outputTemplate": UI_TEMPLATE_URI,
        "openai/widgetAccessible": true
      }
    },
    async ({ raw }) => {
      try {
        const context = checkAuth();
        const client = createApiClient();
        const path = '/api/v1/me';
        const apiResponse = await client.get<ApiResponse<UserProfile>>(path);
        const payload: any = apiResponse.data?.data ?? apiResponse.data;
        const user = payload?.user ?? payload;
        
        const name = user?.name || '';
        const email = user?.email || '';
        const summary = email
          ? `You are authenticated as ${name} (${email}).`
          : 'You are authenticated.';

        const content = raw 
          ? JSON.stringify(payload ?? null, null, 2)
          : summary;

        return { 
          content: [{ type: "text" as const, text: content }],
          structuredContent: withWidgetData({
            user,
            isAuthenticated: true,
            screen: "profile"
          }),
          _meta: {
            apiUrl: config.publicUrl,
            sessionId: context.sessionId
          }
        };
      } catch (error: any) {
        return handleApiError(error);
      }
    }
  );

  setToolSecuritySchemes("whoami", oauth2Schemes());

  server.registerTool(
    "get_token_balance",
    {
      title: "Get Token Balance",
      description: "Get the current token balance of the user.",
      inputSchema: {
        raw: z.boolean().optional().default(false).describe("If true, returns the raw JSON from the API")
      }
    },
    async ({ raw }) => {
      try {
        checkAuth();
        const client = createApiClient();
        const path = '/api/v1/tokens/balance';
        const apiResponse = await client.get<ApiResponse<TokenBalance>>(path);
        const payload: any = apiResponse.data?.data ?? apiResponse.data;
        const data = payload as TokenBalance;
        
        if (raw) {
          const rawPayload = payload ?? null;
          return { content: [{ type: "text" as const, text: JSON.stringify(rawPayload, null, 2) }] };
        }

        return { 
          content: [{ type: "text" as const, text: `You currently have ${data.balance ?? 0} tokens available.` }]
        };
      } catch (error: any) {
        return handleApiError(error);
      }
    }
  );

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