#include <WiFi.h>
#include <esp_now.h>
#include <esp_mac.h>          // for esp_read_mac / ESP_MAC_WIFI_STA (IDF v5.x)
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

// ================= Display =================
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define SCREEN_ADDRESS 0x3C
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// ================= Touch Buttons (your mapping) =================
#define N_TOUCH 6
#define THRESHOLD 100000UL
int  touch_pins[N_TOUCH]         = {9, 8, 7, 2, 3, 1};
int  touch_values[N_TOUCH]       = {0, 0, 0, 0, 0, 0};
bool pin_touched_now[N_TOUCH]    = {false, false, false, false, false, false};
bool pin_touched_past[N_TOUCH]   = {false, false, false, false, false, false};

// Controls: P0=left, P1=right, P2=jump (P3..P5 unused)
#define IDX_LEFT  1
#define IDX_RIGHT 0
#define IDX_JUMP  2

// ================= Player / Net Packet =================
struct PlayerState {
  int x;
  int y;
  int vx;
  int vy;
  bool onGround;
};

struct NetPacket {
  PlayerState st;
  bool ready;
};

PlayerState myState, otherState;
volatile bool otherReady = false;   // becomes true when READY received
bool isP1 = false;                  // decided by self MAC
uint8_t selfMac[6] = {0};
uint8_t peerMac[6] = {0};

// Your two boards' MAC addresses (Wi-Fi STA)
uint8_t macP1[] = {0xD8, 0x3B, 0xDA, 0x75, 0x05, 0xAC};
uint8_t macP2[] = {0xD8, 0x3B, 0xDA, 0x75, 0xE1, 0x9C};

// ================= Touch =================
void update_touch() {
  for (int i = 0; i < N_TOUCH; i++) {
    int v = touchRead(touch_pins[i]);
    touch_values[i] = v;
    pin_touched_past[i] = pin_touched_now[i];
    pin_touched_now[i]  = (v > (int)THRESHOLD);
  }
}

// ================= Simple Physics =================
#define PLAYER_SIZE    8
#define MOVE_SPEED     2
#define JUMP_VELOCITY -6
#define GRAVITY        1
#define MAX_FALL_SPEED 4
#define GROUND_Y       (SCREEN_HEIGHT - PLAYER_SIZE)

void resetPlayer(PlayerState &p, bool leftSide) {
  p.x = leftSide ? 20 : (SCREEN_WIDTH - 20 - PLAYER_SIZE);
  p.y = GROUND_Y;
  p.vx = 0;
  p.vy = 0;
  p.onGround = true;
}

void updatePlayer(PlayerState &p, bool left, bool right, bool jump) {
  if (left && !right)      p.vx = -MOVE_SPEED;
  else if (right && !left) p.vx =  MOVE_SPEED;
  else                     p.vx =  0;

  if (jump && p.onGround) {
    p.vy = JUMP_VELOCITY;
    p.onGround = false;
  }

  p.vy += GRAVITY;
  if (p.vy > MAX_FALL_SPEED) p.vy = MAX_FALL_SPEED;

  p.x += p.vx;
  p.y += p.vy;

  if (p.y >= GROUND_Y) { p.y = GROUND_Y; p.vy = 0; p.onGround = true; }
  if (p.x < 0) p.x = 0;
  if (p.x > SCREEN_WIDTH - PLAYER_SIZE) p.x = SCREEN_WIDTH - PLAYER_SIZE;
}

// ================= ESP-NOW callbacks (IDF v5.x signatures) =================
void onDataSent(const wifi_tx_info_t *info, esp_now_send_status_t status) {
  // Optional debug:
  // Serial.println(status == ESP_NOW_SEND_SUCCESS ? "Send OK" : "Send FAIL");
}

void onDataRecv(const esp_now_recv_info_t *info, const uint8_t *data, int len) {
  if (len == (int)sizeof(NetPacket)) {
    NetPacket pkt;
    memcpy(&pkt, data, sizeof(NetPacket));
    otherState = pkt.st;
    if (pkt.ready) otherReady = true;
  }
}

// ================= ESP-NOW Init =================
bool initESPNowAndPeer() {
  if (esp_now_init() != ESP_OK) {
    Serial.println("ESP-NOW init failed");
    return false;
  }
  esp_now_register_send_cb(onDataSent);
  esp_now_register_recv_cb(onDataRecv);

  // Decide role from self MAC and set peer
  if (memcmp(selfMac, macP1, 6) == 0) {
    isP1 = true;
    memcpy(peerMac, macP2, 6);
  } else if (memcmp(selfMac, macP2, 6) == 0) {
    isP1 = false;
    memcpy(peerMac, macP1, 6);
  } else {
    // Unknown board: default to P2 and try peer as P1
    isP1 = false;
    memcpy(peerMac, macP1, 6);
  }

  esp_now_peer_info_t peerInfo = {};
  memcpy(peerInfo.peer_addr, peerMac, 6);
  peerInfo.channel = 0;
  peerInfo.encrypt = false;
  if (esp_now_add_peer(&peerInfo) != ESP_OK) {
    Serial.println("Failed to add peer");
    return false;
  }
  return true;
}

// ================= Setup =================
void setup() {
  Serial.begin(115200);

  // Optional: Pin 44 input as in your test (prevents conflicts on some XIAO boards)
  pinMode(44, INPUT);

  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println("SSD1306 allocation failed");
    for (;;);
  }
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println("Booting...");
  display.display();

  // --- WiFi STA & MAC (robust) ---
  WiFi.mode(WIFI_STA);
  delay(50);
  // 1) Try IDF API (most reliable)
  esp_err_t er = esp_read_mac(selfMac, ESP_MAC_WIFI_STA);
  // 2) Fallback to Arduino API if needed
  if (er != ESP_OK || (selfMac[0]==0 && selfMac[1]==0 && selfMac[2]==0 && selfMac[3]==0 && selfMac[4]==0 && selfMac[5]==0)) {
    WiFi.macAddress(selfMac);
  }

  Serial.print("Self MAC: ");
  for (int i = 0; i < 6; ++i) {
    if (i) Serial.print(':');
    if (selfMac[i] < 16) Serial.print('0');
    Serial.print(selfMac[i], HEX);
  }
  Serial.println();

  if (!initESPNowAndPeer()) {
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println("ESP-NOW init failed");
    display.display();
    for(;;);
  }

  // Warm-up touch (reduces first-read timeouts on some IDF builds)
  for (int k = 0; k < 5; ++k) {
    for (int i = 0; i < N_TOUCH; ++i) (void)touchRead(touch_pins[i]);
    delay(10);
  }

  resetPlayer(myState, true);
  resetPlayer(otherState, false);
}

// ================= Loop =================
void loop() {
  update_touch();

  // Controls
  bool left  = pin_touched_now[IDX_LEFT];
  bool right = pin_touched_now[IDX_RIGHT];
  bool jump  = pin_touched_now[IDX_JUMP];

  // Update my physics
  updatePlayer(myState, left, right, jump);

  // Send my packet (mark ready=true so peer exits splash)
  NetPacket pkt;
  pkt.st = myState;
  pkt.ready = true;
  esp_now_send(peerMac, (uint8_t*)&pkt, sizeof(NetPacket));

  // ---- Render ----
  display.clearDisplay();

  if (!otherReady) {
    // Waiting splash
    display.setCursor(10, 22);
    display.setTextSize(1);
    display.println("Waiting for Player...");
    display.setCursor(6, 36);
    display.print("My MAC ");
    for (int i = 0; i < 6; ++i) {
      if (i) display.print(':');
      if (selfMac[i] < 16) display.print('0');
      display.print(selfMac[i], HEX);
    }
  } else {
    // Ground line (optional visual)
    display.drawFastHLine(0, SCREEN_HEIGHT - 1, SCREEN_WIDTH, SSD1306_WHITE);

    // P1 = SQUARE, P2 = CIRCLE (same color)
    if (isP1) {
      display.fillRect(myState.x, myState.y, PLAYER_SIZE, PLAYER_SIZE, SSD1306_WHITE);
      display.fillCircle(otherState.x + PLAYER_SIZE/2, otherState.y + PLAYER_SIZE/2, PLAYER_SIZE/2, SSD1306_WHITE);
    } else {
      display.fillCircle(myState.x + PLAYER_SIZE/2, myState.y + PLAYER_SIZE/2, PLAYER_SIZE/2, SSD1306_WHITE);
      display.fillRect(otherState.x, otherState.y, PLAYER_SIZE, PLAYER_SIZE, SSD1306_WHITE);
    }
  }

  display.display();
  delay(50);
}