// src/services/db/budget.db.test.ts
import { describe, it, expect, vi, beforeEach, Mock } from 'vitest';
import { ForeignKeyConstraintError } from './errors.db';

// Un-mock the module we are testing to ensure we use the real implementation.
vi.unmock('./budget.db');

import { BudgetRepository } from './budget.db';
import type { Pool, PoolClient } from 'pg';
import { mockPoolInstance } from '../../tests/setup/tests-setup-unit';
import type { Budget, SpendingByCategory } from '../../types';

// Mock the logger to prevent console output during tests
vi.mock('../logger.server', () => ({
  logger: {
    info: vi.fn(),
    warn: vi.fn(),
    error: vi.fn(), // Keep warn for other tests that might use it
    debug: vi.fn(),
  },
}));
import { logger as mockLogger } from '../logger.server';

// Mock the withTransaction helper
vi.mock('./connection.db', async (importOriginal) => {
  const actual = await importOriginal<typeof import('./connection.db')>();
  return { ...actual, withTransaction: vi.fn() };
});

const { mockedAwardAchievement } = vi.hoisted(() => ({
  mockedAwardAchievement: vi.fn(),
}));

// Mock the gamification repository, as createBudget calls it.
vi.mock('./gamification.db', () => ({
  GamificationRepository: class {
    awardAchievement = mockedAwardAchievement;
  },
}));

import { withTransaction } from './connection.db';

describe('Budget DB Service', () => {
  let budgetRepo: BudgetRepository;

  beforeEach(() => {
    vi.clearAllMocks();
    // Instantiate the repository with the mock pool for each test
    budgetRepo = new BudgetRepository(mockPoolInstance as unknown as Pool);
  });

  describe('getBudgetsForUser', () => {
    it('should execute the correct SELECT query and return budgets', async () => {
      const mockBudgets: Budget[] = [{ budget_id: 1, user_id: 'user-123', name: 'Groceries', amount_cents: 50000, period: 'monthly', start_date: '2024-01-01' }];
      mockPoolInstance.query.mockResolvedValue({ rows: mockBudgets });
 
      const result = await budgetRepo.getBudgetsForUser('user-123', mockLogger);

      expect(mockPoolInstance.query).toHaveBeenCalledWith('SELECT * FROM public.budgets WHERE user_id = $1 ORDER BY start_date DESC', ['user-123']);
      expect(result).toEqual(mockBudgets);
    });

    it('should return an empty array if the user has no budgets', async () => {
      mockPoolInstance.query.mockResolvedValue({ rows: [] });
      const result = await budgetRepo.getBudgetsForUser('user-123', mockLogger);
      expect(result).toEqual([]);
      expect(mockPoolInstance.query).toHaveBeenCalledWith(expect.any(String), ['user-123']);
    });

    it('should throw an error if the database query fails', async () => {
      const dbError = new Error('DB Error');
      mockPoolInstance.query.mockRejectedValue(dbError);
      await expect(budgetRepo.getBudgetsForUser('user-123', mockLogger)).rejects.toThrow('Failed to retrieve budgets.');
      expect(mockLogger.error).toHaveBeenCalledWith({ err: dbError, userId: 'user-123' }, 'Database error in getBudgetsForUser');
    });
  });

  describe('createBudget', () => {
    it('should execute an INSERT query and return the new budget', async () => {
      const budgetData = { name: 'Groceries', amount_cents: 50000, period: 'monthly' as const, start_date: '2024-01-01' };
      const mockCreatedBudget: Budget = { budget_id: 1, user_id: 'user-123', ...budgetData };
      
      // Create a mock client that we can reference both inside and outside the transaction mock.
      const mockClient = { query: vi.fn() };

      vi.mocked(withTransaction).mockImplementation(async (callback) => {
        // Configure the mock client's query behavior for this specific test.
        (mockClient.query as Mock)
          .mockResolvedValueOnce({ rows: [mockCreatedBudget] }) // For the INSERT...RETURNING
          .mockResolvedValueOnce({ rows: [] }); // For award_achievement
        return callback(mockClient as unknown as PoolClient);
      });
      
      const result = await budgetRepo.createBudget('user-123', budgetData, mockLogger);
      
      // Now we can assert directly on the mockClient we created.
      expect(mockClient.query).toHaveBeenCalledWith(expect.stringContaining('INSERT INTO public.budgets'), expect.any(Array));
      expect(mockedAwardAchievement).toHaveBeenCalledWith('user-123', 'First Budget Created', mockLogger);
      expect(result).toEqual(mockCreatedBudget);
      expect(withTransaction).toHaveBeenCalledTimes(1);
    });

    it('should throw ForeignKeyConstraintError if user does not exist', async () => {
      const budgetData = { name: 'Groceries', amount_cents: 50000, period: 'monthly' as const, start_date: '2024-01-01' };
      const dbError = new Error('violates foreign key constraint');
      (dbError as Error & { code: string }).code = '23503';

      vi.mocked(withTransaction).mockImplementation(async (callback) => {
        const mockClient = { query: vi.fn() };
        mockClient.query.mockRejectedValueOnce(dbError); // INSERT fails
        await expect(callback(mockClient as unknown as PoolClient)).rejects.toThrow(dbError);
        throw dbError; // Re-throw for the outer expect
      });

      await expect(budgetRepo.createBudget('non-existent-user', budgetData, mockLogger)).rejects.toThrow(ForeignKeyConstraintError);
      await expect(budgetRepo.createBudget('non-existent-user', budgetData, mockLogger)).rejects.toThrow('The specified user does not exist.');
    });

    it('should rollback the transaction if awarding an achievement fails', async () => {
      const budgetData = { name: 'Groceries', amount_cents: 50000, period: 'monthly' as const, start_date: '2024-01-01' };
      const mockCreatedBudget: Budget = { budget_id: 1, user_id: 'user-123', ...budgetData };
      const achievementError = new Error('Achievement award failed');

      mockedAwardAchievement.mockRejectedValueOnce(achievementError);

      vi.mocked(withTransaction).mockImplementation(async (callback) => {
        const mockClient = { query: vi.fn() };
        (mockClient.query as Mock)
          .mockResolvedValueOnce({ rows: [mockCreatedBudget] }); // INSERT...RETURNING
        
        await expect(callback(mockClient as unknown as PoolClient)).rejects.toThrow(achievementError);
        throw achievementError; // Re-throw for the outer expect
      });

      await expect(budgetRepo.createBudget('user-123', budgetData, mockLogger)).rejects.toThrow('Failed to create budget.');
      expect(mockLogger.error).toHaveBeenCalledWith({ err: achievementError, budgetData, userId: 'user-123' }, 'Database error in createBudget');
    });

    it('should throw a generic error if the database query fails', async () => {
      const budgetData = { name: 'Groceries', amount_cents: 50000, period: 'monthly' as const, start_date: '2024-01-01' };
      const dbError = new Error('DB Error');
      vi.mocked(withTransaction).mockImplementation(async (callback) => {
        const mockClient = { query: vi.fn() };
        mockClient.query.mockRejectedValueOnce(dbError); // INSERT fails
        await expect(callback(mockClient as unknown as PoolClient)).rejects.toThrow(dbError);
        throw dbError; // Re-throw for the outer expect
      });
      await expect(budgetRepo.createBudget('user-123', budgetData, mockLogger)).rejects.toThrow('Failed to create budget.');
      expect(mockLogger.error).toHaveBeenCalledWith({ err: dbError, budgetData, userId: 'user-123' }, 'Database error in createBudget');
    });
  });

  describe('updateBudget', () => {
    it('should execute an UPDATE query with COALESCE and return the updated budget', async () => {
      const budgetUpdates = { name: 'Updated Groceries', amount_cents: 55000 };
      const mockUpdatedBudget: Budget = { budget_id: 1, user_id: 'user-123', name: 'Updated Groceries', amount_cents: 55000, period: 'monthly', start_date: '2024-01-01' };
      mockPoolInstance.query.mockResolvedValue({ rows: [mockUpdatedBudget], rowCount: 1 });

      const result = await budgetRepo.updateBudget(1, 'user-123', budgetUpdates, mockLogger);

      expect(mockPoolInstance.query).toHaveBeenCalledWith(
        expect.stringContaining('UPDATE public.budgets SET'),
        [budgetUpdates.name, budgetUpdates.amount_cents, undefined, undefined, 1, 'user-123']
      );
      expect(result).toEqual(mockUpdatedBudget);
    });

    it('should throw an error if no rows are updated', async () => {
      // Arrange: Mock the query to return 0 rows affected
      mockPoolInstance.query.mockResolvedValue({ rows: [], rowCount: 0 });
      
      await expect(budgetRepo.updateBudget(999, 'user-123', { name: 'Fail' }, mockLogger)).rejects.toThrow('Budget not found or user does not have permission to update.');
    });

    it('should throw an error if the database query fails', async () => {
      const dbError = new Error('DB Error');
      mockPoolInstance.query.mockRejectedValue(dbError);
      await expect(budgetRepo.updateBudget(1, 'user-123', { name: 'Fail' }, mockLogger)).rejects.toThrow('Failed to update budget.');
      expect(mockLogger.error).toHaveBeenCalledWith({ err: dbError, budgetId: 1, userId: 'user-123' }, 'Database error in updateBudget');
    });
  });

  describe('deleteBudget', () => {
    it('should execute a DELETE query with user ownership check', async () => {
      mockPoolInstance.query.mockResolvedValue({ rowCount: 1, command: 'DELETE', rows: [] });
      await budgetRepo.deleteBudget(1, 'user-123', mockLogger);
      expect(mockPoolInstance.query).toHaveBeenCalledWith('DELETE FROM public.budgets WHERE budget_id = $1 AND user_id = $2', [1, 'user-123']);
    });

    it('should throw an error if no rows are deleted', async () => {
      // Arrange: Mock the query to return 0 rows affected
      mockPoolInstance.query.mockResolvedValue({ rows: [], rowCount: 0 });
      
      await expect(budgetRepo.deleteBudget(999, 'user-123', mockLogger)).rejects.toThrow('Budget not found or user does not have permission to delete.');
    });

    it('should throw an error if the database query fails', async () => {
      const dbError = new Error('DB Error');
      mockPoolInstance.query.mockRejectedValue(dbError);
      await expect(budgetRepo.deleteBudget(1, 'user-123', mockLogger)).rejects.toThrow('Failed to delete budget.');
      expect(mockLogger.error).toHaveBeenCalledWith({ err: dbError, budgetId: 1, userId: 'user-123' }, 'Database error in deleteBudget');
    });
  });

  describe('getSpendingByCategory', () => {
    it('should call the correct database function and return spending data', async () => {
      const mockSpendingData: SpendingByCategory[] = [{ category_id: 1, category_name: 'Produce', total_spent_cents: 12345 }];
      mockPoolInstance.query.mockResolvedValue({ rows: mockSpendingData });

      const result = await budgetRepo.getSpendingByCategory('user-123', '2024-01-01', '2024-01-31', mockLogger);

      expect(mockPoolInstance.query).toHaveBeenCalledWith('SELECT * FROM public.get_spending_by_category($1, $2, $3)', ['user-123', '2024-01-01', '2024-01-31']);
      expect(result).toEqual(mockSpendingData);
    });

    it('should return an empty array if there is no spending data', async () => {
      mockPoolInstance.query.mockResolvedValue({ rows: [] });
      const result = await budgetRepo.getSpendingByCategory('user-123', '2024-01-01', '2024-01-31', mockLogger);
      expect(result).toEqual([]);
    });

    it('should throw an error if the database query fails', async () => {
      const dbError = new Error('DB Error');
      mockPoolInstance.query.mockRejectedValue(dbError);
      await expect(budgetRepo.getSpendingByCategory('user-123', '2024-01-01', '2024-01-31', mockLogger)).rejects.toThrow('Failed to get spending analysis.');
      expect(mockLogger.error).toHaveBeenCalledWith({ err: dbError, userId: 'user-123', startDate: '2024-01-01', endDate: '2024-01-31' }, 'Database error in getSpendingByCategory');
    });
  });
});