use std::collections::BTreeSet;
use std::time::Instant;
use bevy::app::App;
use bevy::color::palettes::basic::{RED, WHITE};
use bevy::color::palettes::css::LIMEGREEN;
use bevy::math::Vec3Swizzles;
use leafwing_input_manager::prelude::ActionState;
use microlp::{OptimizationDirection, Problem};
use crate::attachment::Parts;
use crate::client::input::ClientAction;
use crate::ecs::thruster::{PartThrusters, Thruster, ThrusterOfPart};
use crate::prelude::*;
use crate::client::input::util::ActionStateExt;
use crate::ecs::Me;
use crate::thrust::ThrustSolution;

pub fn client_thrusters_plugin(app: &mut App) {
    app
        .insert_resource(ThrusterDebugRes(false))
        .insert_resource(ThrustSolution {
            thrusters_on: BTreeSet::default(),
            converged: true,
        })
        .add_systems(Update, draw_thruster_debug)
        .add_systems(Update, solve_thrust);
}

#[derive(Resource, Deref)]
pub struct ThrusterDebugRes(pub bool);

fn draw_thruster_debug(
    thruster_debug_res: Res<ThrusterDebugRes>,
    thrusters: Query<(&Thruster, Entity, &GlobalTransform)>,
    thrust_solution: Res<ThrustSolution>,
    mut gizmos: Gizmos,
) {
    if !thruster_debug_res.0 { return };
    for thruster in thrusters {
        // Draw white if it's just a thruster, bright green if it's in the current thrust solution
        let mut color = if thrust_solution.thrusters_on.contains(&thruster.1) {
            LIMEGREEN
        } else {
            WHITE
        };
        // Exception: if the thrust solution failed to converge, RED
        if !thrust_solution.converged {
            color = RED;
        }
        let rescaled_thrust_vector = thruster.0.thrust_vector / 200.0;
        gizmos.arrow_2d(
            thruster.2.translation().xy(),
            thruster.2.translation().xy() + thruster.2.rotation().mul_vec3(rescaled_thrust_vector.extend(0.0)).xy(),
            color
        );
    }
}

fn solve_thrust(
    me: Query<(Option<&Parts>, &GlobalTransform, Entity), With<Me>>,
    parts: Query<&PartThrusters>,
    thrusters: Query<(&Thruster, &GlobalTransform)>,
    input: Res<ActionState<ClientAction>>,
    mut solution: ResMut<ThrustSolution>,
    mut events: MessageWriter<ThrustSolution>,
) {
    if !(
        input.button_changed(&ClientAction::ThrustForward)
            || input.button_changed(&ClientAction::ThrustBackward)
            || input.button_changed(&ClientAction::TorqueCw)
            || input.button_changed(&ClientAction::TorqueCcw)
            || input.button_changed(&ClientAction::ThrustRight)
            || input.button_changed(&ClientAction::ThrustLeft)
    ) { return; /* no changes, existing thrust solution is valid */ }

    trace!("input changed, recalculating thrust solution");
    let start = Instant::now();
    solution.thrusters_on.clear();
    solution.converged = false;

    let Ok((our_parts, hearty_transform, hearty)) = me.single() else {
        error!("could not solve for thrust: hearty does not exist?");
        error!("failed to solve for thrust after {}ms", start.elapsed().as_millis());
        return;
    };

    // determine our target vector:
    // unit vector in the intended direction of movement

    // Z-axis torque: this cursed thing is apparently standard
    // +Z == counterclockwise/ccw
    // -Z == clockwise/cw

    let mut target_unit_vector = Vec3::ZERO;

    if input.pressed(&ClientAction::ThrustForward) {
        target_unit_vector += hearty_transform.rotation() * Vec3::new(0.0, 1.0, 0.0);
    }
    if input.pressed(&ClientAction::ThrustBackward) {
        target_unit_vector += hearty_transform.rotation() * Vec3::new(0.0, -1.0, 0.0);
    }
    if input.pressed(&ClientAction::ThrustRight) {
        target_unit_vector += hearty_transform.rotation() * Vec3::new(1.0, 0.0, 0.0);
    }
    if input.pressed(&ClientAction::ThrustLeft) {
        target_unit_vector += hearty_transform.rotation() * Vec3::new(-1.0, 0.0, 0.0);
    }
    if input.pressed(&ClientAction::TorqueCw) {
        target_unit_vector += Vec3::new(0.0, 0.0, -1.0);
    }
    if input.pressed(&ClientAction::TorqueCcw) {
        target_unit_vector += Vec3::new(0.0, 0.0, 1.0);
    }

    if target_unit_vector == Vec3::ZERO {
        trace!("no buttons are pressed; zeroing thrust solution");
        trace!("solved thrust in {}ms", start.elapsed().as_millis());
        solution.converged = true;
        events.write(solution.clone());
        return;
    }

    let mut all_parts = vec![hearty];
    if let Some(parts) = our_parts {
        all_parts.extend(parts.iter());
    }

    // collect all thrusters on our ship, and figure out their thrust vectors
    let mut all_thrusters = vec![];

    for part in &all_parts {
        let Ok(part_thrusters) = parts.get(*part) else {
            continue;
        };
        for thruster_id in &**part_thrusters {
            let Ok((thruster, thruster_transform)) = thrusters.get(*thruster_id) else {
                warn!("issue while solving for thrust: thruster {:?} of part {:?} does not exist? skipping...", *thruster_id, *part);
                continue;
            };

            // determine the thruster force in world space
            let thruster_vector = thruster_transform.rotation().mul_vec3(thruster.thrust_vector.extend(0.0)).xy();

            // determine our xy offset from hearty
            let relative_translation = thruster_transform.translation().xy() - hearty_transform.translation().xy();
            // determine our rotational offset from hearty
            let relative_rotation = thruster_transform.rotation() * -hearty_transform.rotation();

            let thruster_torque = relative_translation.extend(0.0).cross(thruster_vector.extend(0.0)).z;

            // magically assemble the worldspace vector! for the solver (not shipspace)
            let target_vector = thruster_vector.extend(thruster_torque);

            all_thrusters.push((thruster_id, target_vector));
        }
    }

    // calculate thrust and torque values
    trace!("found {} thrusters, computing coefficients", all_thrusters.len());

    for thruster in &all_thrusters {
        trace!("thruster on ship: {:?}", thruster);
    }

    let coefficients = all_thrusters.iter()
        .map(|u| target_unit_vector.dot(u.1))
        .collect::<Vec<_>>();

    trace!("preparing model");
    let mut problem = Problem::new(OptimizationDirection::Maximize);

    // add variables to problem
    let variables = coefficients.iter()
        .map(|u| problem.add_binary_var(*u as f64))
        .collect::<Vec<_>>();

    trace!("prepared {} variables; solving", variables.len());

    let ssolution = match problem.solve() {
        Ok(soln) => soln,
        Err(e) => {
            match e {
                microlp::Error::Infeasible => {
                    error!("failed to solve for thrust: constraints cannot be satisfied");
                    error!("failed to solve for thrust after {}ms", start.elapsed().as_millis());
                    return;
                },
                microlp::Error::Unbounded => {
                    error!("failed to solve for thrust: system is unbounded");
                    error!("failed to solve for thrust after {}ms", start.elapsed().as_millis());
                    return;
                },
                microlp::Error::InternalError(e) => {
                    error!("failed to solve for thrust: solver encountered internal error: {e}");
                    error!("failed to solve for thrust after {}ms", start.elapsed().as_millis());
                    return;
                }
            }
        }
    };

    trace!("found thrust solution!");
    trace!("solution alignment (higher is better): {}", ssolution.objective());

    let mut new_soln = ThrustSolution {
        thrusters_on: BTreeSet::default(),
        converged: true
    };

    for thruster in all_thrusters.iter().enumerate() {
        trace!("solution: thruster #{} ({:?}): {}", thruster.0, thruster.1.0, ssolution.var_value_rounded(variables[thruster.0]));
        if ssolution.var_value_rounded(variables[thruster.0]) == 1.0 {
            new_soln.thrusters_on.insert(*thruster.1.0);
        }
    }

    let elapsed = start.elapsed();
    debug!(?elapsed, ?target_unit_vector, "solved for thrust");
    *solution = new_soln;
    events.write(solution.clone());
    return;
}