import edu.wpi.first.math.controller.PIDController;
import edu.wpi.first.wpilibj2.command.Command;

public class DriveDistanceCommand extends Command {

    private final DriveSubsystem m_drive;
    private final double m_targetMeters;
    private final double m_speed;

    // PID controller to keep the robot tracking straight.
    // Even with encoders on both sides, the robot can drift left/right
    // if one side has slightly more friction. A small kP on heading error
    // from the gyro corrects this without needing a full TurnToAngleCommand.
    private final PIDController m_straightController =
        new PIDController(0.02, 0, 0);
    private double m_initialHeading;

    public DriveDistanceCommand(DriveSubsystem drive, double meters, double speed) {
        m_drive        = drive;
        m_targetMeters = meters;
        m_speed        = speed;
        addRequirements(drive);
    }

    @Override
    public void initialize() {
        m_drive.resetEncoders();
        // Capture heading at start so we can correct drift.
        // We want to maintain this heading throughout the drive.
        m_initialHeading = m_drive.getHeadingDegrees();
        m_straightController.reset();
        m_straightController.setSetpoint(m_initialHeading);
    }

    @Override
    public void execute() {
        // Heading correction: PID output pushes the rotation input
        // to keep the robot pointing in the same direction it started.
        double headingError = m_drive.getHeadingDegrees() - m_initialHeading;
        double rotationCorrection = -m_straightController.calculate(headingError);
        m_drive.arcadeDrive(m_speed, rotationCorrection);
    }

    @Override
    public boolean isFinished() {
        // getAverageDistanceMeters() returns the average of left and right
        // encoder readings after the reset in initialize().
        return Math.abs(m_drive.getAverageDistanceMeters()) >= m_targetMeters;
    }

    @Override
    public void end(boolean interrupted) {
        m_drive.stop();
    }
}

public class TurnToAngleCommand extends Command {

    private final DriveSubsystem m_drive;
    private final double m_targetDegrees;

    // kP: how aggressively to correct heading error.
    // 0.01 means 1% motor output per degree of error.
    // At 10° error → 10% output. At 1° error → 1% output.
    // Start here and increase until you see overshoot, then reduce.
    private final PIDController m_turnController =
        new PIDController(0.01, 0, 0.001);

    public TurnToAngleCommand(DriveSubsystem drive, double degrees) {
        m_drive         = drive;
        m_targetDegrees = degrees;
        addRequirements(drive);
        // Tolerance: finish when within ±2° of target AND not moving fast
        m_turnController.setTolerance(2.0, 5.0); // degrees, degrees/sec
    }

    @Override
    public void initialize() {
        m_turnController.reset();
        m_turnController.setSetpoint(m_targetDegrees);
        // Note: we do NOT reset the gyro here.
        // The gyro measures absolute field heading. Resetting it mid-auto
        // would corrupt odometry and any subsequent heading-based commands.
        // Instead, pass the gyro's current reading and the desired absolute heading.
    }

    @Override
    public void execute() {
        double rotationOutput = m_turnController.calculate(m_drive.getHeadingDegrees());
        // Clamp output so the robot doesn't turn at full speed even when far off target
        rotationOutput = Math.copySign(Math.min(Math.abs(rotationOutput), 0.6), rotationOutput);
        m_drive.arcadeDrive(0, rotationOutput);
    }

    @Override
    public boolean isFinished() {
        // atSetpoint() is true only when BOTH position AND velocity are within tolerance.
        // This prevents finishing while the robot is still rotating through the target.
        return m_turnController.atSetpoint();
    }

    @Override
    public void end(boolean interrupted) {
        m_drive.stop();
    }
}

public class IntakeUntilSensorCommand extends Command {

    private final IntakeSubsystem m_intake;

    public IntakeUntilSensorCommand(IntakeSubsystem intake) {
        m_intake = intake;
        addRequirements(intake);
    }

    @Override
    public void initialize() {
        m_intake.runIntake(0.6);
        // Start the intake motors in initialize() rather than execute()
        // because the speed is constant. This is acceptable here —
        // the motor command doesn't depend on any loop calculation.
    }

    @Override
    public void execute() {
        // Nothing needed — motors are running at constant speed.
        // You could add dashboard telemetry here for debugging.
        SmartDashboard.putBoolean("Intake/SensorTripped", m_intake.hasPiece());
    }

    @Override
    public boolean isFinished() {
        return m_intake.hasPiece();
        // hasPiece() returns true when the beam break is tripped.
        // The command ends immediately when the game piece interrupts the beam.
        // No timer needed. No guess about how long intake takes.
    }

    @Override
    public void end(boolean interrupted) {
        m_intake.stopIntake();
        // If interrupted (e.g., auto period ended before piece was acquired),
        // stop the motors — don't leave them running.
    }
}

public Command getAutonomousCommand() {
    double startHeading = m_drive.getHeadingDegrees();
    // Capture starting heading once so relative turns are computed correctly.

    return Commands.sequence(

        // 1. Drive 2 meters forward at 50% speed, with 4s timeout safety net
        new DriveDistanceCommand(m_drive, 2.0, 0.5).withTimeout(4.0),

        // 2. Run intake until piece acquired, or 3s max
        new IntakeUntilSensorCommand(m_intake).withTimeout(3.0),

        // 3. Turn 90° right from starting heading (absolute target, not reset)
        new TurnToAngleCommand(m_drive, startHeading + 90.0).withTimeout(2.5),

        // 4. Drive 1.5 meters forward to scoring zone at 40% speed
        new DriveDistanceCommand(m_drive, 1.5, 0.4).withTimeout(4.0),

        // 5. Score (runs intake in reverse for 0.5 seconds — timed is fine here)
        Commands.run(() -> m_intake.runIntake(-0.8), m_intake).withTimeout(0.5)
    );
    // Note: the score step uses a timed command — and that's appropriate!
    // Ejecting a game piece doesn't have a sensor to detect "done."
    // The hybrid approach (sensor-based where possible, timed as fallback)
    // is exactly the pattern used by championship teams.
}