import { createLexer } from '../src/core/lexer.js';
import { createParser } from '../src/core/parser.js';
import { createInterpreter } from '../src/core/interpreter.js';

function runBabaYaga(code) {
  const lexer = createLexer(code);
  const tokens = lexer.allTokens();
  const parser = createParser(tokens);
  const ast = parser.parse();
  
  const outputs = [];
  const debugOutputs = [];
  
  const host = {
    io: {
      out: (...args) => outputs.push(args.join(' ')),
      debug: (...args) => debugOutputs.push(args.join(' ')),
      in: () => '',
    },
  };
  
  const interpreter = createInterpreter(ast, host);
  const result = interpreter.interpret();
  
  return { outputs, debugOutputs, result };
}

describe('Functional Programming Enhancements', () => {
  
  describe('Scan Operations', () => {
    test('scan with addition function', () => {
      const code = `
        addFunc : acc x -> acc + x;
        numbers : [1, 2, 3, 4, 5];
        result : scan addFunc 0 numbers;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('0,1,3,6,10,15');
    });

    test('cumsum utility function', () => {
      const code = `
        numbers : [1, 2, 3, 4, 5];
        result : cumsum numbers;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('0,1,3,6,10,15');
    });

    test('cumprod utility function', () => {
      const code = `
        numbers : [1, 2, 3, 4];
        result : cumprod numbers;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('1,1,2,6,24');
    });

    test('scan with multiplication function', () => {
      const code = `
        mulFunc : acc x -> acc * x;
        numbers : [2, 3, 4];
        result : scan mulFunc 1 numbers;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('1,2,6,24');
    });
  });

  describe('Array Indexing Operations', () => {
    test('at function selects elements at indices', () => {
      const code = `
        data : [10, 20, 30, 40, 50];
        indices : [0, 2, 4];
        result : at indices data;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('10,30,50');
    });

    test('where function finds matching indices', () => {
      const code = `
        data : [10, 21, 30, 43, 50];
        evenPredicate : x -> x % 2 = 0;
        result : where evenPredicate data;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('0,2,4');
    });

    test('take function gets first n elements', () => {
      const code = `
        data : [1, 2, 3, 4, 5, 6];
        result : take 3 data;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('1,2,3');
    });

    test('drop function removes first n elements', () => {
      const code = `
        data : [1, 2, 3, 4, 5, 6];
        result : drop 3 data;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('4,5,6');
    });

    test('at with empty indices returns empty array', () => {
      const code = `
        data : [1, 2, 3];
        indices : [];
        result : at indices data;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('');
    });

    test('take with zero returns empty array', () => {
      const code = `
        data : [1, 2, 3];
        result : take 0 data;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('');
    });
  });

  describe('Function Combinators', () => {
    test('flip reverses function argument order', () => {
      const code = `
        subtract : x y -> x - y;
        flippedSubtract : flip subtract;
        result : flippedSubtract 3 10;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('7'); // 10 - 3 = 7
    });

    test('apply applies function to value', () => {
      const code = `
        double : x -> x * 2;
        result : apply double 7;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('14');
    });

    test('pipe pipes value through function', () => {
      const code = `
        triple : x -> x * 3;
        result : pipe 4 triple;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('12');
    });

    test('compose creates function composition', () => {
      const code = `
        increment : x -> x + 1;
        double : x -> x * 2;
        composed : compose increment double;
        result : composed 5;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('11'); // increment(double(5)) = increment(10) = 11
    });

    test('combinators work with curried functions', () => {
      const code = `
        add : x -> y -> x + y;
        add5 : add 5;
        flippedAdd5 : flip add5;
        result : flippedAdd5 3;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('8'); // Should still work: 5 + 3 = 8
    });
  });

  describe('Broadcasting Operations', () => {
    test('broadcast applies scalar operation to array', () => {
      const code = `
        addOp : x y -> x + y;
        numbers : [1, 2, 3, 4];
        result : broadcast addOp 10 numbers;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('11,12,13,14');
    });

    test('zipWith applies operation element-wise', () => {
      const code = `
        mulOp : x y -> x * y;
        array1 : [1, 2, 3];
        array2 : [4, 5, 6];
        result : zipWith mulOp array1 array2;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('4,10,18');
    });

    test('zipWith handles arrays of different lengths', () => {
      const code = `
        addOp : x y -> x + y;
        array1 : [1, 2, 3, 4, 5];
        array2 : [10, 20, 30];
        result : zipWith addOp array1 array2;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('11,22,33'); // Only processes minimum length
    });

    test('reshape creates 2D matrix', () => {
      const code = `
        flatArray : [1, 2, 3, 4, 5, 6];
        result : reshape [2, 3] flatArray;
        // Check that result is a 2x3 matrix
        row1 : result.0;
        row2 : result.1;
        io.out row1;
        io.out row2;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('1,2,3'); // First row
      expect(outputs[1]).toBe('4,5,6'); // Second row
    });

    test('broadcast with subtraction', () => {
      const code = `
        subOp : x y -> x - y;
        numbers : [10, 20, 30];
        result : broadcast subOp 5 numbers;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('-5,-15,-25'); // 5 - 10, 5 - 20, 5 - 30
    });
  });

  describe('Monadic Operations', () => {
    test('flatMap flattens mapped results', () => {
      const code = `
        duplicateFunc : x -> [x, x];
        original : [1, 2, 3];
        result : flatMap duplicateFunc original;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('1,1,2,2,3,3');
    });

    test('flatMap with range generation', () => {
      const code = `
        rangeFunc : x -> range 1 x;
        original : [2, 3];
        result : flatMap rangeFunc original;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('1,2,1,2,3');
    });

    test('flatMap with empty results', () => {
      const code = `
        emptyFunc : x -> [];
        original : [1, 2, 3];
        result : flatMap emptyFunc original;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('');
    });

    test('flatMap with mixed result lengths', () => {
      const code = `
        variableFunc : x -> when x is
          1 then [x]
          2 then [x, x]
          _ then [x, x, x];
        original : [1, 2, 3];
        result : flatMap variableFunc original;
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('1,2,2,3,3,3');
    });
  });

  describe('Pattern Guards', () => {
    test('basic pattern guards with numeric conditions', () => {
      const code = `
        classify : x ->
          when x is
            n if (n > 0) then "positive"
            n if (n < 0) then "negative"
            0 then "zero";
        
        result1 : classify 5;
        result2 : classify -3;
        result3 : classify 0;
        io.out result1;
        io.out result2;
        io.out result3;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('positive');
      expect(outputs[1]).toBe('negative');
      expect(outputs[2]).toBe('zero');
    });

    test('pattern guards with range conditions', () => {
      const code = `
        categorizeAge : age ->
          when age is
            a if (a >= 0 and a < 18) then "minor"
            a if (a >= 18 and a < 65) then "adult"
            a if (a >= 65) then "senior"
            _ then "invalid";
        
        result1 : categorizeAge 16;
        result2 : categorizeAge 30;
        result3 : categorizeAge 70;
        io.out result1;
        io.out result2;
        io.out result3;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('minor');
      expect(outputs[1]).toBe('adult');
      expect(outputs[2]).toBe('senior');
    });

    test('pattern guards with complex conditions', () => {
      const code = `
        gradeStudent : score ->
          when score is
            s if (s >= 90) then "A"
            s if (s >= 80 and s < 90) then "B"
            s if (s >= 70 and s < 80) then "C"
            s if (s < 70) then "F"
            _ then "Invalid";
        
        result1 : gradeStudent 95;
        result2 : gradeStudent 85;
        result3 : gradeStudent 75;
        result4 : gradeStudent 65;
        io.out result1;
        io.out result2;
        io.out result3;
        io.out result4;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('A');
      expect(outputs[1]).toBe('B');
      expect(outputs[2]).toBe('C');
      expect(outputs[3]).toBe('F');
    });

    test('pattern guards with wildcard patterns', () => {
      const code = `
        checkRange : x ->
          when x is
            _ if (x >= 1 and x <= 10) then "small"
            _ if (x >= 11 and x <= 100) then "medium"
            _ if (x > 100) then "large"
            _ then "invalid";
        
        result1 : checkRange 5;
        result2 : checkRange 50;
        result3 : checkRange 150;
        result4 : checkRange -5;
        io.out result1;
        io.out result2;
        io.out result3;
        io.out result4;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('small');
      expect(outputs[1]).toBe('medium');
      expect(outputs[2]).toBe('large');
      expect(outputs[3]).toBe('invalid');
    });

    test('pattern guards fail when condition is false', () => {
      const code = `
        testGuard : x ->
          when x is
            n if (n > 10) then "big"
            _ then "small";
        
        result1 : testGuard 15;
        result2 : testGuard 5;
        io.out result1;
        io.out result2;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('big');
      expect(outputs[1]).toBe('small');
    });
  });

  describe('Integration Tests', () => {
    test('combining scan and broadcast operations', () => {
      const code = `
        numbers : [1, 2, 3, 4];
        cumulative : cumsum numbers;
        addTen : broadcast (x y -> x + y) 10 cumulative;
        io.out addTen;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('10,11,13,16,20'); // cumsum [1,2,3,4] = [0,1,3,6,10], then +10 each
    });

    test('combining flatMap with array indexing', () => {
      const code = `
        data : [[1, 2], [3, 4, 5], [6]];
        flattened : flatMap (x -> x) data;
        evens : where (x -> x % 2 = 0) flattened;
        evenValues : at evens flattened;
        io.out evenValues;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('2,4,6');
    });

    test('combining pattern guards with functional operations', () => {
      const code = `
        processNumbers : numbers ->
          with (
            classified : map (n -> when n is
              x if (x > 0) then "pos"
              x if (x < 0) then "neg"
              0 then "zero") numbers;
            positives : filter (n -> n > 0) numbers;
            posSum : reduce (acc x -> acc + x) 0 positives;
          ) ->
            {classifications: classified, sum: posSum};
        
        result : processNumbers [-2, 0, 3, -1, 5];
        io.out result.sum;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('8'); // 3 + 5 = 8
    });

    test('complex pipeline with multiple new features', () => {
      const code = `
        data : [1, 2, 3, 4, 5];
        
        // Use scan to get cumulative sums
        cumSums : cumsum data;
        
        // Use broadcast to multiply by 2
        doubled : broadcast (x y -> x * y) 2 cumSums;
        
        // Use where to find indices of values > 10
        bigIndices : where (x -> x > 10) doubled;
        
        // Use at to get those values
        bigValues : at bigIndices doubled;
        
        io.out bigValues;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('12,20,30'); // Values > 10 from [0,2,6,12,20,30]
    });
  });

  describe('Error Handling', () => {
    test('at throws error for out of bounds index', () => {
      const code = `
        data : [1, 2, 3];
        indices : [0, 5];
        result : at indices data;
      `;
      expect(() => runBabaYaga(code)).toThrow(/Index out of bounds|Can't find variable/);
    });

    test('reshape throws error for incompatible dimensions', () => {
      const code = `
        data : [1, 2, 3, 4, 5];
        result : reshape [2, 3] data;
      `;
      expect(() => runBabaYaga(code)).toThrow('Cannot reshape array');
    });

    test('scan requires function as first argument', () => {
      const code = `
        result : scan 42 0 [1, 2, 3];
      `;
      expect(() => runBabaYaga(code)).toThrow('Scan expects a function');
    });

    test('broadcast requires function as first argument', () => {
      const code = `
        result : broadcast "not a function" 5 [1, 2, 3];
      `;
      expect(() => runBabaYaga(code)).toThrow('broadcast expects a function');
    });

    test('where requires function as first argument', () => {
      const code = `
        result : where "not a function" [1, 2, 3];
      `;
      expect(() => runBabaYaga(code)).toThrow('where expects a function');
    });

    test('flatMap requires function as first argument', () => {
      const code = `
        result : flatMap 42 [1, 2, 3];
      `;
      expect(() => runBabaYaga(code)).toThrow('flatMap expects a function');
    });

    test('take with negative number throws error', () => {
      const code = `
        result : take -1 [1, 2, 3];
      `;
      expect(() => runBabaYaga(code)).toThrow('take expects a non-negative number');
    });

    test('drop with negative number throws error', () => {
      const code = `
        result : drop -1 [1, 2, 3];
      `;
      expect(() => runBabaYaga(code)).toThrow('drop expects a non-negative number');
    });
  });

  describe('Edge Cases', () => {
    test('scan with empty array', () => {
      const code = `
        addFunc : acc x -> acc + x;
        result : scan addFunc 0 [];
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('0'); // Just the initial value
    });

    test('broadcast with empty array', () => {
      const code = `
        addOp : x y -> x + y;
        result : broadcast addOp 5 [];
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe(''); // Empty result
    });

    test('zipWith with empty arrays', () => {
      const code = `
        addOp : x y -> x + y;
        result : zipWith addOp [] [];
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe(''); // Empty result
    });

    test('where with no matches', () => {
      const code = `
        neverTrue : x -> false;
        result : where neverTrue [1, 2, 3];
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe(''); // No matching indices
    });

    test('flatMap with single-element arrays', () => {
      const code = `
        wrapFunc : x -> [x];
        result : flatMap wrapFunc [1, 2, 3];
        io.out result;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('1,2,3'); // Should flatten to original
    });

    test('pattern guards with complex boolean expressions', () => {
      const code = `
        complexTest : x ->
          when x is
            n if ((n > 5) and (n < 15) and (n % 2 = 0)) then "even between 5 and 15"
            n if ((n > 0) or (n < -10)) then "positive or very negative"
            _ then "other";
        
        result1 : complexTest 8;
        result2 : complexTest 3;
        result3 : complexTest -15;
        result4 : complexTest -5;
        io.out result1;
        io.out result2;
        io.out result3;
        io.out result4;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('even between 5 and 15'); // 8 matches first condition
      expect(outputs[1]).toBe('positive or very negative'); // 3 is positive
      expect(outputs[2]).toBe('positive or very negative'); // -15 is very negative
      expect(outputs[3]).toBe('other'); // -5 doesn't match any condition
    });

    test('combinators with identity functions', () => {
      const code = `
        identity : x -> x;
        doubled : x -> x * 2;
        
        // Compose with identity should be equivalent to original function
        composedWithId : compose identity doubled;
        result1 : composedWithId 5;
        
        // Apply identity should return original value
        result2 : apply identity 42;
        
        // Pipe through identity should return original value
        result3 : pipe 7 identity;
        
        io.out result1;
        io.out result2;
        io.out result3;
      `;
      const { outputs } = runBabaYaga(code);
      expect(outputs[0]).toBe('10'); // identity(doubled(5)) = 10
      expect(outputs[1]).toBe('42'); // identity(42) = 42
      expect(outputs[2]).toBe('7');  // pipe 7 identity = 7
    });
  });
});