#include <WiFi.h>
#include <WiFiUdp.h>
#include "MidiData.h"

MidiData midi("192.168.8.173", 5000);
const char* downloadSSID = "YOUR_WIFI_SSID";
const char* downloadPassword = "YOUR_WIFI_PASSWORD";

const char* ssid = "YOUR_WIFI_SSID";
const char* password = "YOUR_WIFI_PASSWORD";
const unsigned int localPort = 2808;
WiFiUDP Udp;
char packetBuffer[255];

int DownloadStatus = 0;
bool wifiInitDone = false;

#define SERVO_PIN    4

#define SERVO_FREQUENCY  50
#define SERVO_RESOLUTION 12
#define PWM_CHANNEL      0

#define SERVO_OPEN_POSITION   280
#define SERVO_CLOSED_POSITION 325
const int MIN_SERVO_COUNTS = 250;
const int MAX_SERVO_COUNTS = 350;
const int SERVO_MOVEMENT_TIME_MS = 150;
const int SERVO_UPDATE_INTERVAL_MS = 10;
const int NOTE_HOLD_MIN_MS = 10;

int NOTE_COUNT = 0;
int* Midi_Note = nullptr;
int* start_time = nullptr;
int* noteDuration = nullptr;

uint32_t songTimebase = 0;
size_t currentNote = 0;
bool songPlaying = false;
bool fluteEnabled = false;
uint32_t lastFluteToggleTime = 0;

bool* noteActiveFlag = nullptr;
bool anyNoteActive = false;


bool servoMoving = false;
struct ServoState {
    int pin;
    float currentPos;
    float targetPos;
    float speed;
};
ServoState Servos[1];
unsigned long lastServoUpdate = 0;


void moveServo(int servoNumber, int target, int movementTime)
{
    if (servoNumber < 0 || servoNumber >= 1) return;
    int targetClamped = constrain(target, MIN_SERVO_COUNTS, MAX_SERVO_COUNTS);
    ServoState& servo = Servos[servoNumber];
    servo.targetPos = targetClamped;

    bool targetIsOpen = (targetClamped >= SERVO_OPEN_POSITION - 5);
    float distance = abs(servo.targetPos - servo.currentPos);

    if (distance < 0.01f) {
        servo.currentPos = servo.targetPos;
        servo.speed = 0.0f;
        servoMoving = false;
        return;
    }

    int updateTimes = movementTime / SERVO_UPDATE_INTERVAL_MS;
    if (updateTimes < 1) updateTimes = 1;
    servo.speed = distance / updateTimes;
    if (servo.speed < 0.2f) servo.speed = 0.2f;
    servoMoving = true;
}


void updateServo(ServoState& servo)
{
    if (servo.currentPos == servo.targetPos) return;

    if (servo.currentPos < servo.targetPos) {
        servo.currentPos += servo.speed;
        if (servo.currentPos >= servo.targetPos) servo.currentPos = servo.targetPos;
    } else {
        servo.currentPos -= servo.speed;
        if (servo.currentPos <= servo.targetPos) servo.currentPos = servo.targetPos;
    }
    ledcWrite(servo.pin, static_cast<int>(servo.currentPos));

    if (servo.currentPos == servo.targetPos && servoMoving) {
        servoMoving = false;
    }
}


void serviceServo()
{
    unsigned long currentTime = millis();
    if (currentTime - lastServoUpdate >= SERVO_UPDATE_INTERVAL_MS) {
        updateServo(Servos[0]);
        lastServoUpdate = currentTime;
    }
}


void setServoOpen(bool open)
{
    int targetPosition = open ? SERVO_OPEN_POSITION : SERVO_CLOSED_POSITION;
    moveServo(0, targetPosition, SERVO_MOVEMENT_TIME_MS);
}


void setFluteEnabled(bool enabled)
{
    unsigned long now = millis();
    if (enabled == fluteEnabled || (now - lastFluteToggleTime) < NOTE_HOLD_MIN_MS)
        return;

    fluteEnabled = enabled;
    lastFluteToggleTime = now;
    if (enabled) {
        setServoOpen(true);
    } else {
        setServoOpen(false);
    }
}


void stopSong()
{
    songPlaying = false;
    currentNote = 0;
    anyNoteActive = false;
    fluteEnabled = false;
    setFluteEnabled(false);
    if(noteActiveFlag) memset(noteActiveFlag, 0, NOTE_COUNT * sizeof(bool));
}

void startSong()
{
    if (NOTE_COUNT <= 0) {
        Serial.println("[SONG] Cannot start: no MIDI notes loaded.");
        return;
    }
    stopSong();
    songTimebase = millis();
    songPlaying = true;
    Serial.printf("[SONG] Playback start, total notes: %d\n", NOTE_COUNT);
}

void playMidiSong()
{
    if (!songPlaying) return;
    uint32_t elapsedMs = millis() - songTimebase;
    anyNoteActive = false;

    for(size_t i = 0; i < NOTE_COUNT; i++){
        uint32_t noteStart = (uint32_t)start_time[i];
        uint32_t noteEnd = noteStart + (uint32_t)noteDuration[i];
        bool isPlaying = (elapsedMs >= noteStart) && (elapsedMs < noteEnd);
        noteActiveFlag[i] = isPlaying;
        if(isPlaying) anyNoteActive = true;
    }

    while(currentNote < NOTE_COUNT){
        uint32_t nextStart = (uint32_t)start_time[currentNote];
        if(elapsedMs < nextStart) break;

        Serial.printf("[MIDI] Note #%u Triggered, pitch:%d, start:%dms, dur:%dms\n",
            currentNote, Midi_Note[currentNote], start_time[currentNote], noteDuration[currentNote]);
        currentNote++;
    }

    if(anyNoteActive && !fluteEnabled){
        setFluteEnabled(true);
    }
    if(!anyNoteActive && fluteEnabled){
        setFluteEnabled(false);
    }

    if(currentNote >= NOTE_COUNT && !anyNoteActive){
        stopSong();
        Serial.println("[SONG] Playback finished, waiting for G command");
    }
}


void connectToDownloadNetwork()
{
    Serial.print("[WIFI] Connecting download network");
    WiFi.mode(WIFI_STA);
    WiFi.begin(downloadSSID, downloadPassword);
    while (WiFi.status() != WL_CONNECTED) {
        Serial.print(".");
        delay(500);
    }
    Serial.println();
    Serial.print("[WIFI] Download IP: ");
    Serial.println(WiFi.localIP());
}

void downloadMidiSong()
{
    Serial.println("[MIDI] Fetch MIDI track Recorder...");
    midi.fetchByTrackName("Song.mid", "Recorder");
    NOTE_COUNT = midi.getNoteCount();
    Serial.printf("[MIDI] Server return %d notes\n", NOTE_COUNT);
    if (NOTE_COUNT <= 0) {
        Serial.println("[MIDI] No valid notes!");
        return;
    }

    delete[] Midi_Note;
    delete[] start_time;
    delete[] noteDuration;
    delete[] noteActiveFlag;

    Midi_Note = new int[NOTE_COUNT];
    start_time = new int[NOTE_COUNT];
    noteDuration = new int[NOTE_COUNT];
    noteActiveFlag = new bool[NOTE_COUNT]{false};

    if (!Midi_Note || !start_time || !noteDuration || !noteActiveFlag) {
        Serial.println("[MIDI] Memory alloc failed!");
        NOTE_COUNT = 0;
        return;
    }

    midi.printSummary();
    for (int i = 0; i < NOTE_COUNT; i++) {
        const NoteEvent* note = midi.getNote(i);
        if (!note) {
            Midi_Note[i] = start_time[i] = noteDuration[i] = 0;
            noteActiveFlag[i] = false;
            continue;
        }
        Midi_Note[i] = note->pitch;
        start_time[i] = note->startTimeMs;
        noteDuration[i] = note->durationMs;
    }
    Serial.println("[MIDI] MIDI data load complete");
}

void connectToConductorNetwork()
{
    Serial.println("[WIFI] Switch to control network");
    WiFi.disconnect();
    delay(300);
    WiFi.config(IPAddress(192, 168, 0, 78), IPAddress(192, 168, 0, 1), IPAddress(255, 255, 255, 0));
    WiFi.begin(ssid, password);
    while (WiFi.status() != WL_CONNECTED) {
        Serial.print(".");
        delay(500);
    }
    Serial.println();
    Serial.print("[WIFI] Control IP: ");
    Serial.println(WiFi.localIP());
    if (Udp.begin(localPort)) {
        Serial.printf("[UDP] Listen port %u\n", localPort);
    } else {
        Serial.println("[UDP] Start failed");
    }
}


void processUdpCommands()
{
    int packetSize = Udp.parsePacket();
    if (packetSize <= 0) return;
    int len = Udp.read(packetBuffer, sizeof(packetBuffer) - 1);
    if (len <= 0) return;
    packetBuffer[len] = '\0';
    char cmd = packetBuffer[0];
    Serial.print("[UDP] Recv cmd: ");
    Serial.println(packetBuffer);

    switch (cmd) {
        case 'T': case 't':
            Serial.println("[UDP] TEST cmd");
            break;
        case 'G': case 'g':
            Serial.println("[UDP] GO, restart song");
            startSong();
            break;
        default:
            Serial.printf("[UDP] Unknown cmd: %c\n", cmd);
            break;
    }
}


void setup()
{
    Serial.begin(921600);
    delay(500);
    Serial.println("\nSingle Servo Flute Controller Start");

    ledcAttachChannel(SERVO_PIN, SERVO_FREQUENCY, SERVO_RESOLUTION, PWM_CHANNEL);
    Servos[0] = {SERVO_PIN, (float)SERVO_CLOSED_POSITION, (float)SERVO_CLOSED_POSITION, 0.0f};
    ledcWrite(SERVO_PIN, SERVO_CLOSED_POSITION);
    fluteEnabled = false;
    servoMoving = false;
    lastServoUpdate = millis();

    connectToDownloadNetwork();
    downloadMidiSong();
    DownloadStatus = 1;

    connectToConductorNetwork();
    DownloadStatus = 2;
    wifiInitDone = true;
    Serial.println("[SYSTEM] Ready, wait UDP G command");
}


void loop()
{
    serviceServo();
    if(songPlaying) playMidiSong();
    if(DownloadStatus == 2) processUdpCommands();
    delay(1);
}