import { describe, it, expect } from 'vitest'
import { solveFlow, balanceSupplies, type FlowNodeInput, type FlowEdgeInput } from '@diagram/flow'

describe('balanceSupplies', () => {
  it('scales demand to match supply', () => {
    const b = balanceSupplies([
      { id: 's', supply: 10 },
      { id: 'a', supply: -5 },
      { id: 'b', supply: -5 }
    ])
    expect(b.get('s')).toBe(10)
    expect((b.get('a') ?? 0) + (b.get('b') ?? 0)).toBeCloseTo(-10)
  })

  it('scales up when demand exceeds supply', () => {
    const b = balanceSupplies([
      { id: 's', supply: 4 },
      { id: 'a', supply: -8 }
    ])
    expect(b.get('a')).toBeCloseTo(-4)
  })
})

describe('solveFlow (tree)', () => {
  // source(10) -> j -> two outlets(5 each)
  const nodes: FlowNodeInput[] = [
    { id: 'src', supply: 10 },
    { id: 'j', supply: 0 },
    { id: 'o1', supply: -5 },
    { id: 'o2', supply: -5 }
  ]
  const edges: FlowEdgeInput[] = [
    { id: 'e1', source: 'src', target: 'j' },
    { id: 'e2', source: 'j', target: 'o1' },
    { id: 'e3', source: 'j', target: 'o2' }
  ]

  it('conserves mass (zero residual on a tree)', () => {
    expect(solveFlow(nodes, edges).residual).toBeCloseTo(0, 6)
  })

  it('carries full supply on the trunk and splits on branches', () => {
    const sol = solveFlow(nodes, edges)
    expect(sol.edges.get('e1')!.magnitude).toBeCloseTo(10)
    expect(sol.edges.get('e2')!.magnitude).toBeCloseTo(5)
    expect(sol.edges.get('e3')!.magnitude).toBeCloseTo(5)
    expect(sol.maxMagnitude).toBeCloseTo(10)
  })

  it('orients flow from source toward sinks', () => {
    const sol = solveFlow(nodes, edges)
    // e1 is src->j, flow should go forward (source to target)
    expect(sol.edges.get('e1')!.forward).toBe(true)
    // e2 is j->o1, forward means j -> outlet
    expect(sol.edges.get('e2')!.forward).toBe(true)
  })

  it('reverses direction when an edge is drawn against the flow', () => {
    const reversed: FlowEdgeInput[] = [{ id: 'e1', source: 'j', target: 'src' }, edges[1], edges[2]]
    const sol = solveFlow(nodes, reversed)
    expect(sol.edges.get('e1')!.forward).toBe(false)
    expect(sol.edges.get('e1')!.magnitude).toBeCloseTo(10)
  })
})

describe('solveFlow (loop)', () => {
  it('still conserves mass with a cycle present', () => {
    const nodes: FlowNodeInput[] = [
      { id: 'src', supply: 6 },
      { id: 'a', supply: 0 },
      { id: 'b', supply: 0 },
      { id: 'sink', supply: -6 }
    ]
    const edges: FlowEdgeInput[] = [
      { id: 'e1', source: 'src', target: 'a' },
      { id: 'e2', source: 'a', target: 'b' },
      { id: 'e3', source: 'b', target: 'sink' },
      { id: 'e4', source: 'a', target: 'sink' } // extra loop edge
    ]
    expect(solveFlow(nodes, edges).residual).toBeCloseTo(0, 6)
  })
})