import fs from 'node:fs/promises';
import path from 'node:path';
import {fileURLToPath} from 'node:url';
import {NodeIO, Document, getBounds} from '@gltf-transform/core';
import {ALL_EXTENSIONS} from '@gltf-transform/extensions';
import {prune, dedup, flatten, transformMesh, cloneDocument, mergeDocuments, simplify, compactPrimitive, unpartition} from '@gltf-transform/functions';
import {MeshoptSimplifier} from 'meshoptimizer';
import sharp from 'sharp';

const root=path.resolve(path.dirname(fileURLToPath(import.meta.url)),'../../Content/Vegitations');
const manifest=JSON.parse(await fs.readFile(path.join(root,'sources.json'),'utf8'));
const ids=[...new Set(manifest.files.filter(f=>f.path.endsWith('/source/model.gltf')).map(f=>f.path.split('/')[0]))];
const io=new NodeIO().registerExtensions(ALL_EXTENSIONS);
await MeshoptSimplifier.ready;
for(const id of ids) {
  const selected=process.argv.slice(2).filter(a=>!a.startsWith('--'));
  if(selected.length&&!selected.includes(id))continue;
  const out=path.join(root,id,'prepared');await fs.mkdir(out,{recursive:true});
  const cards=JSON.parse(await fs.readFile(path.join(out,'cards.json'),'utf8').catch(()=>'null'));
  const sourceOnly=process.argv.includes('--source-only');
  const stamp=JSON.stringify({version:5,cards:!!cards,files:manifest.files.filter(f=>f.path.startsWith(id+'/')).map(f=>f.md5)});
  if(sourceOnly&&await fs.stat(path.join(root,id,'source/bake-source.glb')).catch(()=>null))continue;
  if(!sourceOnly&&await fs.readFile(path.join(out,'stamp.json'),'utf8').catch(()=>null)===stamp){console.log('Cached',id);continue;}
  console.log('Preparing',id);
  const doc=await io.read(path.join(root,id,'source/model.gltf'));
  await doc.transform(flatten());
  const scene=doc.getRoot().getDefaultScene()??doc.getRoot().listScenes()[0];
  const all=scene.listChildren();
  const chosen=all.length<=3?all:[all[0],all[Math.floor(all.length/3)],all[Math.floor(all.length*2/3)]];
  for(const node of all)if(!chosen.includes(node))node.dispose();
  // Apply source transforms and give each variant a gravity-aligned, bottom-centred origin.
  for(let i=0;i<chosen.length;i++) {
    const node=chosen[i],mesh=node.getMesh();
    if(!mesh)throw new Error('Expected a static mesh at root: '+id);
    transformMesh(mesh,node.getWorldMatrix());node.setMatrix([1,0,0,0,0,1,0,0,0,0,1,0,0,0,0,1]);
    const b=getBounds(node),offset=[-(b.min[0]+b.max[0])*.5,-b.min[1],-(b.min[2]+b.max[2])*.5];
    transformMesh(mesh,[1,0,0,0,0,1,0,0,0,0,1,0,...offset,1]);
    node.setName('variant_'+i);
  }
  if(sourceOnly) {
    await doc.transform(unpartition());await io.write(path.join(root,id,'source/bake-source.glb'),doc);
    console.log('Prepared full-detail baking source',id);continue;
  }
  // Dense geometric foliage is baked in small spatial clusters, preserving its volume
  // instead of reducing millions of individual needles/leaves to isolated triangles.
  if(cards)for(let variant=0;variant<chosen.length;variant++) {
    const mesh=chosen[variant].getMesh();
    for(const p of [...mesh.listPrimitives()])if(/twig|leaves/.test(p.getMaterial().getName().toLowerCase())){mesh.removePrimitive(p);p.dispose();}
    const color=doc.createTexture('canopy').setImage(await fs.readFile(path.join(out,`cards-${variant}-color.png`))).setMimeType('image/png');
    const normal=doc.createTexture('canopy_normal').setImage(await fs.readFile(path.join(out,`cards-${variant}-normal.png`))).setMimeType('image/png');
    const material=doc.createMaterial('baked_canopy').setBaseColorTexture(color).setNormalTexture(normal)
      .setMetallicFactor(0).setRoughnessFactor(.9).setAlphaMode('MASK').setAlphaCutoff(.4).setDoubleSided(true);
    const pos=[],norm=[],uv=[],indices=[];
    for(const card of cards.Variants[variant]) {
      const index=pos.length/3;
      for(const [x,y] of [[-.5,-.5],[.5,-.5],[.5,.5],[-.5,.5]]) {
        pos.push(...card.Center.map((v,k)=>v+(card.Right[k]*x+card.Up[k]*y)*card.Size));
        norm.push(...card.Normal);
        uv.push(card.Uv[0]+(x+.5)*card.Uv[2],card.Uv[1]+(.5-y)*card.Uv[3]);
      }
      indices.push(index,index+1,index+2,index,index+2,index+3);
    }
    const buffer=doc.getRoot().listBuffers()[0];
    const accessor=(type,array)=>doc.createAccessor().setType(type).setArray(array).setBuffer(buffer);
    mesh.addPrimitive(doc.createPrimitive().setAttribute('POSITION',accessor('VEC3',new Float32Array(pos)))
      .setAttribute('NORMAL',accessor('VEC3',new Float32Array(norm))).setAttribute('TEXCOORD_0',accessor('VEC2',new Float32Array(uv)))
      .setIndices(accessor('SCALAR',new Uint32Array(indices))).setMaterial(material));
  }
  await doc.transform(prune());
  const masks=manifest.files.filter(f=>f.path.startsWith(id+'/source/masks/')&&f.path.toLowerCase().includes('alpha'));
  for(const material of doc.getRoot().listMaterials()) {
    const name=material.getName().toLowerCase();
    if(name==='baked_canopy')continue;
    const alpha=masks.find(f=>path.basename(f.path).toLowerCase()==='alpha.png'||name.includes(path.basename(f.path).toLowerCase().replace('_alpha.png','')));
    const foliage=material.getAlphaMode()!=='OPAQUE'||/twig|leaves|grass|fern|shrub|weed/.test(name);
    material.setMetallicFactor(0).setDoubleSided(foliage);
    if(foliage&&alpha&&material.getBaseColorTexture()) {
      const texture=material.getBaseColorTexture();
      const info=await sharp(texture.getImage()).metadata();
      const a=await sharp(await fs.readFile(path.join(root,alpha.path))).resize(info.width,info.height).extractChannel(0).toBuffer();
      const rgba=await sharp(texture.getImage()).removeAlpha().joinChannel(a).png().toBuffer();
      const replacement=doc.createTexture(name+'_alpha').setImage(rgba).setMimeType('image/png');
      material.setBaseColorTexture(replacement).setAlphaMode('MASK').setAlphaCutoff(.4);
    }
  }
  const tree=/tree|sapling/.test(id),budgets=tree?[12000,2000,350]:/shrub/.test(id)?[12000,3000,600]:[2500,800,200];
  const merged=new Document();const stats=[];
  for(let lod=0;lod<3;lod++) {
    const copy=cloneDocument(doc),sc=copy.getRoot().listScenes()[0];sc.setName('LOD'+lod);
    for(const node of sc.listChildren()) {
      const primitives=node.getMesh().listPrimitives();
      const total=primitives.reduce((n,p)=>n+(p.getIndices()?.getCount()??p.getAttribute('POSITION').getCount())/3,0);
      const canopy=primitives.filter(p=>p.getMaterial().getName()==='baked_canopy').reduce((n,p)=>n+p.getIndices().getCount()/3,0);
      for(const p of primitives) {
        const count=p.getIndices()?.getCount()??p.getAttribute('POSITION').getCount();
        const target=p.getMaterial().getName()==='baked_canopy'?count:Math.max(6,Math.floor((budgets[lod]-canopy)*count/Math.max(1,total-canopy)/3)*3);
        if(count<=target)continue;
        // The library performs topology/attribute-preserving reduction first.
        // simplify() acts on the whole document below; retain the target per primitive for a bounded fallback.
        p.setExtras({...p.getExtras(),target});
      }
    }
    let original=0;for(const m of copy.getRoot().listMeshes())for(const p of m.listPrimitives())original+=p.getIndices().getCount()/3;
    const retained=[];
    for(const mesh of copy.getRoot().listMeshes())for(const p of [...mesh.listPrimitives()])
      if(p.getMaterial().getName()==='baked_canopy'){mesh.removePrimitive(p);retained.push([mesh,p]);}
    await copy.transform(simplify({simplifier:MeshoptSimplifier,ratio:Math.min(1,budgets[lod]*chosen.length/original),error:.01*(lod+1),lockBorder:false}));
    for(const [mesh,p] of retained)mesh.addPrimitive(p);
    for(const m of copy.getRoot().listMeshes())for(const p of m.listPrimitives()) {
      const indices=new Uint32Array(p.getIndices().getArray()),positions=new Float32Array(p.getAttribute('POSITION').getArray());
      const target=p.getExtras().target;
      if(target&&indices.length>target*1.1) {
        let [reduced]=MeshoptSimplifier.simplify(indices,positions,3,target,1,['Permissive','Prune']);
        if(reduced.length>target*1.1)[reduced]=MeshoptSimplifier.simplifySloppy(indices,positions,3,null,target,1);
        p.setIndices(copy.createAccessor().setType('SCALAR').setArray(reduced).setBuffer(copy.getRoot().listBuffers()[0]));
        compactPrimitive(p);
      }
      p.setExtras({});
    }
    await copy.transform(prune());
    let triangles=0;for(const m of copy.getRoot().listMeshes())for(const p of m.listPrimitives())triangles+=p.getIndices().getCount()/3;
    stats.push(triangles);mergeDocuments(merged,copy);
  }
  merged.getRoot().setDefaultScene(merged.getRoot().listScenes()[0]);
  await merged.transform(dedup(),unpartition(),prune());
  for(const ext of merged.getRoot().listExtensionsUsed())ext.dispose();
  await io.write(path.join(out,'model.glb'),merged);
  const dryFile=manifest.files.find(f=>f.path.startsWith(id+'/source/masks/dry_diff'));
  if(dryFile) {
    const dry=cloneDocument(merged);
    const mask=masks[0];
    const image=await sharp(await fs.readFile(path.join(root,dryFile.path))).metadata();
    const a=await sharp(await fs.readFile(path.join(root,mask.path))).resize(image.width,image.height).extractChannel(0).toBuffer();
    const rgba=await sharp(await fs.readFile(path.join(root,dryFile.path))).removeAlpha().joinChannel(a).png().toBuffer();
    const texture=dry.createTexture('dry_grass').setImage(rgba).setMimeType('image/png');
    for(const m of dry.getRoot().listMaterials())m.setBaseColorTexture(texture);
    await dry.transform(prune());await io.write(path.join(out,'dry.glb'),dry);
  }
  const variantBounds=chosen.map(n=>getBounds(n));
  await fs.writeFile(path.join(out,'preparation.json'),JSON.stringify({source:'https://polyhaven.com/a/'+id,license:'CC0',budgets,triangles:stats,variants:variantBounds},null,2));
  await fs.writeFile(path.join(out,'stamp.json'),stamp);
  console.log('Prepared',id,stats);
}