import { useMemo, useRef } from 'react';
import { useFrame } from '@react-three/fiber';
import * as THREE from 'three';
import type { RocketConfig, SimulationState } from './types';
import { engineState, separationAge, separationTransform, stageGeometry } from './sceneMath';

const DARK = '#20262d';
const METAL = '#8b96a1';

/**
 * Animated exhaust plume. The whole plume is anchored at the engine bell (local
 * y = 0) and grows DOWNWARD; we scale an anchored group (not a center-origin
 * mesh) so the flame root always stays glued to the nozzle regardless of
 * throttle/flicker. `nozzleY` is where the engine bell sits in the parent.
 */
function EnginePlume({
  radius,
  length,
  throttle,
  active,
  running,
  nozzleY = 0,
}: {
  radius: number;
  length: number;
  throttle: number;
  active: boolean;
  running: boolean;
  nozzleY?: number;
}) {
  const anchorRef = useRef<THREE.Group>(null);
  const lightRef = useRef<THREE.PointLight>(null);

  useFrame(({ clock }) => {
    if (!anchorRef.current || !running) return;
    const flicker = 0.85 + Math.sin(clock.elapsedTime * 45) * 0.12 + Math.random() * 0.06;
    const stretch = flicker * (0.5 + throttle * 0.7);
    // Only Y is animated; the group origin stays at the nozzle so the flame
    // never detaches from the tail.
    anchorRef.current.scale.set(0.75 + throttle * 0.35, stretch, 0.75 + throttle * 0.35);
    if (lightRef.current) lightRef.current.intensity = 5 + flicker * 3;
  });

  if (!active) return null;
  // Cones are built with their tip up and base down, positioned so the base
  // (top of cone) sits exactly at the nozzle and the tip points away.
  return (
    <group ref={anchorRef} position={[0, nozzleY, 0]}>
      <mesh position={[0, -length * 0.9, 0]} rotation={[Math.PI, 0, 0]}>
        <coneGeometry args={[radius * 2.1, length * 1.8, 20, 1, true]} />
        <meshBasicMaterial color="#ff7a2a" transparent opacity={0.32} blending={THREE.AdditiveBlending} depthWrite={false} side={THREE.DoubleSide} />
      </mesh>
      <mesh position={[0, -length * 0.5, 0]} rotation={[Math.PI, 0, 0]}>
        <coneGeometry args={[radius * 1.15, length, 20, 1, true]} />
        <meshBasicMaterial color="#fff2c2" transparent opacity={0.9} blending={THREE.AdditiveBlending} depthWrite={false} side={THREE.DoubleSide} />
      </mesh>
      <pointLight ref={lightRef} color="#ff9a3c" intensity={6} distance={radius * 30} decay={2} position={[0, -length * 0.4, 0]} />
    </group>
  );
}

/** Nozzle cluster drawn as a dark boat-tail plus a few nozzle cones. */
function EngineCluster({ radius, count }: { radius: number; count: number }) {
  const nozzles = Math.min(count, 9);
  const ring = useMemo(() => {
    if (nozzles <= 1) return [[0, 0]] as Array<[number, number]>;
    const points: Array<[number, number]> = [[0, 0]];
    const outer = nozzles - 1;
    for (let i = 0; i < outer; i += 1) {
      const angle = (i / outer) * Math.PI * 2;
      points.push([Math.cos(angle) * radius * 0.5, Math.sin(angle) * radius * 0.5]);
    }
    return points;
  }, [nozzles, radius]);

  return (
    <group>
      <mesh position={[0, radius * 0.4, 0]}>
        <cylinderGeometry args={[radius * 0.82, radius, radius * 0.8, 24]} />
        <meshStandardMaterial color={DARK} metalness={0.6} roughness={0.5} />
      </mesh>
      {ring.map(([x, z], i) => (
        <mesh key={i} position={[x, -radius * 0.2, z]} rotation={[Math.PI, 0, 0]}>
          <coneGeometry args={[radius * 0.16, radius * 0.55, 12, 1, true]} />
          <meshStandardMaterial color="#3a4149" metalness={0.7} roughness={0.35} side={THREE.DoubleSide} />
        </mesh>
      ))}
    </group>
  );
}

/** Grid-fin / stabiliser fins around the base of the first stage. */
function Fins({ radius, height }: { radius: number; height: number }) {
  return (
    <group>
      {[0, 1, 2, 3].map((i) => (
        <mesh key={i} position={[0, height * 0.5, 0]} rotation={[0, (i / 4) * Math.PI * 2, 0]}>
          <boxGeometry args={[radius * 0.12, height, radius * 2.4]} />
          <meshStandardMaterial color={DARK} metalness={0.5} roughness={0.6} />
        </mesh>
      ))}
    </group>
  );
}

interface StageProps {
  radius: number;
  height: number;
  color: string;
}

/** First-stage airframe with fins and engine cluster; base sits at local y=0. */
function FirstStageBody({ radius, height, color, engineCount }: StageProps & { engineCount: number }) {
  return (
    <group>
      <mesh position={[0, height * 0.5, 0]}>
        <cylinderGeometry args={[radius, radius, height, 32]} />
        <meshStandardMaterial color={color} metalness={0.35} roughness={0.45} />
      </mesh>
      {/* Livery band near the top of the stage */}
      <mesh position={[0, height * 0.9, 0]}>
        <cylinderGeometry args={[radius * 1.005, radius * 1.005, height * 0.05, 32]} />
        <meshStandardMaterial color={DARK} metalness={0.4} roughness={0.5} />
      </mesh>
      <group position={[0, height * 0.16, 0]}>
        <Fins radius={radius} height={height * 0.22} />
      </group>
      <EngineCluster radius={radius} count={engineCount} />
    </group>
  );
}

/** Interstage + second stage + payload fairing/nose; base sits at local y=0. */
function UpperStack({
  radius,
  interstageHeight,
  stage2Height,
  noseHeight,
  color,
  engineCount,
}: {
  radius: number;
  interstageHeight: number;
  stage2Height: number;
  noseHeight: number;
  color: string;
  engineCount: number;
}) {
  const stage2Base = interstageHeight;
  const noseBase = interstageHeight + stage2Height;
  return (
    <group>
      {/* Interstage (dark) */}
      <mesh position={[0, interstageHeight * 0.5, 0]}>
        <cylinderGeometry args={[radius * 0.98, radius, interstageHeight, 32]} />
        <meshStandardMaterial color={DARK} metalness={0.5} roughness={0.5} />
      </mesh>
      {/* Second-stage engine bell tucked under the interstage */}
      <group position={[0, stage2Base, 0]}>
        <EngineCluster radius={radius * 0.7} count={Math.min(engineCount, 3)} />
      </group>
      {/* Second-stage body */}
      <mesh position={[0, stage2Base + stage2Height * 0.5, 0]}>
        <cylinderGeometry args={[radius * 0.96, radius * 0.98, stage2Height, 32]} />
        <meshStandardMaterial color={color} metalness={0.35} roughness={0.45} />
      </mesh>
      {/* Payload fairing / nose cone */}
      <mesh position={[0, noseBase + noseHeight * 0.5, 0]}>
        <coneGeometry args={[radius * 0.96, noseHeight, 32]} />
        <meshStandardMaterial color={METAL} metalness={0.4} roughness={0.4} />
      </mesh>
    </group>
  );
}

function PayloadSatellite({ radius, panelProgress }: { radius: number; panelProgress: number }) {
  const satelliteRef = useRef<THREE.Group>(null);
  const easedProgress = 1 - Math.pow(1 - panelProgress, 3);

  useFrame((_, delta) => {
    if (satelliteRef.current) satelliteRef.current.rotation.y += delta * 0.18;
  });

  return (
    <group ref={satelliteRef}>
      <mesh>
        <boxGeometry args={[radius * 1.35, radius * 1.15, radius * 1.2]} />
        <meshStandardMaterial color="#d2a84f" metalness={0.65} roughness={0.38} />
      </mesh>
      <mesh position={[0, radius * 0.76, 0]}>
        <cylinderGeometry args={[radius * 0.28, radius * 0.42, radius * 0.42, 20]} />
        <meshStandardMaterial color="#d9ddda" metalness={0.65} roughness={0.28} />
      </mesh>
      <mesh position={[0, radius * 1.08, 0]} rotation={[Math.PI / 2, 0, 0]}>
        <sphereGeometry args={[radius * 0.54, 24, 12, 0, Math.PI * 2, 0, Math.PI / 2]} />
        <meshStandardMaterial color="#eff2ef" metalness={0.35} roughness={0.4} side={THREE.BackSide} />
      </mesh>
      {[-1, 1].map((direction) => (
        <group key={direction} position={[direction * radius * 0.78, 0, 0]}>
          <mesh
            position={[direction * radius * 1.55 * easedProgress, 0, 0]}
            scale={[Math.max(0.04, easedProgress), 1, 1]}
          >
            <boxGeometry args={[radius * 3.1, radius * 0.72, radius * 0.08]} />
            <meshStandardMaterial color="#24507a" metalness={0.25} roughness={0.55} />
          </mesh>
          <mesh position={[direction * radius * 0.08, 0, 0]}>
            <boxGeometry args={[radius * 0.18, radius * 0.18, radius * 0.16]} />
            <meshStandardMaterial color={DARK} metalness={0.7} roughness={0.35} />
          </mesh>
        </group>
      ))}
    </group>
  );
}

interface RocketModelProps {
  rocket: RocketConfig;
  state: SimulationState;
}

/**
 * Full launch vehicle. Each discarded stage drifts away from the active stack;
 * after orbital insertion, the payload separates and deploys its solar arrays.
 */
export function RocketModel({ rocket, state }: RocketModelProps) {
  const geo = useMemo(() => stageGeometry(rocket), [rocket]);
  const engines = engineState(state);
  const age = separationAge(state);
  const sep = separationTransform(age);
  const deployEvent = state.events.find((event) => event.id === 'deploy');
  const deployAge = deployEvent ? Math.max(0, state.time - deployEvent.time) : null;
  const deployProgress = deployAge === null ? 0 : Math.min(1, deployAge / 2.4);

  const firstStageRef = useRef<THREE.Group>(null);
  const upperStackRef = useRef<THREE.Group>(null);

  // The upper stack stays anchored until payload deployment. The first stage
  // sits directly below it (base at local y=0) until stage separation.
  const upperBaseY = geo.stage1Height;
  const payloadY = upperBaseY + geo.interstageHeight + geo.stage2Height + geo.noseHeight * 0.58;

  useFrame(() => {
    if (firstStageRef.current) {
      if (sep) {
        firstStageRef.current.position.set(sep.drift, -sep.drop, 0);
        firstStageRef.current.rotation.set(sep.tumble * 0.6, sep.tumble * 0.3, sep.tumble);
        firstStageRef.current.visible = sep.opacity > 0.02;
      } else {
        firstStageRef.current.position.set(0, 0, 0);
        firstStageRef.current.rotation.set(0, 0, 0);
        firstStageRef.current.visible = true;
      }
    }

    if (upperStackRef.current) {
      if (deployAge === null) {
        upperStackRef.current.position.set(0, upperBaseY, 0);
        upperStackRef.current.rotation.set(0, 0, 0);
        upperStackRef.current.visible = true;
      } else {
        upperStackRef.current.position.set(-deployAge * 0.72, upperBaseY - deployAge * 0.82, deployAge * 0.18);
        upperStackRef.current.rotation.set(deployAge * 0.16, deployAge * 0.08, -deployAge * 0.3);
        upperStackRef.current.visible = deployAge < 5.2;
      }
    }
  });

  return (
    <group>
      {/* Upper stack: active through insertion, then discarded after deployment. */}
      <group ref={upperStackRef} position={[0, upperBaseY, 0]}>
        <UpperStack
          radius={geo.radius}
          interstageHeight={geo.interstageHeight}
          stage2Height={geo.stage2Height}
          noseHeight={geo.noseHeight}
          color={rocket.color}
          engineCount={rocket.stage_2.engine_count}
        />
        {/* Stage-2 plume: fires from the second-stage bell (above interstage) */}
        <group position={[0, geo.interstageHeight, 0]}>
          <EnginePlume
            radius={geo.radius * 0.7}
            length={geo.radius * 7}
            nozzleY={-geo.radius * 0.7 * 0.4}
            throttle={state.throttle}
            active={engines.stage2}
            running={state.isRunning}
          />
        </group>
      </group>

      {deployAge !== null && (
        <group position={[deployProgress * 1.8, payloadY + deployProgress * 2.6, 0]} scale={1.25}>
          <PayloadSatellite radius={geo.radius} panelProgress={deployProgress} />
        </group>
      )}

      {/* First stage: attached below the stack, or detached and falling */}
      <group ref={firstStageRef}>
        <FirstStageBody
          radius={geo.radius}
          height={geo.stage1Height}
          color={rocket.color}
          engineCount={rocket.stage_1.engine_count}
        />
        {/* Stage-1 plume: fires from the engine cluster at the base (y≈0) */}
        <EnginePlume
          radius={geo.radius}
          length={geo.radius * 10}
          nozzleY={-geo.radius * 0.4}
          throttle={state.throttle}
          active={engines.stage1}
          running={state.isRunning}
        />
      </group>
    </group>
  );
}