import type { GustoClient } from '../clients/gusto.js';
import type { FederalTax, StateTax } from '../types/index.js';

export function createTaxTools(client: GustoClient) {
  return {
    // List federal tax setup for employees
    list_federal_taxes: async (args: { company_id: string }) => {
      const employees = await client.getAllPaginated(
        `/companies/${args.company_id}/employees`
      );
      const federalTaxes = [];

      for (const employee of employees as any[]) {
        try {
          const tax = await client.get<FederalTax>(
            `/employees/${employee.id}/federal_taxes`
          );
          federalTaxes.push({ ...tax, employee_id: employee.id });
        } catch (error) {
          // Employee might not have federal tax setup yet
        }
      }

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

    // Get federal tax for an employee
    get_federal_tax: async (args: { employee_id: string }) => {
      const tax = await client.get<FederalTax>(`/employees/${args.employee_id}/federal_taxes`);
      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify(tax, null, 2),
          },
        ],
      };
    },

    // Update federal tax withholding for an employee
    update_federal_tax: async (args: {
      employee_id: string;
      version: string;
      filing_status?: string;
      allowances?: number;
      extra_withholding?: string;
      two_jobs?: boolean;
      dependents_amount?: string;
      other_income?: string;
      deductions?: string;
    }) => {
      const tax = await client.put<FederalTax>(
        `/employees/${args.employee_id}/federal_taxes`,
        {
          version: args.version,
          filing_status: args.filing_status,
          allowances: args.allowances,
          extra_withholding: args.extra_withholding,
          two_jobs: args.two_jobs,
          dependents_amount: args.dependents_amount,
          other_income: args.other_income,
          deductions: args.deductions,
        }
      );
      return {
        content: [
          {
            type: 'text',
            text: `Federal tax updated successfully:\n${JSON.stringify(tax, null, 2)}`,
          },
        ],
      };
    },

    // List state taxes for an employee
    list_state_taxes: async (args: { employee_id: string }) => {
      const taxes = await client.get<StateTax[]>(
        `/employees/${args.employee_id}/state_taxes`
      );
      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify(taxes, null, 2),
          },
        ],
      };
    },

    // Update state tax withholding for an employee
    update_state_tax: async (args: {
      employee_id: string;
      state: string;
      version: string;
      filing_status?: string;
      allowances?: number;
      extra_withholding?: string;
      supplemental_withholding?: string;
    }) => {
      const tax = await client.put<StateTax>(
        `/employees/${args.employee_id}/state_taxes/${args.state}`,
        {
          version: args.version,
          filing_status: args.filing_status,
          allowances: args.allowances,
          extra_withholding: args.extra_withholding,
          supplemental_withholding: args.supplemental_withholding,
        }
      );
      return {
        content: [
          {
            type: 'text',
            text: `State tax updated successfully:\n${JSON.stringify(tax, null, 2)}`,
          },
        ],
      };
    },
  };
}

export const taxToolDefinitions = [
  {
    name: 'list_federal_taxes',
    description: 'List federal tax withholding setup for all employees in a company',
    inputSchema: {
      type: 'object',
      properties: {
        company_id: { type: 'string', description: 'Company ID' },
      },
      required: ['company_id'],
    },
  },
  {
    name: 'get_federal_tax',
    description: 'Get federal tax withholding information for an employee',
    inputSchema: {
      type: 'object',
      properties: {
        employee_id: { type: 'string', description: 'Employee ID' },
      },
      required: ['employee_id'],
    },
  },
  {
    name: 'update_federal_tax',
    description: 'Update federal tax withholding for an employee',
    inputSchema: {
      type: 'object',
      properties: {
        employee_id: { type: 'string', description: 'Employee ID' },
        version: { type: 'string', description: 'Version for optimistic locking' },
        filing_status: {
          type: 'string',
          description: 'Filing status (Single, Married, Head of Household)',
        },
        allowances: { type: 'number', description: 'Number of allowances' },
        extra_withholding: { type: 'string', description: 'Extra withholding amount' },
        two_jobs: { type: 'boolean', description: 'Whether employee has two jobs' },
        dependents_amount: { type: 'string', description: 'Dependents credit amount' },
        other_income: { type: 'string', description: 'Other income amount' },
        deductions: { type: 'string', description: 'Deductions amount' },
      },
      required: ['employee_id', 'version'],
    },
  },
  {
    name: 'list_state_taxes',
    description: 'List state tax withholding for an employee',
    inputSchema: {
      type: 'object',
      properties: {
        employee_id: { type: 'string', description: 'Employee ID' },
      },
      required: ['employee_id'],
    },
  },
  {
    name: 'update_state_tax',
    description: 'Update state tax withholding for an employee',
    inputSchema: {
      type: 'object',
      properties: {
        employee_id: { type: 'string', description: 'Employee ID' },
        state: { type: 'string', description: 'State code (e.g., CA, NY)' },
        version: { type: 'string', description: 'Version for optimistic locking' },
        filing_status: { type: 'string', description: 'State filing status' },
        allowances: { type: 'number', description: 'Number of allowances' },
        extra_withholding: { type: 'string', description: 'Extra withholding amount' },
        supplemental_withholding: {
          type: 'string',
          description: 'Supplemental withholding rate',
        },
      },
      required: ['employee_id', 'state', 'version'],
    },
  },
];
