import { useMemo, useRef } from 'react';
import { useFrame } from '@react-three/fiber';
import * as THREE from 'three';
import type { FlowSegment, PlacedPart, SimResult } from '../types';
import { useBuilder } from '../store/builderStore';
import { useSimulation } from '../simulation/flowSimulator';
import {
  ensureWaterVisual,
  freezeWaterClock,
  getWaterTime,
  getWaterVisual,
  tickWaterClock,
  waterVisuals,
} from '../simulation/waterState';

/**
 * FlowArrows — mount root for all time-based water visual systems:
 *
 *  - WaterDriver (always mounted): advances a global fill wavefront each
 *    frame and writes per-part fill state into the shared water map
 *    (simulation/waterState), which the meshes in PartBody read in their
 *    own useFrame callbacks. No React state is touched per frame.
 *  - Direction arrows (gated by the showFlow toggle): small cones marching
 *    along flowing segments, fading in/out at segment ends and appearing
 *    only once the pipe has actually filled with water.
 */
export function FlowArrows() {
  const showFlow = useBuilder((s) => s.showFlow);
  const parts = useBuilder((s) => s.parts);
  const sim = useSimulation();
  return (
    <group>
      <WaterDriver sim={sim} parts={parts} />
      {showFlow &&
        sim.segments.map((seg, i) => <ArrowStream key={`${seg.partId}-${i}`} seg={seg} />)}
    </group>
  );
}

// ---------------------------------------------------------------------------
// Water fill driver
// ---------------------------------------------------------------------------

// Module-level so remounts / panel toggles don't reset the animation.
let waveFront = 0;

const clamp01 = (x: number) => (x < 0 ? 0 : x > 1 ? 1 : x);
const clamp = (x: number, lo: number, hi: number) => (x < lo ? lo : x > hi ? hi : x);

const GAL_PER_FT3 = 7.48;

/** Approximate internal water volume (gal) — sets how fast a part fills. */
function partVolumeGal(part: PlacedPart, pathLen: number): number {
  const p = part.params;
  const d = Math.max(0.25, p.diameter ?? 1);
  switch (part.type) {
    case 'pipe':
    case 'elbow':
    case 'tee':
    case 'valve':
      // Bore volume: π·r² × path length (d inches → ft).
      return Math.PI * Math.pow(d / 24, 2) * pathLen * GAL_PER_FT3;
    case 'pump':
      return 0.06;
    case 'tower':
      return 0.1 * (p.height ?? 4);
    case 'tray':
      return 0.08 * (p.length ?? 3) * (p.width ?? 1.2);
    case 'wallPanel':
      return 0.25;
    case 'emitter':
      return 0.02;
    case 'drain':
      return 0.05;
    case 'reservoir':
      return 2;
    default:
      return 0.1;
  }
}

interface WaveSpan {
  dist: number;
  len: number;
  /** Wavefront speed through this part (ft/s) ≈ len · GPH / volume. */
  speed: number;
}

function WaterDriver({ sim, parts }: { sim: SimResult; parts: Record<string, PlacedPart> }) {
  const simRef = useRef(sim);
  simRef.current = sim;

  // Per-part fill data, recomputed only when the sim solution changes.
  const wave = useMemo(() => {
    const partIds = Object.keys(sim.fillLength);
    const spans: WaveSpan[] = [];
    for (const pid of partIds) {
      const dist = sim.fillDistance[pid];
      const flow = sim.flows[pid] ?? 0;
      const part = parts[pid];
      if (dist === undefined || flow <= 0 || !part) continue;
      const len = Math.max(0.2, sim.fillLength[pid]);
      const vol = Math.max(0.004, partVolumeGal(part, len));
      spans.push({ dist, len, speed: clamp((len * (flow / 3600)) / vol, 0.5, 5) });
    }
    return { partIds, spans };
  }, [sim, parts]);
  const waveRef = useRef(wave);
  waveRef.current = wave;

  useFrame((_, rawDt) => {
    // Master play/pause: freeze ALL water animation in place.
    if (!useBuilder.getState().simRunning) {
      freezeWaterClock();
      return;
    }
    const s = simRef.current;
    const dt = Math.min(rawDt, 0.1);
    tickWaterClock(dt);

    // Advance the wavefront while pumping; retreat (drain back) when idle.
    // The front moves at the speed of the slowest part it is currently
    // filling (∝ GPH / part volume), so fat tanks fill slower than thin pipes.
    const flowing = s.totalGph > 0.5;
    if (flowing) {
      let speed = Infinity;
      for (const span of waveRef.current.spans) {
        if (waveFront >= span.dist && waveFront < span.dist + span.len) {
          speed = Math.min(speed, span.speed);
        }
      }
      if (!Number.isFinite(speed)) speed = 2; // bridging a junction gap
      waveFront = Math.min(waveFront + speed * dt, s.totalPathLength + 1);
    } else {
      waveFront = Math.max(waveFront - 5 * dt, 0);
    }

    for (const pid of waveRef.current.partIds) ensureWaterVisual(pid);

    for (const [pid, v] of waterVisuals) {
      if (!(pid in s.fillLength)) {
        // Part was deleted — let its water vanish, then drop the entry.
        v.fill = Math.max(0, v.fill - dt * 1.5);
        if (v.fill <= 0) waterVisuals.delete(pid);
        continue;
      }
      const flow = s.flows[pid] ?? 0;
      const backed = s.backedUp.has(pid);
      const dist = s.fillDistance[pid];
      const len = Math.max(0.2, s.fillLength[pid]);

      if (flow > 0 && dist !== undefined) {
        // A part D ft from the pump fills once the wavefront passes D.
        v.fill = clamp01((waveFront - dist) / len);
      } else if (backed) {
        v.fill = Math.min(1, v.fill + dt * 0.6); // stagnant water backing up
      } else {
        v.fill = Math.max(0, v.fill - dt * 0.9); // no supply — drain out
      }

      v.flow = flow;
      v.backedUp = backed;
      v.overflow = false;
      v.spill = false;
      v.entry = s.entryConnector[pid];
      v.state = backed
        ? 'backedUp'
        : v.fill <= 0.002
          ? 'empty'
          : v.fill >= 0.998
            ? 'full'
            : 'filling';
    }

    for (const d of s.drains) {
      if (!d.overflowing) continue;
      const v = waterVisuals.get(d.partId);
      if (v) v.overflow = true;
    }
    for (const pot of s.netPots) {
      if (!pot.overflowing) continue;
      const v = ensureWaterVisual(pot.partId);
      v.fill = 1;
      v.flow = pot.inflowGph;
      v.backedUp = true;
      v.overflow = false;
      v.spill = true;
      v.state = 'backedUp';
    }
  });

  return null;
}

// ---------------------------------------------------------------------------
// Direction arrows
// ---------------------------------------------------------------------------

const UP = new THREE.Vector3(0, 1, 0);

function ArrowStream({ seg }: { seg: FlowSegment }) {
  const groupRef = useRef<THREE.Group>(null);

  const { from, dir, len, quat, count, speed } = useMemo(() => {
    const from = new THREE.Vector3(...seg.from);
    const to = new THREE.Vector3(...seg.to);
    const dir = to.clone().sub(from);
    const len = dir.length();
    dir.normalize();
    // Arrow speed tracks water VELOCITY (∝ GPH / d²), not raw GPH — the same
    // flow squeezed through a narrow pipe visibly rushes.
    const dia = Math.max(0.5, seg.diameter ?? 1);
    return {
      from,
      dir,
      len,
      quat: new THREE.Quaternion().setFromUnitVectors(UP, dir),
      count: Math.max(1, Math.round(len / 0.55)),
      speed: 0.35 + Math.min(2.5, seg.gph / (dia * dia) / 220),
    };
  }, [seg]);

  useFrame(() => {
    const g = groupRef.current;
    if (!g) return;
    const t = getWaterTime() * speed;
    // Arrows only appear once the water has actually reached this part.
    const fill = getWaterVisual(seg.partId)?.fill ?? 1;
    g.children.forEach((child, i) => {
      const f = (((i + t) % count) + count) % count; // 0..count
      const u = f / count;
      child.position.copy(from).addScaledVector(dir, u * len);
      const mat = (child as THREE.Mesh).material as THREE.MeshBasicMaterial;
      // Fade in/out near the segment ends.
      mat.opacity = 0.85 * fill * Math.min(1, Math.min(u, 1 - u) * 4 + 0.15);
    });
  });

  if (len < 0.2) return null;

  return (
    <group ref={groupRef}>
      {Array.from({ length: count }, (_, i) => (
        <mesh key={i} quaternion={quat}>
          <coneGeometry args={[0.05, 0.14, 8]} />
          <meshBasicMaterial color="#38bdf8" toneMapped={false} transparent opacity={0.85} />
        </mesh>
      ))}
    </group>
  );
}