#include "lumacs/keybinding.hpp"
#include "lumacs/command_system.hpp" // Include for CommandSystem
#include <sstream>
#include <algorithm>
#include <cctype>

namespace lumacs {

// ============================================================================
// Key Implementation
// ============================================================================

namespace {
    // Helper to convert string to BaseKey
    BaseKey string_to_base_key(const std::string& s) {
        if (s.length() == 1) {
            char c = s[0];
            if (c >= 'a' && c <= 'z') return static_cast<BaseKey>(static_cast<int>(BaseKey::A) + (c - 'a'));
            if (c >= 'A' && c <= 'Z') return static_cast<BaseKey>(static_cast<int>(BaseKey::A) + (c - 'A'));
            if (c >= '0' && c <= '9') return static_cast<BaseKey>(static_cast<int>(BaseKey::D0) + (c - '0'));
        }
        if (s == "Space") return BaseKey::Space;
        if (s == "Return") return BaseKey::Return;
        if (s == "Tab") return BaseKey::Tab;
        if (s == "Escape") return BaseKey::Escape;
        if (s == "Backspace") return BaseKey::Backspace;
        if (s == "Delete") return BaseKey::Delete;
        if (s == "ArrowUp") return BaseKey::ArrowUp;
        if (s == "ArrowDown") return BaseKey::ArrowDown;
        if (s == "ArrowLeft") return BaseKey::ArrowLeft;
        if (s == "ArrowRight") return BaseKey::ArrowRight;
        if (s == "Home") return BaseKey::Home;
        if (s == "End") return BaseKey::End;
        if (s == "PageUp") return BaseKey::PageUp;
        if (s == "PageDown") return BaseKey::PageDown;
        if (s == "Insert") return BaseKey::Insert;
        if (s == ";") return BaseKey::Semicolon;
        if (s == "=") return BaseKey::Equal;
        if (s == ",") return BaseKey::Comma;
        if (s == "-") return BaseKey::Minus;
        if (s == ".") return BaseKey::Period;
        if (s == "/") return BaseKey::Slash;
        if (s == "`") return BaseKey::Backtick;
        if (s == "[") return BaseKey::LeftBracket;
        if (s == "\\") return BaseKey::Backslash;
        if (s == "]") return BaseKey::RightBracket;
        if (s == "'") return BaseKey::Quote;
        // Add more keys as needed
        return BaseKey::Unknown;
    }

    // Helper to convert BaseKey to string
    std::string base_key_to_string(BaseKey bk) {
        if (bk >= BaseKey::A && bk <= BaseKey::Z) return std::string(1, static_cast<char>(static_cast<int>(bk) - static_cast<int>(BaseKey::A) + 'a'));
        if (bk >= BaseKey::D0 && bk <= BaseKey::D9) return std::string(1, static_cast<char>(static_cast<int>(bk) - static_cast<int>(BaseKey::D0) + '0'));
        
        switch (bk) {
            case BaseKey::Space: return "Space";
            case BaseKey::Return: return "Return";
            case BaseKey::Tab: return "Tab";
            case BaseKey::Escape: return "Escape";
            case BaseKey::Backspace: return "Backspace";
            case BaseKey::Delete: return "Delete";
            case BaseKey::ArrowUp: return "ArrowUp";
            case BaseKey::ArrowDown: return "ArrowDown";
            case BaseKey::ArrowLeft: return "ArrowLeft";
            case BaseKey::ArrowRight: return "ArrowRight";
            case BaseKey::Home: return "Home";
            case BaseKey::End: return "End";
            case BaseKey::PageUp: return "PageUp";
            case BaseKey::PageDown: return "PageDown";
            case BaseKey::Insert: return "Insert";
            case BaseKey::F1: return "F1"; // Handle F keys properly
            case BaseKey::F2: return "F2";
            case BaseKey::F3: return "F3";
            case BaseKey::F4: return "F4";
            case BaseKey::F5: return "F5";
            case BaseKey::F6: return "F6";
            case BaseKey::F7: return "F7";
            case BaseKey::F8: return "F8";
            case BaseKey::F9: return "F9";
            case BaseKey::F10: return "F10";
            case BaseKey::F11: return "F11";
            case BaseKey::F12: return "F12";
            case BaseKey::Semicolon: return ";";
            case BaseKey::Equal: return "=";
            case BaseKey::Comma: return ",";
            case BaseKey::Minus: return "-";
            case BaseKey::Period: return ".";
            case BaseKey::Slash: return "/";
            case BaseKey::Backtick: return "`";
            case BaseKey::LeftBracket: return "[";
            case BaseKey::Backslash: return "\\";
            case BaseKey::RightBracket: return "]";
            case BaseKey::Quote: return "'";
            default: return "Unknown";
        }
    }
} // anonymous namespace


Key::Key(const std::string& key_name_str) {
    *this = parse(key_name_str);
}

Key Key::parse(const std::string& key_str) {
    Key key;
    std::string remaining = key_str;
    
    // Parse modifiers
    while (true) {
        if (remaining.length() >= 2 && remaining.substr(0, 2) == "C-") {
            key.modifiers = static_cast<Modifier>(static_cast<uint16_t>(key.modifiers) | static_cast<uint16_t>(Modifier::Control));
            remaining = remaining.substr(2);
        } else if (remaining.length() >= 2 && (remaining.substr(0, 2) == "M-" || remaining.substr(0, 2) == "A-")) {
            key.modifiers = static_cast<Modifier>(static_cast<uint16_t>(key.modifiers) | static_cast<uint16_t>(Modifier::Meta));
            remaining = remaining.substr(2);
        } else if (remaining.length() >= 2 && remaining.substr(0, 2) == "S-") {
            key.modifiers = static_cast<Modifier>(static_cast<uint16_t>(key.modifiers) | static_cast<uint16_t>(Modifier::Shift));
            remaining = remaining.substr(2);
        } else {
            break;
        }
    }
    
    key.base_key = string_to_base_key(remaining);
    
    // If Shift is pressed and the base key is a letter, convert base_key to uppercase if necessary,
    // but the BaseKey enum already handles 'a' through 'z' as lowercase by default for simplicity
    // If the original input was 'S-a', string_to_base_key("a") gives BaseKey::A.
    // So the explicit shift modifier should be enough.
    
    return key;
}

std::string Key::to_string() const {
    std::string result;
    
    if (static_cast<uint16_t>(modifiers) & static_cast<uint16_t>(Modifier::Control)) result += "C-";
    if (static_cast<uint16_t>(modifiers) & static_cast<uint16_t>(Modifier::Meta)) result += "M-";
    if (static_cast<uint16_t>(modifiers) & static_cast<uint16_t>(Modifier::Shift)) result += "S-";
    
    result += base_key_to_string(base_key);
    
    return result;
}

bool Key::operator==(const Key& other) const {
    return base_key == other.base_key && modifiers == other.modifiers;
}

bool Key::operator<(const Key& other) const {
    if (base_key != other.base_key) return base_key < other.base_key;
    return modifiers < other.modifiers;
}

// ============================================================================
// KeySequence Implementation
// ============================================================================

KeySequence::KeySequence(const std::vector<Key>& keys) : keys_(keys) {
}

KeySequence::KeySequence(const std::string& key_sequence_str) {
    if (key_sequence_str.empty()) {
        return;
    }
    
    // Split by spaces, but handle quoted strings
    std::istringstream iss(key_sequence_str);
    std::string token;
    
    while (iss >> token) {
        keys_.emplace_back(Key::parse(token));
    }
}

void KeySequence::add_key(const Key& key) {
    keys_.push_back(key);
}

void KeySequence::add_key(const std::string& key_str) {
    keys_.emplace_back(Key::parse(key_str));
}

bool KeySequence::is_prefix_of(const KeySequence& other) const {
    if (keys_.size() > other.keys_.size()) {
        return false;
    }
    
    return std::equal(keys_.begin(), keys_.end(), other.keys_.begin());
}

bool KeySequence::has_prefix(const KeySequence& prefix) const {
    return prefix.is_prefix_of(*this);
}

std::string KeySequence::to_string() const {
    if (keys_.empty()) {
        return "";
    }
    
    std::string result = keys_[0].to_string();
    
    for (size_t i = 1; i < keys_.size(); ++i) {
        result += " " + keys_[i].to_string();
    }
    
    return result;
}

void KeySequence::clear() {
    keys_.clear();
}

KeySequence KeySequence::subsequence(size_t start_index) const {
    if (start_index >= keys_.size()) {
        return KeySequence();
    }
    
    std::vector<Key> sub_keys(keys_.begin() + start_index, keys_.end());
    return KeySequence(sub_keys);
}

bool KeySequence::operator==(const KeySequence& other) const {
    return keys_ == other.keys_;
}

bool KeySequence::operator<(const KeySequence& other) const {
    return keys_ < other.keys_;
}

// ============================================================================
// KeyBinding Implementation
// ============================================================================

KeyBinding::KeyBinding(const KeySequence& seq, std::string cmd_name, std::string desc)
    : sequence(seq), command_name(std::move(cmd_name)), description(std::move(desc)) {
}

KeyBinding::KeyBinding(const std::string& key_str, std::string cmd_name, std::string desc)
    : sequence(key_str), command_name(std::move(cmd_name)), description(std::move(desc)) {
}

// ============================================================================
// KeyBindingManager Implementation
// ============================================================================

KeyBindingManager::KeyBindingManager(CommandSystem* command_system) 
    : root_(std::make_unique<TrieNode>()), // Initialize root_
      sequence_timeout_(std::chrono::milliseconds(1000)), 
      command_system_(command_system) {
}

void KeyBindingManager::bind(const KeySequence& sequence, std::string cmd_name, const std::string& description) {
    if (sequence.empty() || cmd_name.empty()) {
        return;
    }
    
    TrieNode* current = root_.get();
    for (const Key& key : sequence.keys()) {
        if (current->children.find(key) == current->children.end()) {
            current->children[key] = std::make_unique<TrieNode>();
        }
        current = current->children[key].get();
        current->ref_count++; // Increment ref count for this path
    }
    current->binding = KeyBinding(sequence, std::move(cmd_name), description);
}

void KeyBindingManager::bind(const std::string& key_str, std::string cmd_name, const std::string& description) {
    bind(KeySequence(key_str), std::move(cmd_name), description);
}

void KeyBindingManager::unbind(const KeySequence& sequence) {
    if (sequence.empty()) return;

    // Find the node corresponding to the end of the sequence
    TrieNode* current = root_.get();
    std::vector<TrieNode*> path_nodes; // To decrement ref_counts later
    path_nodes.push_back(current);

    for (const Key& key : sequence.keys()) {
        if (current->children.find(key) == current->children.end()) {
            return; // Sequence not found
        }
        current = current->children[key].get();
        path_nodes.push_back(current);
    }

    if (!current->binding) {
        return; // No exact binding at this node
    }

    // Remove the binding
    current->binding.reset();

    // Decrement ref counts along the path, and remove nodes if ref_count becomes 0
    // Traverse backwards from the node just before the binding
    for (int i = static_cast<int>(sequence.length()) - 1; i >= 0; --i) {
        TrieNode* parent = path_nodes[i];
        TrieNode* child = path_nodes[i+1];
        
        child->ref_count--; // Decrement ref count
        
        // If child's ref_count is 0 AND it has no own binding, we can remove it
        if (child->ref_count == 0 && !child->binding.has_value() && child->children.empty()) {
            parent->children.erase(sequence.keys()[i]);
        } else {
            // Stop if the node is still referenced or has a binding/children
            break;
        }
    }
}


void KeyBindingManager::unbind(const std::string& key_str) {
    unbind(KeySequence(key_str));
}

KeyProcessingResult KeyBindingManager::process_key(const Key& key) {
    // Check for timeout first
    if (is_building_sequence() && has_sequence_timed_out()) {
        clear_partial_sequence();
        // Fall through to process this key as a new sequence
    }
    
    // Add this key to the current sequence
    current_sequence_.add_key(key);
    
    // Update timestamp
    sequence_start_time_ = std::chrono::steady_clock::now();
    
    // Check for exact binding
    std::optional<KeyBinding> exact_binding = find_exact_binding(current_sequence_);
    if (exact_binding.has_value()) {
        // Found exact match
        clear_partial_sequence();
        
        try {
            // Execute the command via the CommandSystem
            if (command_system_) {
                CommandResult cmd_result = command_system_->execute(exact_binding->command_name, {}); // No args for now
                return {cmd_result.success ? KeyResult::Executed : KeyResult::Failed, cmd_result};
            }
            return {KeyResult::Failed, CommandResult{false, "No CommandSystem available"}};
        } catch (const std::exception& e) {
            return {KeyResult::Failed, CommandResult{false, std::string("Command execution failed: ") + e.what()}};
        } catch (...) {
            return {KeyResult::Failed, CommandResult{false, "Command execution failed with unknown error"}};
        }
    }
    
    // Check if this could be a prefix for other bindings
    if (has_prefix_bindings_impl(current_sequence_)) {
        return {KeyResult::Partial, std::nullopt};
    }
    
    // No binding found, clear sequence and return unbound
    clear_partial_sequence();
    return {KeyResult::Unbound, std::nullopt};
}

KeyProcessingResult KeyBindingManager::process_key(const std::string& key_str) {
    return process_key(Key::parse(key_str));
}

bool KeyBindingManager::has_binding(const KeySequence& sequence) const {
    // A sequence has a binding if it's an exact binding or a prefix for other bindings
    TrieNode* node = find_node(sequence);
    if (!node) return false;
    
    return node->binding.has_value() || !node->children.empty();
}

bool KeyBindingManager::has_exact_binding(const KeySequence& sequence) const {
    TrieNode* node = find_node(sequence);
    return node && node->binding.has_value();
}

bool KeyBindingManager::has_prefix_bindings(const KeySequence& sequence) const {
    TrieNode* node = find_node(sequence);
    return node && !node->children.empty();
}

void KeyBindingManager::clear_partial_sequence() {
    current_sequence_.clear();
    sequence_start_time_ = std::chrono::steady_clock::time_point{};
}

bool KeyBindingManager::is_building_sequence() const {
    return !current_sequence_.empty();
}

std::string KeyBindingManager::current_sequence_display() const {
    if (current_sequence_.empty()) {
        return "";
    }
    
    return current_sequence_.to_string() + "-";
}

std::vector<KeyBinding> KeyBindingManager::get_all_bindings() const {
    std::vector<KeyBinding> result;
    collect_bindings(root_.get(), KeySequence(), result);
    return result;
}

void KeyBindingManager::set_sequence_timeout(std::chrono::milliseconds timeout) {
    sequence_timeout_ = timeout;
}

bool KeyBindingManager::has_sequence_timed_out() const {
    if (!is_building_sequence()) {
        return false;
    }
    
    auto now = std::chrono::steady_clock::now();
    return (now - sequence_start_time_) > sequence_timeout_;
}

// Helper function to find a node in the trie
KeyBindingManager::TrieNode* KeyBindingManager::find_node(const KeySequence& sequence) const {
    TrieNode* current = root_.get();
    for (const Key& key : sequence.keys()) {
        if (current->children.find(key) == current->children.end()) {
            return nullptr;
        }
        current = current->children[key].get();
    }
    return current;
}

// Helper to traverse trie and collect all bindings
void KeyBindingManager::collect_bindings(TrieNode* node, KeySequence current_prefix, std::vector<KeyBinding>& result) const {
    if (!node) return;

    if (node->binding.has_value()) {
        result.push_back(node->binding.value());
    }

    for (const auto& pair : node->children) {
        KeySequence next_prefix = current_prefix;
        next_prefix.add_key(pair.first);
        collect_bindings(pair.second.get(), next_prefix, result);
    }
}

std::optional<KeyBinding> KeyBindingManager::find_exact_binding(const KeySequence& sequence) const {
    TrieNode* node = find_node(sequence);
    if (node && node->binding.has_value()) {
        return node->binding;
    }
    return std::nullopt;
}

bool KeyBindingManager::has_prefix_bindings_impl(const KeySequence& sequence) const {
    TrieNode* node = find_node(sequence);
    return node && !node->children.empty();
}

} // namespace lumacs

} // namespace lumacs