#include <Arduino.h>
#include <stm32f1xx_hal.h> // Required for HAL register definitions and clock functions
#include <HardwareTimer.h> // Include the HardwareTimer library

// --- Pin Definitions (Modify to match your wiring) ---
#define STEP_PIN_X     PA4
#define DIR_PIN_X      PA5
#define ENABLE_PIN_X   PA6 // Optional X Enable

#define STEP_PIN_Y     PB0 // Example: Use different pins for Y
#define DIR_PIN_Y      PB1
#define ENABLE_PIN_Y   PB5 // Optional Y Enable (Choose a suitable pin)

// --- Stepper Configuration ---
#define STEPS_PER_REVOLUTION 200
#define MICROSTEPS           1
#define STEP_PULSE_MICROSECONDS 5

// --- Global State Variables ---
volatile bool isRunningX = false;
volatile long stepsToGoX = 0;
uint32_t stepPeriodTicksX = 65535; // Default slowest
bool currentDirectionX = 0;

volatile bool isRunningY = false;
volatile long stepsToGoY = 0;
uint32_t stepPeriodTicksY = 65535; // Default slowest
bool currentDirectionY = 0;

uint32_t tim2_effective_clock_freq = 0; // Store calculated TIM2 *effective* clock frequency

// --- HardwareTimer Objects ---
HardwareTimer Timer2(TIM2); // Step generation timer
HardwareTimer Timer3(TIM3); // Pulse width timer (One-Pulse Mode)

// --- Callback Functions (Replaces ISRs) ---
void stepperISR_callback(); // Forward declaration
void resetISR_callback();   // Forward declaration


// --- Stepper Timer Initialization Function ---
void stepper_timers_init() {

    // == Configure Timer 2 (TIM2) for Stepper Driver Interrupt ==
    TIM_HandleTypeDef* StepperTimer_Handle = Timer2.getHandle();
    TIM_TypeDef *StepperTimer = StepperTimer_Handle->Instance; // Correct way to get TIM_TypeDef*

    __HAL_RCC_TIM2_CLK_ENABLE(); // Enable clock if not already done by library
    StepperTimer->CR1 &= ~TIM_CR1_CEN;
    StepperTimer->CR1 &= ~(TIM_CR1_DIR | TIM_CR1_CMS | TIM_CR1_CKD); // UP, Edge, Div1
    StepperTimer->CCER &= ~TIM_CCER_CC1E; // Disable CC1 output pin explicitly

    uint32_t tim2_raw_freq = Timer2.getTimerClkFreq();
    uint32_t tim2_prescaler_factor = (tim2_raw_freq / 1000000); // Aim for 1MHz effective clock
    if (tim2_prescaler_factor == 0) tim2_prescaler_factor = 1;
    Timer2.setPrescaleFactor(tim2_prescaler_factor);
    tim2_effective_clock_freq = Timer2.getTimerClkFreq() / Timer2.getPrescaleFactor();
    Serial.print("TIM2 Effective Clock Freq (Hz): ");
    Serial.println(tim2_effective_clock_freq);

    Timer2.setOverflow(0xFFFF + 1); // Max period (ARR = 0xFFFF)
    Timer2.setCaptureCompare(1, 1, TICK_COMPARE_FORMAT); // Initial dummy compare value
    Timer2.attachInterrupt(1, stepperISR_callback); // Attach to Ch1 Compare


    // == Configure Timer 3 (TIM3) for Stepper Port Reset Interrupt ==
    TIM_HandleTypeDef* ResetTimer_Handle = Timer3.getHandle();
    TIM_TypeDef *ResetTimer = ResetTimer_Handle->Instance; // Correct way to get TIM_TypeDef*

     __HAL_RCC_TIM3_CLK_ENABLE(); // Enable clock if not already done by library
    ResetTimer->CR1 &= ~TIM_CR1_CEN;
    ResetTimer->CR1 &= ~(TIM_CR1_DIR | TIM_CR1_CMS | TIM_CR1_CKD); // UP, Edge, Div1
    ResetTimer->CR1 |= TIM_CR1_OPM;  // *** Enable One-Pulse Mode ***

    uint32_t tim3_raw_freq = Timer3.getTimerClkFreq();
    uint16_t reset_timer_psc_factor = (tim3_raw_freq / 1000000); // Aim for 1MHz
    if (reset_timer_psc_factor == 0) reset_timer_psc_factor = 1;
    Timer3.setPrescaleFactor(reset_timer_psc_factor);

    uint16_t pulse_duration_ticks = STEP_PULSE_MICROSECONDS;
    if (pulse_duration_ticks == 0) pulse_duration_ticks = 1;
    Timer3.setOverflow(pulse_duration_ticks);
    Timer3.attachInterrupt(resetISR_callback); // Attach to Update Event

    // Priorities
    HAL_NVIC_SetPriority(TIM2_IRQn, 1, 0); // Lower priority for step generation
    HAL_NVIC_SetPriority(TIM3_IRQn, 0, 0); // Higher priority for pulse reset
    // Note: Library might manage enabling NVIC, but ensure they are enabled
    HAL_NVIC_EnableIRQ(TIM2_IRQn);
    HAL_NVIC_EnableIRQ(TIM3_IRQn);
}

// --- GPIO Initialization ---
void setup_gpio() {
    // Axis X
    pinMode(STEP_PIN_X, OUTPUT);
    pinMode(DIR_PIN_X, OUTPUT);
    digitalWrite(STEP_PIN_X, LOW);
    digitalWrite(DIR_PIN_X, currentDirectionX ? HIGH : LOW);
#ifdef ENABLE_PIN_X
    pinMode(ENABLE_PIN_X, OUTPUT);
    digitalWrite(ENABLE_PIN_X, HIGH); // Disable driver initially
#endif

    // Axis Y
    pinMode(STEP_PIN_Y, OUTPUT);
    pinMode(DIR_PIN_Y, OUTPUT);
    digitalWrite(STEP_PIN_Y, LOW);
    digitalWrite(DIR_PIN_Y, currentDirectionY ? HIGH : LOW);
#ifdef ENABLE_PIN_Y
    pinMode(ENABLE_PIN_Y, OUTPUT);
    digitalWrite(ENABLE_PIN_Y, HIGH); // Disable driver initially
#endif
}

// --- Speed Control ---
// Sets step period ticks for a specific axis (X or Y)
void setAxisSpeed(char axis, float stepsPerSecond) {
    if (stepsPerSecond <= 0.0) {
        Serial.print("Warning: Speed for axis "); Serial.print(axis); Serial.println(" set to near zero.");
        if (axis == 'x') stepPeriodTicksX = 0xFFFF;
        else if (axis == 'y') stepPeriodTicksY = 0xFFFF;
        return;
    }
    if (tim2_effective_clock_freq == 0) {
       tim2_effective_clock_freq = Timer2.getTimerClkFreq() / Timer2.getPrescaleFactor();
        if (tim2_effective_clock_freq == 0) {
            Serial.println("Error: TIM2 effective clock frequency not determined.");
            return;
        }
    }

    uint32_t ticks = (uint32_t)((float)tim2_effective_clock_freq / stepsPerSecond);
    uint32_t* targetPeriodPtr = (axis == 'x') ? &stepPeriodTicksX : &stepPeriodTicksY;

    if (ticks > 0xFFFF) {
        *targetPeriodPtr = 0xFFFF;
        Serial.print("Warning: Speed too slow for "); Serial.print(axis); Serial.println(", clamping to max period.");
    } else if (ticks < 10) {
        *targetPeriodPtr = 10;
        Serial.print("Warning: Speed too high for "); Serial.print(axis); Serial.println(", clamping to min period.");
    }
     else {
        *targetPeriodPtr = ticks;
    }
    Serial.print("Set "); Serial.print(axis); Serial.print(" Step Period Ticks: ");
    Serial.println(*targetPeriodPtr);

    // If the axis is running, we might want to update the main timer's CCR immediately
    // This requires careful handling if both are running - see ISR
}

// --- Enable/Disable Control ---
void setAxisEnable(char axis, bool enable) {
    uint8_t pin;
    const char* axisName;

    if (axis == 'x') {
#ifdef ENABLE_PIN_X
        pin = ENABLE_PIN_X;
        axisName = "X";
#else
        Serial.println("Enable Pin X not defined."); return;
#endif
    } else if (axis == 'y') {
#ifdef ENABLE_PIN_Y
        pin = ENABLE_PIN_Y;
        axisName = "Y";
#else
         Serial.println("Enable Pin Y not defined."); return;
#endif
    } else {
        Serial.println("Invalid axis for enable."); return;
    }

    // Most drivers are Active LOW enable
    digitalWrite(pin, enable ? LOW : HIGH);
    Serial.print("Driver "); Serial.print(axisName);
    Serial.println(enable ? " Enabled." : " Disabled.");
}


// --- Movement Control ---
// Starts a move for a specific axis if it's not already running
void moveAxisSteps(char axis, long steps) {
    // Select axis-specific variables
    volatile bool* isRunningPtr = (axis == 'x') ? &isRunningX : &isRunningY;
    volatile long* stepsToGoPtr = (axis == 'x') ? &stepsToGoX : &stepsToGoY;
    uint8_t dirPin = (axis == 'x') ? DIR_PIN_X : DIR_PIN_Y;
    bool* currentDirPtr = (axis == 'x') ? &currentDirectionX : &currentDirectionY; // Use global default dir
    // bool* currentDirPtr = nullptr; // Pointer to the default direction bool

    // if (axis == 'x') {
    // //    dirPin = DIR_PIN_X;
    //    currentDirPtr = &currentDirectionX; // Get ADDRESS of currentDirectionX
    // } else if (axis == 'y') {
    // //    dirPin = DIR_PIN_Y;
    //    currentDirPtr = &currentDirectionY; // Get ADDRESS of currentDirectionY
    // } else {
    //   Serial.println("Internal Error: Invalid axis passed to moveAxisSteps");
    //   return;
    // }

    if (*isRunningPtr) {
        Serial.print("Error: Axis "); Serial.print(axis); Serial.println(" is already running.");
        return;
    }
    if (steps == 0) {
        Serial.print("Error: Cannot move 0 steps for axis "); Serial.println(axis);
        return;
    }

    bool moveDir; // Direction for this specific move
    if (steps > 0) {
        moveDir = 0; // Forward relative to default 'currentDirection' for the axis maybe? Or absolute 0/1? Let's use absolute 0/1.
        // Or: moveDir = *currentDirPtr; // Move in the default direction
    } else {
        moveDir = 1; // Reverse relative to default
        // Or: moveDir = !(*currentDirPtr); // Move opposite the default direction
        steps = -steps;
    }
    digitalWrite(dirPin, moveDir ? HIGH : LOW);
    delayMicroseconds(10); // Increase delay slightly for safety

    *stepsToGoPtr = steps;

    Serial.print("Moving "); Serial.print(axis); Serial.print(" by ");
    Serial.print(steps); Serial.print(" steps, Direction: ");
    Serial.println(moveDir ? "Reverse" : "Forward");

    // Enable the specific driver
    setAxisEnable(axis, true);

    // Set the running flag for this axis
    *isRunningPtr = true;

    // Start Timer2 ONLY if it's not already running for the other axis
    if (!(axis == 'x' ? isRunningY : isRunningX)) { // If the *other* axis is NOT running
        Timer2.pause();
        Timer2.setCount(0);
        // Set initial compare based on the axis that is starting
        uint32_t initial_period = (axis == 'x') ? stepPeriodTicksX : stepPeriodTicksY;
        Timer2.setCaptureCompare(1, (initial_period > 0 ? initial_period : 1), TICK_COMPARE_FORMAT);
        // Timer2.refresh(); // Maybe needed?
        Timer2.resume();
        // Serial.print("Starting Timer2 for axis "); Serial.println(axis); // Debug
    } else {
        // Timer is already running, ISR will handle adding this axis
        // We might need to force an immediate CCR update if this axis is faster
        uint32_t otherPeriod = (axis == 'x') ? stepPeriodTicksY : stepPeriodTicksX;
        uint32_t myPeriod = (axis == 'x') ? stepPeriodTicksX : stepPeriodTicksY;
        if (myPeriod < otherPeriod) {
             // Force update CCR potentially sooner
             uint32_t current_count = Timer2.getCount();
             uint32_t next_compare = (current_count + myPeriod);
             uint32_t period = Timer2.getOverflow();
             if(period > 0) next_compare %= period; else next_compare &= 0xFFFF;
             if (next_compare == 0) next_compare = 1;
             Timer2.setCaptureCompare(1, next_compare, TICK_COMPARE_FORMAT);
             // Serial.println("Forcing faster CCR update"); // Debug
        }
        // Serial.print("Timer2 already running, axis "); Serial.print(axis); Serial.println(" added."); //Debug
    }
}

// --- Arduino Setup ---
void setup() {
    Serial.begin(115200);
    while (!Serial);
    Serial.println("\nSTM32 Dual Axis Stepper Control Initializing...");

    setup_gpio();
    stepper_timers_init();

    setAxisSpeed('x', 400); // Set default speeds
    setAxisSpeed('y', 400);

    Serial.println("Initialization Complete. Ready for commands:");
    Serial.println("  x[+/-]<steps>: Move X axis (e.g., x+2000, x-500)");
    Serial.println("  y[+/-]<steps>: Move Y axis (e.g., y+1000, y-300)");
    Serial.println("  sX<speed>    : Set X speed (steps/sec)");
    Serial.println("  sY<speed>    : Set Y speed (steps/sec)");
    Serial.println("  s<speed>     : Set both X and Y speed");
    Serial.println("  eX<0|1>      : Enable(0)/Disable(1) X Driver");
    Serial.println("  eY<0|1>      : Enable(0)/Disable(1) Y Driver");
    Serial.println("  e<0|1>       : Enable(0)/Disable(1) Both Drivers");
}

// --- Arduino Loop ---
void loop() {
    if (Serial.available() > 0) {
        String command = Serial.readStringUntil('\n');
        command.trim();
        // Serial.print("Rcvd: "); Serial.println(command); // Debug

        if (command.length() > 1) {
            char cmdAxis = tolower(command.charAt(0));
            char cmdType = tolower(command.charAt(0)); // Default if no axis specified
            String valueStr = command.substring(1);
            long valueInt = 0;
            float valueFloat = 0.0;

            // Check for axis specific commands (x, y)
            if (cmdAxis == 'x' || cmdAxis == 'y') {
                cmdType = tolower(command.charAt(1)); // e.g., x+ means cmdType='+'
                valueStr = command.substring(2);
            }

            // Parse value
            if (cmdType == '+' || cmdType == '-') {
                valueInt = valueStr.toInt();
                if (cmdType == '-') valueInt = -valueInt; // Apply sign
            } else if (cmdType == 's') {
                valueFloat = valueStr.toFloat();
            } else if (cmdType == 'e') {
                valueInt = valueStr.toInt();
            }


            // --- Execute Command ---
            switch (cmdAxis) {
                case 'x':
                    if (cmdType == '+' || cmdType == '-') {
                         if (valueInt != 0) moveAxisSteps('x', valueInt);
                         else Serial.println("X Move: Steps cannot be zero.");
                    } else if (cmdType == 's') {
                        if (valueFloat > 0) setAxisSpeed('x', valueFloat);
                        else Serial.println("X Speed: Must be positive.");
                    } else if (cmdType == 'e') {
                        setAxisEnable('x', (valueInt == 0)); // 0=Enable, 1=Disable
                    } else { Serial.println("Invalid X command type."); }
                    break;

                case 'y':
                     if (cmdType == '+' || cmdType == '-') {
                         if (valueInt != 0) moveAxisSteps('y', valueInt);
                         else Serial.println("Y Move: Steps cannot be zero.");
                    } else if (cmdType == 's') {
                        if (valueFloat > 0) setAxisSpeed('y', valueFloat);
                        else Serial.println("Y Speed: Must be positive.");
                    } else if (cmdType == 'e') {
                        setAxisEnable('y', (valueInt == 0)); // 0=Enable, 1=Disable
                    } else { Serial.println("Invalid Y command type."); }
                    break;

                case 's': // Global Speed Set 's<speed>'
                    valueFloat = valueStr.toFloat();
                    if (valueFloat > 0) {
                        setAxisSpeed('x', valueFloat);
                        setAxisSpeed('y', valueFloat);
                    } else { Serial.println("Speed: Must be positive."); }
                    break;

                 case 'e': // Global Enable Set 'e<0|1>'
                    valueInt = valueStr.toInt();
                    setAxisEnable('x', (valueInt == 0));
                    setAxisEnable('y', (valueInt == 0));
                    break;

                default:
                    Serial.println("Unknown command or axis.");
                    break;
            }
        } else if (command.length() > 0){
            Serial.println("Invalid command format.");
        }
    }
    // Optional auto-disable logic can go here
}


// --- Callback Function Definitions ---

// TIM2 Callback - Step Pulse Generation
void stepperISR_callback() {
    // Determine the CCR value for the *next* interrupt *before* modifying state
    uint32_t next_period_ticks = 0xFFFF; // Default to slowest if nothing ends up running
    bool x_will_run = (isRunningX && stepsToGoX > 0); // Will X be running after this step?
    bool y_will_run = (isRunningY && stepsToGoY > 0); // Will Y be running after this step?

    // Find the shortest period among axes that WILL be running next cycle
    if (x_will_run) next_period_ticks = stepPeriodTicksX;
    if (y_will_run && stepPeriodTicksY < next_period_ticks) next_period_ticks = stepPeriodTicksY;


    bool x_pulsed = false;
    bool y_pulsed = false;

    // --- X Axis Step Logic ---
    // Simple logic: if running, take a step.
    // Assumes timer interrupt rate is set by the faster axis (or CCR update)
    if (isRunningX) {
        digitalWrite(STEP_PIN_X, HIGH);
        x_pulsed = true;
        if (stepsToGoX > 0) stepsToGoX--; // Decrement if positive
        if (stepsToGoX <= 0) {            // Check if just finished
            isRunningX = false;
            stepsToGoX = 0;
        }
    }

    // --- Y Axis Step Logic ---
    if (isRunningY) {
        digitalWrite(STEP_PIN_Y, HIGH);
        y_pulsed = true;
        if (stepsToGoY > 0) stepsToGoY--;
        if (stepsToGoY <= 0) {
            isRunningY = false;
            stepsToGoY = 0;
        }
    }

    // --- Start Reset Timer if any axis pulsed ---
    if (x_pulsed || y_pulsed) {
       Timer3.pause();
       Timer3.setCount(0);
       Timer3.resume();
    }

    // --- Update Timer 2 CCR or Stop ---
    if (!isRunningX && !isRunningY) { // Check if BOTH are now stopped
        Timer2.pause(); // Stop timer completely
    } else {
       // Calculate and set next compare value
       uint32_t current_count = Timer2.getCount();
       uint32_t next_compare_value = (current_count + next_period_ticks); // Use shortest period of running axes
       uint32_t period = Timer2.getOverflow(); // Get ARR+1
       if (period > 0) { next_compare_value = next_compare_value % period; }
       else { next_compare_value &= 0xFFFF; }
       if (next_compare_value == 0) next_compare_value = 1; // Avoid compare at 0

       Timer2.setCaptureCompare(1, next_compare_value, TICK_COMPARE_FORMAT);
    }
}

// TIM3 Callback - Reset Step Pulses
void resetISR_callback() {
    // Reset both step pins LOW
    digitalWrite(STEP_PIN_X, LOW);
    digitalWrite(STEP_PIN_Y, LOW);
    // Timer stops automatically due to One-Pulse Mode (OPM)
}