import { describe, it, expect } from 'vitest';
import { parse } from '../../push3-transpiler/src/parser';
import { isValid } from '../mutate';
import {
  generateSeedVariant,
  selectVariants,
  enumerateVariants,
  getTaxThresholds,
  STAKED_THRESHOLDS,
  PENALTY_THRESHOLDS,
  BULL_VARIANTS,
  BEAR_VARIANTS,
  TAX_DISTRIBUTIONS,
  type SeedVariantParams,
  type TaxDistribution,
} from '../seed-generator';

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

function parseAndValidate(src: string): boolean {
  try {
    const program = parse(src);
    return isValid(program);
  } catch {
    return false;
  }
}

const DEFAULT_PARAMS: SeedVariantParams = {
  stakedThreshold: 91,
  penaltyThreshold: 50,
  bull: BULL_VARIANTS[0]!,
  bear: BEAR_VARIANTS[0]!,
  taxDistribution: 'exponential',
};

// ---------------------------------------------------------------------------
// getTaxThresholds
// ---------------------------------------------------------------------------

describe('getTaxThresholds', () => {
  for (const dist of TAX_DISTRIBUTIONS) {
    it(`returns 29 values for ${dist}`, () => {
      const t = getTaxThresholds(dist as TaxDistribution);
      expect(t).toHaveLength(29);
    });

    it(`all ${dist} thresholds are positive`, () => {
      const t = getTaxThresholds(dist as TaxDistribution);
      for (const v of t) {
        expect(v).toBeGreaterThan(0n);
      }
    });

    it(`${dist} thresholds are strictly increasing`, () => {
      const t = getTaxThresholds(dist as TaxDistribution);
      for (let i = 1; i < t.length; i++) {
        expect(t[i]).toBeGreaterThan(t[i - 1]!);
      }
    });

    it(`${dist} thresholds are below 1e18`, () => {
      const t = getTaxThresholds(dist as TaxDistribution);
      const ONE_E18 = 1_000_000_000_000_000_000n;
      for (const v of t) {
        expect(v).toBeLessThan(ONE_E18);
      }
    });
  }

  it('linear thresholds are evenly spaced', () => {
    const t = getTaxThresholds('linear');
    const ONE_E18 = 1_000_000_000_000_000_000n;
    // Each gap should be approximately 1e18/30
    const expected = ONE_E18 / 30n;
    for (let i = 1; i < t.length; i++) {
      const gap = t[i]! - t[i - 1]!;
      // Allow ±1 for integer rounding
      expect(gap - expected).toBeGreaterThanOrEqual(-1n);
      expect(gap - expected).toBeLessThanOrEqual(1n);
    }
  });
});

// ---------------------------------------------------------------------------
// generateSeedVariant — structure and validity
// ---------------------------------------------------------------------------

describe('generateSeedVariant', () => {
  it('generates parseable Push3 with default seed params', () => {
    const src = generateSeedVariant(DEFAULT_PARAMS);
    expect(() => parse(src)).not.toThrow();
  });

  it('generates a valid program with default seed params', () => {
    const src = generateSeedVariant(DEFAULT_PARAMS);
    expect(parseAndValidate(src)).toBe(true);
  });

  it('generates valid programs for all staked thresholds', () => {
    for (const st of STAKED_THRESHOLDS) {
      const src = generateSeedVariant({ ...DEFAULT_PARAMS, stakedThreshold: st });
      expect(parseAndValidate(src)).toBe(true);
    }
  });

  it('generates valid programs for all penalty thresholds', () => {
    for (const pt of PENALTY_THRESHOLDS) {
      const src = generateSeedVariant({ ...DEFAULT_PARAMS, penaltyThreshold: pt });
      expect(parseAndValidate(src)).toBe(true);
    }
  });

  it('generates valid programs for all bull variants', () => {
    for (const bull of BULL_VARIANTS) {
      const src = generateSeedVariant({ ...DEFAULT_PARAMS, bull });
      expect(parseAndValidate(src)).toBe(true);
    }
  });

  it('generates valid programs for all bear variants', () => {
    for (const bear of BEAR_VARIANTS) {
      const src = generateSeedVariant({ ...DEFAULT_PARAMS, bear });
      expect(parseAndValidate(src)).toBe(true);
    }
  });

  it('generates valid programs for all tax distributions', () => {
    for (const taxDistribution of TAX_DISTRIBUTIONS) {
      const src = generateSeedVariant({ ...DEFAULT_PARAMS, taxDistribution });
      expect(parseAndValidate(src)).toBe(true);
    }
  });

  it('default params reproduce optimizer_seed.push3 constants', () => {
    const src = generateSeedVariant(DEFAULT_PARAMS);
    // Staked threshold
    expect(src).toContain('STAKED 91 DYADIC.>');
    // Penalty threshold
    expect(src).toContain('50 DYADIC.<');
    // Bull outputs (DD=1e18, AW=20, AS=1e18, CI=0)
    expect(src).toContain('1000000000000000000 20 1000000000000000000 0');
    // Bear outputs (DD=0.3e18, AW=100, AS=0.3e18, CI=0)
    expect(src).toContain('300000000000000000 100 300000000000000000 0');
  });

  it('different staked thresholds produce distinct programs', () => {
    const programs = STAKED_THRESHOLDS.map(st =>
      generateSeedVariant({ ...DEFAULT_PARAMS, stakedThreshold: st }),
    );
    const unique = new Set(programs);
    expect(unique.size).toBe(STAKED_THRESHOLDS.length);
  });

  it('different tax distributions produce distinct programs', () => {
    const programs = TAX_DISTRIBUTIONS.map(td =>
      generateSeedVariant({ ...DEFAULT_PARAMS, taxDistribution: td }),
    );
    const unique = new Set(programs);
    expect(unique.size).toBe(TAX_DISTRIBUTIONS.length);
  });
});

// ---------------------------------------------------------------------------
// enumerateVariants
// ---------------------------------------------------------------------------

describe('enumerateVariants', () => {
  it('returns 1152 combinations', () => {
    const all = enumerateVariants();
    const expected =
      STAKED_THRESHOLDS.length *
      PENALTY_THRESHOLDS.length *
      BULL_VARIANTS.length *
      BEAR_VARIANTS.length *
      TAX_DISTRIBUTIONS.length;
    expect(all).toHaveLength(expected);
  });

  it('all combinations are distinct', () => {
    const all = enumerateVariants();
    const keys = new Set(
      all.map(v =>
        JSON.stringify([
          v.stakedThreshold,
          v.penaltyThreshold,
          v.bull.anchorWidth,
          v.bear.anchorWidth,
          v.taxDistribution,
        ]),
      ),
    );
    expect(keys.size).toBe(all.length);
  });

  it('covers all staked thresholds', () => {
    const all = enumerateVariants();
    const seen = new Set(all.map(v => v.stakedThreshold));
    for (const st of STAKED_THRESHOLDS) {
      expect(seen.has(st)).toBe(true);
    }
  });

  it('covers all tax distributions', () => {
    const all = enumerateVariants();
    const seen = new Set(all.map(v => v.taxDistribution));
    for (const td of TAX_DISTRIBUTIONS) {
      expect(seen.has(td)).toBe(true);
    }
  });
});

// ---------------------------------------------------------------------------
// selectVariants
// ---------------------------------------------------------------------------

describe('selectVariants', () => {
  it('returns exactly n variants for n < total', () => {
    expect(selectVariants(1)).toHaveLength(1);
    expect(selectVariants(10)).toHaveLength(10);
    expect(selectVariants(50)).toHaveLength(50);
  });

  it('returns all variants when n >= total', () => {
    const total = enumerateVariants().length;
    expect(selectVariants(total)).toHaveLength(total);
    expect(selectVariants(total + 100)).toHaveLength(total);
  });

  it('throws for n < 1', () => {
    expect(() => selectVariants(0)).toThrow();
    expect(() => selectVariants(-1)).toThrow();
  });

  it('all selected variants produce valid programs', () => {
    const variants = selectVariants(20);
    for (const v of variants) {
      const src = generateSeedVariant(v);
      expect(parseAndValidate(src)).toBe(true);
    }
  });

  it('n=6 covers all staked thresholds via even stride', () => {
    // enumerateVariants puts STAKED_THRESHOLDS outermost (192 entries per threshold).
    // stride = 1152/6 = 192 → one representative per staked% value.
    const variants = selectVariants(6);
    expect(variants).toHaveLength(6);
    const stakedValues = new Set(variants.map(v => v.stakedThreshold));
    expect(stakedValues.size).toBe(STAKED_THRESHOLDS.length);
    for (const st of STAKED_THRESHOLDS) {
      expect(stakedValues.has(st)).toBe(true);
    }
    // Each representative is a valid Push3 program
    for (const v of variants) {
      expect(parseAndValidate(generateSeedVariant(v))).toBe(true);
    }
  });

  it('produces diverse programs (no duplicates in typical pop size)', () => {
    const variants = selectVariants(24);
    const programs = variants.map(v => generateSeedVariant(v));
    const unique = new Set(programs);
    expect(unique.size).toBe(24);
  });
});