#include "FluidSim/Machine.hpp"
#include "MachineTestControls.hpp"
#include <algorithm>
#include <chrono>
#include <cmath>
#include <fstream>
#include <iostream>
#include <stdexcept>
using namespace fluidsim::port;
namespace {
void check(bool value,const char* why){if(!value)throw std::runtime_error(why);}
bool close(double a,double b,double tolerance=1e-5){return std::abs(a-b)<=tolerance*std::max({1.,std::abs(a),std::abs(b)});}
template<class F>void rejects(F&& f,const char* why){bool caught=false;try{f();}catch(const std::exception&){caught=true;}check(caught,why);}
void run(MachineSimulation& s,int samples){for(int i=0;i<samples;++i){const auto audio=s.process();check(std::isfinite(audio.left)&&std::isfinite(audio.right),"Mechanical machine generated invalid audio");}}
const ShaftState& shaft(const MachineSnapshot& state,std::uint32_t id){for(const auto& s:state.shafts)if(s.id==id)return s;throw std::runtime_error("Missing shaft snapshot");}
void balance(const ShaftState& s){check(close(s.initial_energy+s.external_work+s.drive_work,s.kinetic_energy+s.friction_loss+s.brake_loss,1e-8),"Compiled machine lost rotational work");}
std::uint32_t cylinder(const MachineDocument& d){return std::find_if(d.parts.begin(),d.parts.end(),[](const auto& p){return p.kind==MachinePartKind::Cylinder;})->id;}
}
int main()try{
    MachineDocument bench;bench.name="Independent mechanical bench";
    const auto a=bench.add(MachinePartKind::Shaft,{}).id,b=bench.add(MachinePartKind::Shaft,{0,0,.5f}).id;
    bench.find(a)->shaft.inertia=2;bench.find(b)->shaft.inertia=4;
    const auto wheel=bench.add(MachinePartKind::Flywheel,{.3f,0,0}).id;bench.find(wheel)->shaft.inertia=1;bench.connect({wheel,0},{a,0});
    const auto motor=bench.add(MachinePartKind::DriveMotor,{-.3f,0,0}).id;bench.find(motor)->drive={0,6,false,true};bench.connect({a,0},{motor,0});
    const auto reverse=bench.add(MachinePartKind::DriveMotor,{-.3f,0,.5f}).id;bench.find(reverse)->drive={0,-4,false,true};bench.connect({b,0},{reverse,0});
    const auto brake=bench.add(MachinePartKind::Brake,{.3f,0,.2f}).id;bench.find(brake)->brake.constant=1;bench.connect({a,0},{brake,0});
    bench.probes={{1,a,"First RPM",MachineQuantity::Rpm},{2,b,"Second RPM",MachineQuantity::Rpm},{3,motor,"Motor torque",MachineQuantity::Torque},{4,brake,"Absorbed power",MachineQuantity::Power}};
    validate_machine(bench);check(machine_part_volume(*bench.find(a))==0&&machine_part_volume(*bench.find(motor))==0,"Mechanical parts added gas storage");
    check(!machine_connection_preview(bench,{a,0},{b,0}).allowed,"Direct shaft-to-shaft rigid constraint was accepted without a coupling");
    check(!machine_connection_preview(bench,{b,0},{wheel,0}).allowed,"Flywheel was attached to two independent shafts");
    for(int rate:{1000,4000}){MachineSimulation sim(bench,rate);run(sim,rate);const auto state=sim.snapshot();check(state.cells==0&&state.shafts.size()==2,"Pure mechanical machine required a hidden gas network");
        check(close(shaft(state,a).speed,5./3)&&close(shaft(state,b).speed,-1),"Compiled shafts share state or use the wrong effective inertia");
        check(close(shaft(state,a).inertia,3),"Compiled flywheel inertia was counted incorrectly");const auto values=sim.measure();check(close(values.values[0],5./3*60/(2*pi))&&close(values.values[1],-60/(2*pi))&&values.values[2]==6&&values.values[3]>=0,"Rotational probes lost their named targets or sign conventions");balance(shaft(state,a));balance(shaft(state,b));
        auto tuned=bench;tuned.find(motor)->drive.torque=9;const auto time=state.time;sim.apply_tuning(tuned);check(sim.snapshot().time==time,"Live motor tuning restarted the simulation");run(sim,rate);check(shaft(sim.snapshot(),a).speed>4,"Live motor tuning did not affect the running shaft");
        tuned.find(wheel)->shaft.inertia=2;rejects([&]{sim.apply_tuning(tuned);},"Flywheel inertia was rebuilt during Test");
    }
    auto detached=bench;detached.disconnect({wheel,0});rejects([&]{MachineSimulation invalid(detached);},"Unattached flywheel ran silently");
    const auto detached_checks=machine_checks(detached);check(std::any_of(detached_checks.begin(),detached_checks.end(),[&](const auto& c){return c.part==wheel&&c.title=="No shaft attached";}),"Unattached flywheel lacks a navigable check");

    auto first=machine_engine_example();machine_test::drive(first,1200);const auto first_cylinder=cylinder(first),first_shaft=machine_shaft_for(first,first_cylinder);
    auto second=machine_engine_example();machine_test::drive(second,2400);const auto second_cylinder=cylinder(second);
    std::vector<std::uint32_t> selected;for(const auto& p:second.parts)selected.push_back(p.id);const auto assembly=machine_extract_assembly(second,selected,"Second engine");
    const auto ids=machine_insert_assembly(first,assembly,{1.5f,.3f,0});std::uint32_t copied_cylinder=0,copied_motor=0;
    for(std::size_t i=0;i<assembly.parts.size();++i){if(assembly.parts[i].id==second_cylinder)copied_cylinder=ids[i];if(assembly.parts[i].kind==MachinePartKind::DriveMotor)copied_motor=ids[i];}
    const auto second_shaft=machine_shaft_for(first,copied_cylinder);check(first_shaft&&second_shaft&&first_shaft!=second_shaft,"Assembly insertion aliased an existing shaft");
    first.probes={{1,first_cylinder,"First engine RPM",MachineQuantity::Rpm},{2,copied_cylinder,"Second engine RPM",MachineQuantity::Rpm}};
    MachineSimulation engines(first);run(engines,11025);const auto state=engines.snapshot();check(close(shaft(state,first_shaft).speed,1200*2*pi/60)&&close(shaft(state,second_shaft).speed,2400*2*pi/60),"Independent engines did not retain their speeds");
    check(close(engines.measure().values[0],1200)&&close(engines.measure().values[1],2400),"Cylinder probes used a global RPM");
    auto live=first;live.find(copied_motor)->drive.target_speed=3600*2*pi/60;engines.apply_tuning(live);check(engines.snapshot().time==state.time,"Second engine tuning reset the first engine");run(engines,128);check(close(engines.measure().values[0],1200)&&close(engines.measure().values[1],3600),"Tuning one motor changed another engine");balance(shaft(engines.snapshot(),first_shaft));balance(shaft(engines.snapshot(),second_shaft));
    auto mixed=first;std::uint32_t controller=0;for(const auto& link:mixed.links)if(link.kind==MachineConnectionKind::Control&&link.b.part==first_cylinder)controller=link.a.part;mixed.disconnect({copied_cylinder,2});mixed.connect({controller,0},{copied_cylinder,2});rejects([&]{MachineSimulation invalid(mixed);},"An ECU silently controlled cylinders on different shafts");
    auto unbound=first;unbound.disconnect({first_cylinder,machine_port_index(*unbound.find(first_cylinder),"shaft")});rejects([&]{MachineSimulation invalid(unbound);},"Cylinder without a shaft ran on an implicit crank");
    auto reverse_engine=first;reverse_engine.find(first_shaft)->shaft.initial_speed=-10;rejects([&]{MachineSimulation invalid(reverse_engine);},"Unsupported reverse engine timing was enabled");
    auto moved=first;std::vector<std::vector<Vec3>> paths;for(const auto& p:first.parts)if(p.kind==MachinePartKind::Pipe)paths.push_back(p.points);moved.transform_selection({first_shaft},{},{.2f,.1f,0},0);validate_machine(moved);std::size_t index=0;for(const auto& p:moved.parts)if(p.kind==MachinePartKind::Pipe)check(p.points==paths[index++],"Moving a shaft silently changed gas geometry");
    const auto checks=machine_checks(unbound);check(std::any_of(checks.begin(),checks.end(),[&](const auto& c){return c.part==first_cylinder&&c.title=="No shaft attached";}),"Unattached cylinder lacks a navigable check");
    {
        std::vector<std::uint32_t> all;for(const auto& p:bench.parts)all.push_back(p.id);const auto module=machine_extract_assembly(bench,all,"Mechanical module");auto destination=bench;const auto copied=machine_insert_assembly(destination,module,{2,0,0});
        check(destination.probes.size()==8&&destination.probes[4].part==copied[0]&&destination.probes[4].part!=a,"Copied rotational instrument did not follow its copied shaft");
        const auto source_motor=std::find_if(module.parts.begin(),module.parts.end(),[&](const auto& p){return p.id==motor;})-module.parts.begin();destination.find(copied[source_motor])->drive.torque=12;
        MachineSimulation copied_sim(destination,1000);run(copied_sim,1000);check(copied_sim.measure().values[4]>copied_sim.measure().values[0]*2,"Copied motor or instrument still referenced the original shaft");
        const auto loose=machine_extract_assembly(bench,{wheel},"Unattached flywheel");const auto inserted=machine_insert_assembly(destination,loose,{3,0,0});check(machine_shaft_for(destination,inserted.front())==0,"External assembly attachment accidentally matched a destination ID");rejects([&]{MachineSimulation invalid(destination);},"Unresolved external shaft attachment was silently repaired");
    }
    const auto file=std::filesystem::temp_directory_path()/("fluidsim-mechanical-"+std::to_string(std::chrono::steady_clock::now().time_since_epoch().count())+".machine.toml");save_machine(bench,file);const auto restored=load_machine(file);std::filesystem::remove(file);
    check(restored.parts.size()==bench.parts.size()&&restored.links.size()==bench.links.size()&&restored.primary_shaft==bench.primary_shaft,"Explicit mechanical topology did not round-trip");MachineSimulation restored_sim(restored,1000);run(restored_sim,1000);check(close(shaft(restored_sim.snapshot(),a).speed,5./3),"Saved motor/flywheel configuration changed behavior");
    {std::ofstream legacy(file);legacy<<R"(format = 4
driven = true
drive_rpm = 600.0
shaft_inertia = 1.0
friction_constant = 2.0
friction_viscous = 0.0
parts = []
[[probes]]
id = 1
part = 0
name = 'Legacy output torque'
quantity = 7
[[probes]]
id = 2
part = 0
name = 'Legacy output power'
quantity = 8
)";}
    auto legacy=load_machine(file);check(legacy.probes[0].generated_output&&legacy.probes[0].part==legacy.primary_shaft,"Legacy torque gauge lost its measurement definition");MachineSimulation legacy_sim(legacy,1000);run(legacy_sim,10);check(close(legacy_sim.measure().values[0],-2)&&close(legacy_sim.measure().values[1],-2*600*2*pi/60),"Legacy gauges included speed-drive reaction torque");save_machine(legacy,file);check(load_machine(file).probes.front().generated_output,"Legacy gauge mode was not saved");std::filesystem::remove(file);
    std::cout<<"Mechanical parts, gas torque coupling, independent engines and targeted instruments passed\n";return 0;
}catch(const std::exception& error){std::cerr<<error.what()<<'\n';return 1;}