import * as THREE from "three";
import { OrbitControls } from "three/examples/jsm/controls/OrbitControls.js";
import { RoomEnvironment } from "three/examples/jsm/environments/RoomEnvironment.js";
import { createDemoLoop } from "../demoLoop";
import { observeThreeResize } from "../threeResize";

// A shared 3D iris scene used by every pupil demo in Part 2: a dished,
// procedurally textured iris under a glassy cornea dome, framed by the white
// sclera and lit by a studio environment. The flat two-circle shader from the
// first draft is replaced by real geometry, shading, and reflections in the
// same style as the cornea article.

export interface IrisParams {
	pupilRadius: number; // normalized radius on the iris disc (~0.08 to 0.34)
	pigmentation: number; // 0 brown, 1 hazel, 2 blue, 3 albino
	showMuscles: boolean;
	muscleReveal?: number; // screen-space fraction, 0–1; defaults to the whole iris
}

export interface IrisScene {
	params: IrisParams;
	setOnFrame: (
		cb: ((dt: number, elapsed: number, reducedMotion: boolean) => void) | null,
	) => void;
	setDebug: (cb: (() => Record<string, unknown>) | null) => void;
	invalidate: () => void;
	dispose: () => void;
}

const vertexShader = `
varying vec2 vUv;
varying vec3 vNormal;
varying vec4 vClipPosition;
void main() {
  vUv = uv;
  vNormal = normalize(normalMatrix * normal);
  gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
  vClipPosition = gl_Position;
}
`;

const fragmentShader = `
uniform float uPupilRadius;
uniform int uPigmentation;
uniform bool uShowMuscles;
uniform float uMuscleReveal;

varying vec2 vUv;
varying vec3 vNormal;
varying vec4 vClipPosition;

vec2 hash(vec2 p) {
  p = vec2(dot(p, vec2(127.1, 311.7)), dot(p, vec2(269.5, 183.3)));
  return -1.0 + 2.0 * fract(sin(p) * 43758.5453123);
}
float noise(in vec2 p) {
  const float K1 = 0.366025404, K2 = 0.211324865;
  vec2 i = floor(p + (p.x + p.y) * K1);
  vec2 a = p - i + (i.x + i.y) * K2;
  float m = step(a.y, a.x);
  vec2 o = vec2(m, 1.0 - m);
  vec2 b = a - o + K2;
  vec2 c = a - 1.0 + 2.0 * K2;
  vec3 h = max(0.5 - vec3(dot(a, a), dot(b, b), dot(c, c)), 0.0);
  vec3 n = h * h * h * h * vec3(dot(a, hash(i + 0.0)), dot(b, hash(i + o)), dot(c, hash(i + 1.0)));
  return dot(n, vec3(70.0));
}
float fbm(vec2 uv) {
  float f = 0.0, w = 0.5;
  for (int i = 0; i < 5; i++) { f += w * noise(uv); uv *= 2.0; w *= 0.5; }
  return f;
}

void main() {
  // Nasally and inferiorly shifted pupil center
  vec2 center = vec2(0.52, 0.48);
  float dist = distance(vUv, center);
  float irisRadius = 0.5;
  if (dist > irisRadius) discard;

  float angle = atan(vUv.y - center.y, vUv.x - center.x);
  
  // Static anatomical irregularity; diameter changes drive the deformation.
  float constriction = (1.0 - smoothstep(0.05, 0.35, uPupilRadius));
  float pupilEllipse = (cos(angle*2.0 - 0.5)*0.003) * (constriction + 0.2); // very subtle elliptical stretch
  float pupilIrregularity = (sin(angle*3.0)*0.0015 + cos(angle*7.0)*0.001) * (constriction + 0.1);
  float effectivePupilRadius = uPupilRadius + pupilEllipse + pupilIrregularity;

  if (dist < effectivePupilRadius) {
    if (uPigmentation == 3) {
      float redGlow = 0.16;
      gl_FragColor = vec4(redGlow, 0.02, 0.03, 1.0);
    } else {
      gl_FragColor = vec4(0.006, 0.006, 0.01, 1.0);
    }
    return;
  }

  float radialNorm = (dist - effectivePupilRadius) / (irisRadius - effectivePupilRadius);

  float stretch = mix(2.5, 6.0, clamp((0.35 - uPupilRadius) / 0.25, 0.0, 1.0));
  vec2 noiseUV = vec2(angle * 12.0, radialNorm * stretch);
  float n1 = fbm(noiseUV);
  float n2 = noise(vec2(angle * 30.0, radialNorm * 15.0));
  float fibers = sin(angle * 80.0 + n1 * 4.0) * 0.5 + 0.5;
  float collarette = 0.32 + 0.08 * noise(vec2(angle * 6.0, 0.0));

  vec3 baseColor;
  if (uPigmentation == 0) {
    baseColor = mix(vec3(0.25, 0.12, 0.04), vec3(0.14, 0.06, 0.02), radialNorm) + vec3(0.12, 0.06, 0.02) * n1;
  } else if (uPigmentation == 1) {
    baseColor = mix(vec3(0.52, 0.33, 0.10), vec3(0.18, 0.38, 0.22), smoothstep(0.1, 0.6, radialNorm)) + vec3(0.15, 0.12, 0.05) * n1 * (1.0 - radialNorm);
  } else if (uPigmentation == 2) {
    baseColor = mix(vec3(0.12, 0.28, 0.48), vec3(0.28, 0.52, 0.76), radialNorm) + vec3(0.10, 0.20, 0.30) * n1;
    if (radialNorm < 0.08) baseColor = mix(vec3(0.35, 0.20, 0.08), baseColor, radialNorm / 0.08);
  } else {
    baseColor = mix(vec3(0.70, 0.30, 0.35), vec3(0.55, 0.22, 0.28), radialNorm) + vec3(0.20, 0.08, 0.10) * n1;
  }

  vec3 finalColor = baseColor * (0.75 + 0.35 * fibers) + vec3(n2 * 0.07);
  float crest = 1.0 - smoothstep(0.0, 0.03, abs(radialNorm - collarette));
  finalColor += vec3(0.12, 0.10, 0.06) * crest;

  // Muscle map overlay: red sphincter ring near the pupil, blue radial dilator.
  if (uShowMuscles && (vClipPosition.x / vClipPosition.w * 0.5 + 0.5) < uMuscleReveal) {
    if (dist < effectivePupilRadius + 0.04) {
      float ring = sin(radialNorm * 120.0) * 0.5 + 0.5;
      finalColor = mix(finalColor, vec3(1.0, 0.2, 0.2), 0.55 + 0.25 * ring);
    } else {
      float spoke = sin(angle * 32.0) * 0.5 + 0.5;
      finalColor = mix(finalColor, vec3(0.2, 0.6, 1.0), 0.4 + 0.3 * spoke);
    }
  }

  // Simple shading from the dished geometry so the iris reads as 3D.
  vec3 lightDir = normalize(vec3(-0.4, 0.6, 0.85));
  float diff = 0.7 + 0.35 * max(dot(normalize(vNormal), lightDir), 0.0);
  finalColor *= diff;

  gl_FragColor = vec4(finalColor, 1.0);
}
`;

// Dished iris disc: a full disc that sinks toward the pupil, so orbiting shows
// real depth. The shader discards the pupil and the region beyond the iris.
function dishedIris() {
	const R = 0.95;
	const dish = 0.12;
	const rings = 72;
	const seg = 128;
	const v: number[] = [];
	const uv: number[] = [];
	const idx: number[] = [];
	for (let i = 0; i <= rings; i++) {
		const rr = (i / rings) * R;
		const z = -0.127 + dish * (1.0 - rr / R); // Edge is at -0.127, center bulges forward
		for (let j = 0; j <= seg; j++) {
			const th = (j / seg) * Math.PI * 2;
			const x = rr * Math.cos(th);
			const y = rr * Math.sin(th);
			v.push(x, y, z);
			uv.push((x / R) * 0.5 + 0.5, (y / R) * 0.5 + 0.5);
		}
	}
	for (let i = 0; i < rings; i++) {
		for (let j = 0; j < seg; j++) {
			const a = i * (seg + 1) + j;
			const b = a + 1;
			const c = (i + 1) * (seg + 1) + j;
			const d = c + 1;
			idx.push(a, b, d, a, d, c);
		}
	}
	const g = new THREE.BufferGeometry();
	g.setAttribute("position", new THREE.Float32BufferAttribute(v, 3));
	g.setAttribute("uv", new THREE.Float32BufferAttribute(uv, 2));
	g.setIndex(idx);
	g.computeVertexNormals();
	return g;
}

// Aspheric cornea cap bulging toward the camera (+Z).
function corneaCap(R: number, maxR: number, apexZ: number) {
	const c = 1 / R;
	const Q = -0.26;
	const rings = 64,
		seg = 96;
	const v: number[] = [];
	const idx: number[] = [];
	for (let i = 0; i <= rings; i++) {
		const rr = (i / rings) * maxR;
		const root = 1 - (1 + Q) * c * c * rr * rr;
		const sag = root >= 0 ? (c * rr * rr) / (1 + Math.sqrt(root)) : 0;
		for (let j = 0; j <= seg; j++) {
			const th = (j / seg) * Math.PI * 2;
			v.push(rr * Math.cos(th), rr * Math.sin(th), apexZ - sag);
		}
	}
	for (let i = 0; i < rings; i++) {
		for (let j = 0; j < seg; j++) {
			const a = i * (seg + 1) + j;
			const b = a + 1;
			const cc = (i + 1) * (seg + 1) + j;
			const d = cc + 1;
			// Outward-facing winding for the anterior surface (+Z).
			idx.push(a, d, b, a, cc, d);
		}
	}
	const g = new THREE.BufferGeometry();
	g.setAttribute("position", new THREE.Float32BufferAttribute(v, 3));
	g.setIndex(idx);
	g.computeVertexNormals();
	return g;
}

export function createIrisScene(
	mount: HTMLDivElement,
	initial: IrisParams,
): IrisScene {
	const params: IrisParams = { ...initial };
	let onFrame:
		| ((dt: number, elapsed: number, reducedMotion: boolean) => void)
		| null = null;
	let debugCb: (() => Record<string, unknown>) | null = null;

	const scene = new THREE.Scene();
	const camera = new THREE.PerspectiveCamera(30, 1, 0.1, 100);
	camera.position.set(0.7, 0.5, 5.2);

	const renderer = new THREE.WebGLRenderer({ antialias: true, alpha: true });
	renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
	renderer.setClearColor(0x000000, 0);
	renderer.toneMapping = THREE.ACESFilmicToneMapping;
	renderer.toneMappingExposure = 1.1;
	mount.appendChild(renderer.domElement);

	const pmrem = new THREE.PMREMGenerator(renderer);
	const environment = new RoomEnvironment();
	const environmentTarget = pmrem.fromScene(environment, 0.04);
	environment.dispose();
	scene.environment = environmentTarget.texture;

	const orbit = new OrbitControls(camera, renderer.domElement);
	orbit.enableDamping = true;
	orbit.dampingFactor = 0.05;
	orbit.enablePan = false;
	orbit.minDistance = 3;
	orbit.maxDistance = 9;
	orbit.target.set(0, 0, 0.1);

	scene.add(new THREE.AmbientLight(0xffffff, 0.35));
	const dir = new THREE.DirectionalLight(0xffffff, 1.6);
	dir.position.set(-3, 4, 5);
	scene.add(dir);

	// Sclera ring removed at user request.

	// Iris.
	const irisGeo = dishedIris();
	const irisMat = new THREE.ShaderMaterial({
		vertexShader,
		fragmentShader,
		uniforms: {
			uPupilRadius: { value: params.pupilRadius },
			uPigmentation: { value: params.pigmentation },
			uShowMuscles: { value: params.showMuscles },
			uMuscleReveal: { value: params.muscleReveal ?? 1 },
		},
		side: THREE.DoubleSide,
	});
	const iris = new THREE.Mesh(irisGeo, irisMat);
	scene.add(iris);

	// Glassy cornea dome over the iris.
	const corneaGeo = corneaCap(1.6, 1.05, 0.25);
	const corneaMat = new THREE.MeshPhysicalMaterial({
		transmission: 1.0,
		thickness: 0.15,
		ior: 1.376,
		roughness: 0.03,
		metalness: 0.0,
		clearcoat: 1.0,
		clearcoatRoughness: 0.03,
		envMapIntensity: 1.3,
		attenuationColor: new THREE.Color(0xeaf4ff),
		attenuationDistance: 6,
		side: THREE.FrontSide,
	});
	const cornea = new THREE.Mesh(corneaGeo, corneaMat);
	scene.add(cornea);

	const resizeObserver = observeThreeResize(mount, renderer, camera, () =>
		loop.invalidate(),
	);

	const loop = createDemoLoop(mount, {
		controls: orbit,
		// Keep the transmissive canvas refreshed while visible; visibility pauses it.
		continuous: () => true,
		draw: (dt, elapsed, reducedMotion) => {
			onFrame?.(dt, elapsed, reducedMotion);
			irisMat.uniforms.uPupilRadius!.value = params.pupilRadius;
			irisMat.uniforms.uPigmentation!.value = params.pigmentation;
			irisMat.uniforms.uShowMuscles!.value = params.showMuscles;
			irisMat.uniforms.uMuscleReveal!.value = params.muscleReveal ?? 1;
			renderer.render(scene, camera);
			return false;
		},
		inspect: () => ({
			model: "Illustrative iris",
			pupilDiameterMm: +(params.pupilRadius * 24).toFixed(3),
			pigmentation: params.pigmentation,
			muscleMap: params.showMuscles,
			...debugCb?.(),
			camera: camera.position.toArray(),
			target: orbit.target.toArray(),
		}),
	});

	return {
		params,
		invalidate: loop.invalidate,
		setOnFrame: (cb) => {
			onFrame = cb;
			loop.invalidate();
		},
		setDebug: (cb) => {
			debugCb = cb;
			loop.inspect();
			loop.invalidate();
		},
		dispose: () => {
			loop.dispose();
			resizeObserver.disconnect();
			orbit.dispose();
			if (renderer.domElement.parentNode)
				renderer.domElement.parentNode.removeChild(renderer.domElement);
			irisGeo.dispose();
			irisMat.dispose();
			corneaGeo.dispose();
			corneaMat.dispose();
			environmentTarget.dispose();
			pmrem.dispose();
			renderer.dispose();
		},
	};
}