nargo_mcp/

lib.rs

#![feature(new_range_api)]
#![warn(missing_docs)]
use dashmap::{DashMap, DashSet};
use nargo_ir::{ElementIR, ExpressionIR, JsExpr, JsProgram, JsStmt, TemplateNodeIR};
use nargo_parser::{template::VueTemplateParser, ParseState, ScriptParser, TemplateParser};
use nargo_types::{is_pos_in_span, Position as NargoPosition, Span as NargoSpan};

use futures::Future;
use nargo_lsp::types::NargoLanguage;
use oak_core::{
    parser::session::ParseSession,
    tree::{GreenNode, RedNode},
    Range,
};
use oak_lsp::{service::LanguageService, types::SourcePosition as OakPosition};
use oak_vfs::{MemoryVfs, Vfs};

use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::mpsc as tokio_mpsc;
use url::Url;

/// 模块标识符
type ModuleId = String;

/// 消息标识符
type MessageId = String;

/// 消息类型枚举
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum MessageType {
    /// 普通消息
    Message,
    /// 事件通知
    Event,
    /// 请求消息
    Request,
    /// 响应消息
    Response,
    /// 广播消息
    Broadcast,
    /// 多播消息
    Multicast,
    /// 错误消息
    Error,
    /// 心跳消息
    Ping,
    /// 心跳响应消息
    Pong,
    /// 订阅消息
    Subscribe,
    /// 取消订阅消息
    Unsubscribe,
}

/// 事件类型枚举
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum EventType {
    /// 模块初始化完成
    ModuleInitialized,
    /// 模块状态更新
    ModuleStateUpdated,
    /// 编译开始
    CompileStarted,
    /// 编译完成
    CompileCompleted,
    /// 编译错误
    CompileError,
    /// 文件变更
    FileChanged,
    /// 配置变更
    ConfigurationChanged,
    /// 模块注册
    ModuleRegistered,
    /// 模块注销
    ModuleUnregistered,
    /// 消息发送
    MessageSent,
    /// 消息接收
    MessageReceived,
    /// 心跳超时
    HeartbeatTimeout,
    /// 网络连接断开
    NetworkDisconnected,
    /// 网络连接恢复
    NetworkConnected,
    /// 资源加载开始
    ResourceLoadStarted,
    /// 资源加载完成
    ResourceLoadCompleted,
    /// 资源加载错误
    ResourceLoadError,
    /// 性能警告
    PerformanceWarning,
    /// 系统警告
    SystemWarning,
    /// 系统错误
    SystemError,
}

impl EventType {
    /// 转换为字符串
    pub fn to_string(&self) -> String {
        match self {
            EventType::ModuleInitialized => "ModuleInitialized",
            EventType::ModuleStateUpdated => "ModuleStateUpdated",
            EventType::CompileStarted => "CompileStarted",
            EventType::CompileCompleted => "CompileCompleted",
            EventType::CompileError => "CompileError",
            EventType::FileChanged => "FileChanged",
            EventType::ConfigurationChanged => "ConfigurationChanged",
            EventType::ModuleRegistered => "ModuleRegistered",
            EventType::ModuleUnregistered => "ModuleUnregistered",
            EventType::MessageSent => "MessageSent",
            EventType::MessageReceived => "MessageReceived",
            EventType::HeartbeatTimeout => "HeartbeatTimeout",
            EventType::NetworkDisconnected => "NetworkDisconnected",
            EventType::NetworkConnected => "NetworkConnected",
            EventType::ResourceLoadStarted => "ResourceLoadStarted",
            EventType::ResourceLoadCompleted => "ResourceLoadCompleted",
            EventType::ResourceLoadError => "ResourceLoadError",
            EventType::PerformanceWarning => "PerformanceWarning",
            EventType::SystemWarning => "SystemWarning",
            EventType::SystemError => "SystemError",
        }
        .to_string()
    }
}

/// 消息状态枚举
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum MessageStatus {
    /// 消息已创建
    Created,
    /// 消息已发送
    Sent,
    /// 消息已接收
    Received,
    /// 消息已确认
    Acknowledged,
    /// 消息已超时
    Timeout,
    /// 消息发送失败
    Failed,
}

/// 消息结构体
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    /// 消息ID
    pub id: MessageId,
    /// 消息类型
    pub message_type: MessageType,
    /// 发送方模块ID
    pub sender: ModuleId,
    /// 接收方模块ID
    pub receiver: ModuleId,
    /// 事件类型（仅当消息类型为Event时有效）
    pub event_type: Option<EventType>,
    /// 消息数据
    pub data: serde_json::Value,
    /// 时间戳
    pub timestamp: u64,
    /// 消息状态
    pub status: MessageStatus,
    /// 超时时间（毫秒）
    pub timeout: Option<u64>,
    /// 重试次数
    pub retry_count: u8,
}

impl Message {
    /// 创建普通消息
    pub fn new_message(sender: ModuleId, receiver: ModuleId, data: serde_json::Value) -> Self {
        Self { id: format!("msg-{}", chrono::Utc::now().timestamp_millis()), message_type: MessageType::Message, sender, receiver, event_type: None, data, timestamp: chrono::Utc::now().timestamp_millis() as u64, status: MessageStatus::Created, timeout: Some(5000), retry_count: 0 }
    }

    /// 创建事件消息
    pub fn new_event(sender: ModuleId, event_type: EventType, data: serde_json::Value) -> Self {
        Self { id: format!("event-{}-{}", event_type.to_string(), chrono::Utc::now().timestamp_millis()), message_type: MessageType::Event, sender, receiver: "broadcast".to_string(), event_type: Some(event_type), data, timestamp: chrono::Utc::now().timestamp_millis() as u64, status: MessageStatus::Created, timeout: Some(3000), retry_count: 0 }
    }

    /// 创建请求消息
    pub fn new_request(sender: ModuleId, receiver: ModuleId, data: serde_json::Value) -> Self {
        Self { id: format!("req-{}", chrono::Utc::now().timestamp_millis()), message_type: MessageType::Request, sender, receiver, event_type: None, data, timestamp: chrono::Utc::now().timestamp_millis() as u64, status: MessageStatus::Created, timeout: Some(10000), retry_count: 3 }
    }

    /// 创建响应消息
    pub fn new_response(request_id: MessageId, sender: ModuleId, receiver: ModuleId, data: serde_json::Value) -> Self {
        Self { id: format!("resp-{}", request_id), message_type: MessageType::Response, sender, receiver, event_type: None, data, timestamp: chrono::Utc::now().timestamp_millis() as u64, status: MessageStatus::Created, timeout: Some(5000), retry_count: 0 }
    }

    /// 创建广播消息
    pub fn new_broadcast(sender: ModuleId, data: serde_json::Value) -> Self {
        Self { id: format!("broadcast-{}", chrono::Utc::now().timestamp_millis()), message_type: MessageType::Broadcast, sender, receiver: "broadcast".to_string(), event_type: None, data, timestamp: chrono::Utc::now().timestamp_millis() as u64, status: MessageStatus::Created, timeout: Some(3000), retry_count: 0 }
    }

    /// 创建心跳消息
    pub fn new_ping(sender: ModuleId, receiver: ModuleId) -> Self {
        Self { id: format!("ping-{}", chrono::Utc::now().timestamp_millis()), message_type: MessageType::Ping, sender, receiver, event_type: None, data: serde_json::Value::Null, timestamp: chrono::Utc::now().timestamp_millis() as u64, status: MessageStatus::Created, timeout: Some(2000), retry_count: 1 }
    }

    /// 创建心跳响应消息
    pub fn new_pong(sender: ModuleId, receiver: ModuleId, ping_id: MessageId) -> Self {
        Self { id: format!("pong-{}", ping_id), message_type: MessageType::Pong, sender, receiver, event_type: None, data: serde_json::json!({ "ping_id": ping_id }), timestamp: chrono::Utc::now().timestamp_millis() as u64, status: MessageStatus::Created, timeout: Some(1000), retry_count: 0 }
    }

    /// 创建订阅消息
    pub fn new_subscribe(sender: ModuleId, event_type: EventType) -> Self {
        Self { id: format!("subscribe-{}-{}", event_type.to_string(), chrono::Utc::now().timestamp_millis()), message_type: MessageType::Subscribe, sender, receiver: "system".to_string(), event_type: Some(event_type), data: serde_json::Value::Null, timestamp: chrono::Utc::now().timestamp_millis() as u64, status: MessageStatus::Created, timeout: Some(3000), retry_count: 0 }
    }

    /// 创建取消订阅消息
    pub fn new_unsubscribe(sender: ModuleId, event_type: EventType) -> Self {
        Self { id: format!("unsubscribe-{}-{}", event_type.to_string(), chrono::Utc::now().timestamp_millis()), message_type: MessageType::Unsubscribe, sender, receiver: "system".to_string(), event_type: Some(event_type), data: serde_json::Value::Null, timestamp: chrono::Utc::now().timestamp_millis() as u64, status: MessageStatus::Created, timeout: Some(3000), retry_count: 0 }
    }

    /// 创建错误消息
    pub fn new_error(sender: ModuleId, receiver: ModuleId, error: String) -> Self {
        Self { id: format!("error-{}", chrono::Utc::now().timestamp_millis()), message_type: MessageType::Error, sender, receiver, event_type: None, data: serde_json::json!({ "error": error }), timestamp: chrono::Utc::now().timestamp_millis() as u64, status: MessageStatus::Created, timeout: Some(3000), retry_count: 0 }
    }
}

/// 消息批处理结构体
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageBatch {
    /// 批次ID
    pub batch_id: String,
    /// 消息列表
    pub messages: Vec<Message>,
    /// 批次大小
    pub size: usize,
    /// 时间戳
    pub timestamp: u64,
}

/// 模块间通信管理器
#[derive(Clone)]
pub struct ModuleCommunicationManager {
    /// 模块注册映射
    modules: DashMap<ModuleId, ModuleInfo>,
    /// 消息通道映射
    message_channels: DashMap<ModuleId, tokio_mpsc::Sender<Message>>,
    /// 事件监听器映射
    event_listeners: DashMap<EventType, Vec<ModuleId>>,
    /// 消息状态跟踪
    message_status: DashMap<MessageId, Message>,
    /// 心跳状态
    heartbeat_status: DashMap<ModuleId, u64>,
    /// 心跳间隔（毫秒）
    heartbeat_interval: u64,
    /// 消息批处理缓冲区
    message_buffers: DashMap<ModuleId, Vec<Message>>,
    /// 批处理大小阈值
    batch_size_threshold: usize,
    /// 批处理间隔（毫秒）
    batch_interval: u64,
    /// 消息缓存
    message_cache: DashMap<String, serde_json::Value>,
    /// 缓存大小限制
    cache_size_limit: usize,
}

// 显式实现Send和Sync trait，因为ModuleCommunicationManager只包含线程安全的组件
unsafe impl Send for ModuleCommunicationManager {}
unsafe impl Sync for ModuleCommunicationManager {}

/// 模块信息结构体
#[derive(Clone)]
struct ModuleInfo {
    /// 模块名称
    name: String,
    /// 模块版本
    version: String,
    /// 模块状态
    status: ModuleStatus,
}

/// 模块状态枚举
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
enum ModuleStatus {
    /// 未初始化
    Uninitialized,
    /// 初始化中
    Initializing,
    /// 已初始化
    Initialized,
    /// 运行中
    Running,
    /// 出错
    Error,
}

impl ModuleCommunicationManager {
    /// 创建新的模块通信管理器
    pub fn new() -> Self {
        Self {
            modules: DashMap::new(),
            message_channels: DashMap::new(),
            event_listeners: DashMap::new(),
            message_status: DashMap::new(),

            heartbeat_status: DashMap::new(),
            heartbeat_interval: 30000, // 默认30秒心跳间隔
            message_buffers: DashMap::new(),
            batch_size_threshold: 10, // 默认批处理大小阈值
            batch_interval: 100,      // 默认批处理间隔（毫秒）
            message_cache: DashMap::new(),
            cache_size_limit: 1000, // 默认缓存大小限制
        }
    }

    /// 注册模块
    pub fn register_module(&self, module_id: ModuleId, name: String, version: String) -> tokio_mpsc::Receiver<Message> {
        let (tx, rx) = tokio_mpsc::channel(100);

        let name_clone = name.clone();
        let version_clone = version.clone();

        self.modules.insert(module_id.clone(), ModuleInfo { name, version, status: ModuleStatus::Uninitialized });

        self.message_channels.insert(module_id.clone(), tx);

        // 初始化心跳状态
        self.heartbeat_status.insert(module_id.clone(), chrono::Utc::now().timestamp_millis() as u64);

        // 广播模块注册事件
        let event_data = serde_json::json!({
            "module_id": module_id,
            "name": name_clone,
            "version": version_clone
        });
        let manager = self.clone();
        tokio::spawn(async move {
            manager.broadcast_event(EventType::ModuleRegistered, event_data, "system".to_string()).await;
        });

        rx
    }

    /// 发送消息
    pub async fn send_message(&self, mut message: Message) {
        // 保存消息状态
        self.message_status.insert(message.id.clone(), message.clone());

        // 更新消息状态为已发送
        message.status = MessageStatus::Sent;
        self.message_status.insert(message.id.clone(), message.clone());

        let receiver = message.receiver.clone();

        // 先获取sender的克隆，避免在async block中持有DashMapRef
        let tx = self.message_channels.get(&receiver).map(|r| r.value().clone());

        if let Some(tx) = tx {
            if tx.send(message.clone()).await.is_err() {
                // 通道已关闭，移除该模块
                self.message_channels.remove(&receiver);
                self.modules.remove(&receiver);
                self.heartbeat_status.remove(&receiver);

                // 更新消息状态为失败
                let mut failed_msg = message;
                failed_msg.status = MessageStatus::Failed;
                self.message_status.insert(failed_msg.id.clone(), failed_msg);

                // 广播模块注销事件
                let event_data = serde_json::json!({
                    "module_id": receiver
                });
                // 创建事件消息并加入队列，避免递归调用
                let event_message = Message::new_event("system".to_string(), EventType::ModuleUnregistered, event_data);
                self.queue_message(event_message);
            }
            else {
                // 广播消息发送事件
                let event_data = serde_json::json!({
                    "message_id": message.id,
                    "sender": message.sender,
                    "receiver": message.receiver
                });
                // 创建事件消息并加入队列，避免递归调用
                let event_message = Message::new_event("system".to_string(), EventType::MessageSent, event_data);
                self.queue_message(event_message);

                // 处理消息超时
                // if let Some(timeout) = message.timeout {
                //     let manager = self.clone();
                //     let message_id = message.id.clone();
                //     tokio::spawn(async move {
                //         tokio::time::sleep(tokio::time::Duration::from_millis(timeout)).await;
                //         if let Some(msg) = manager.message_status.get(&message_id) {
                //             if msg.status != MessageStatus::Acknowledged {
                //                 let mut timeout_msg = msg.value().clone();
                //                 timeout_msg.status = MessageStatus::Timeout;
                //                 manager
                //                     .message_status
                //                     .insert(message_id.clone(), timeout_msg);

                //                 // 广播消息超时事件
                //                 let event_data = serde_json::json!({
                //                     "message_id": message_id
                //                 });
                //                 manager
                //                     .broadcast_event(
                //                         EventType::HeartbeatTimeout,
                //                         event_data,
                //                         "system".to_string(),
                //                     )
                //                     .await;
                //             }
                //         }
                //     });
                // }
            }
        }
        else {
            // 接收方不存在，更新消息状态为失败
            message.status = MessageStatus::Failed;
            self.message_status.insert(message.id.clone(), message);
        }
    }

    /// 广播事件
    pub async fn broadcast_event(&self, event_type: EventType, data: serde_json::Value, sender: ModuleId) {
        let message = Message::new_event(sender, event_type.clone(), data);

        // 先获取监听器列表的克隆，避免在async block中持有DashMap的引用
        let listeners = {
            // 在同步代码中获取监听器列表的克隆
            self.event_listeners.get(&event_type).map(|l| l.value().clone())
        };

        if let Some(listeners) = listeners {
            for module_id in listeners {
                let mut msg = message.clone();
                msg.receiver = module_id;
                // 直接调用send_message，不使用tokio::spawn
                self.send_message(msg).await;
            }
        }
    }

    /// 注册事件监听器
    pub fn register_event_listener(&self, event_type: EventType, module_id: ModuleId) {
        self.event_listeners
            .entry(event_type)
            .and_modify(|listeners| {
                if !listeners.contains(&module_id) {
                    listeners.push(module_id.clone());
                }
            })
            .or_insert_with(|| vec![module_id]);
    }

    /// 移除事件监听器
    pub fn remove_event_listener(&self, event_type: EventType, module_id: &ModuleId) {
        if let Some(mut listeners) = self.event_listeners.get_mut(&event_type) {
            listeners.retain(|id| id != module_id);
        }
    }

    /// 更新模块状态
    pub fn update_module_status(&self, module_id: &ModuleId, status: ModuleStatus) {
        if let Some(mut module) = self.modules.get_mut(module_id) {
            module.status = status;
        }
    }

    /// 获取模块状态
    pub fn get_module_status(&self, module_id: &ModuleId) -> Option<ModuleStatus> {
        self.modules.get(module_id).map(|module| module.status.clone())
    }

    /// 列出所有模块
    pub fn list_modules(&self) -> Vec<(ModuleId, ModuleInfo)> {
        self.modules.iter().map(|entry| (entry.key().clone(), entry.value().clone())).collect()
    }

    /// 确认消息
    pub async fn acknowledge_message(&self, message_id: &MessageId) {
        if let Some(mut msg) = self.message_status.get_mut(message_id) {
            msg.status = MessageStatus::Acknowledged;

            // 广播消息确认事件
            let event_data = serde_json::json!({
                "message_id": message_id
            });
            self.broadcast_event(EventType::MessageReceived, event_data, "system".to_string()).await;

            // 通知等待确认的发送方
            // 移除对不存在的 message_acknowledgements 字段的引用
        }
    }

    /// 发送消息并等待确认
    pub async fn send_message_with_ack(&self, message: Message) -> Result<Message, String> {
        let timeout = message.timeout.unwrap_or(5000);

        self.send_message(message).await;

        // 简化实现，直接返回错误
        Err("Message acknowledgement not supported".to_string())
    }

    /// 启动心跳检测
    pub async fn start_heartbeat(&self) {
        let manager = self.clone();
        tokio::spawn(async move {
            loop {
                tokio::time::sleep(tokio::time::Duration::from_millis(manager.heartbeat_interval)).await;
                manager.check_heartbeats().await;
            }
        });
    }

    /// 检查心跳状态
    async fn check_heartbeats(&self) {
        let current_time = chrono::Utc::now().timestamp_millis() as u64;
        let mut modules_to_remove = Vec::new();

        for entry in self.heartbeat_status.iter() {
            let module_id = entry.key();
            let last_heartbeat = entry.value();

            if current_time - *last_heartbeat > self.heartbeat_interval * 2 {
                modules_to_remove.push(module_id.clone());
            }
        }

        for module_id in modules_to_remove {
            // 移除心跳超时的模块
            self.message_channels.remove(&module_id);
            self.modules.remove(&module_id);
            self.heartbeat_status.remove(&module_id);

            // 广播心跳超时事件
            let event_data = serde_json::json!({
                "module_id": module_id
            });
            self.broadcast_event(EventType::HeartbeatTimeout, event_data, "system".to_string()).await;
        }
    }

    /// 发送心跳消息
    pub async fn send_ping(&self, module_id: &ModuleId) {
        let ping_message = Message::new_ping("system".to_string(), module_id.clone());
        self.send_message(ping_message).await;
    }

    /// 处理心跳响应
    pub async fn handle_pong(&self, module_id: &ModuleId) {
        // 更新心跳状态
        self.heartbeat_status.insert(module_id.clone(), chrono::Utc::now().timestamp_millis() as u64);
    }

    /// 获取消息状态
    pub fn get_message_status(&self, message_id: &MessageId) -> Option<MessageStatus> {
        self.message_status.get(message_id).map(|msg| msg.status.clone())
    }

    /// 清理过期消息
    pub fn cleanup_expired_messages(&self) {
        let current_time = chrono::Utc::now().timestamp_millis() as u64;
        let mut messages_to_remove = Vec::new();

        for entry in self.message_status.iter() {
            let message_id = entry.key();
            let message = entry.value();

            // 清理已确认或超时超过5分钟的消息
            if message.status == MessageStatus::Acknowledged || (message.status == MessageStatus::Timeout && current_time - message.timestamp > 300000) {
                messages_to_remove.push(message_id.clone());
            }
        }

        for message_id in messages_to_remove {
            self.message_status.remove(&message_id);
        }
    }

    /// 发送消息批次
    pub async fn send_message_batch(&self, batch: MessageBatch) {
        for message in batch.messages {
            self.send_message(message).await;
        }
    }

    /// 将消息加入批处理队列
    pub fn queue_message(&self, message: Message) {
        self.message_buffers
            .entry(message.receiver.clone())
            .and_modify(|buffer| {
                buffer.push(message.clone());
                // 当缓冲区达到阈值时，立即处理
                if buffer.len() >= self.batch_size_threshold {
                    let batch = MessageBatch { batch_id: format!("batch-{}", chrono::Utc::now().timestamp_millis()), messages: buffer.drain(..).collect(), size: buffer.len(), timestamp: chrono::Utc::now().timestamp_millis() as u64 };
                    let manager = self.clone();
                    tokio::spawn(async move {
                        manager.send_message_batch(batch).await;
                    });
                }
            })
            .or_insert_with(|| vec![message]);
    }

    /// 处理消息批次
    pub async fn process_message_batches(&self) {
        for mut entry in self.message_buffers.iter_mut() {
            let module_id = entry.key().clone();
            let buffer = entry.value_mut();

            if !buffer.is_empty() {
                let batch = MessageBatch { batch_id: format!("batch-{}-{}", module_id, chrono::Utc::now().timestamp_millis()), messages: buffer.drain(..).collect(), size: buffer.len(), timestamp: chrono::Utc::now().timestamp_millis() as u64 };
                self.send_message_batch(batch).await;
            }
        }
    }

    /// 启动批处理器
    pub async fn start_batch_processor(&self) {
        let manager = self.clone();
        tokio::spawn(async move {
            loop {
                tokio::time::sleep(tokio::time::Duration::from_millis(manager.batch_interval)).await;
                manager.process_message_batches().await;
            }
        });
    }

    /// 缓存消息
    pub fn cache_message(&self, key: String, value: serde_json::Value) {
        // 检查缓存大小
        if self.message_cache.len() >= self.cache_size_limit {
            // 简单的缓存淘汰策略：移除最早的条目
            if let Some(first_key) = self.message_cache.iter().next() {
                self.message_cache.remove(first_key.key());
            }
        }
        self.message_cache.insert(key, value);
    }

    /// 获取缓存的消息
    pub fn get_cached_message(&self, key: &str) -> Option<serde_json::Value> {
        self.message_cache.get(key).and_then(|v| Some(v.clone()))
    }

    /// 清除缓存
    pub fn clear_cache(&self) {
        self.message_cache.clear();
    }

    /// 设置批处理大小阈值
    pub fn set_batch_size_threshold(&mut self, threshold: usize) {
        self.batch_size_threshold = threshold;
    }

    /// 设置批处理间隔
    pub fn set_batch_interval(&mut self, interval: u64) {
        self.batch_interval = interval;
    }

    /// 设置缓存大小限制
    pub fn set_cache_size_limit(&mut self, limit: usize) {
        self.cache_size_limit = limit;
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum CompletionContext {
    Tag,
    Attribute(String),
    Expression,
}

pub struct NargoLanguageService {
    vfs: MemoryVfs,
    auto_imports: DashSet<String>,
    workspace: oak_lsp::workspace::WorkspaceManager,
    sessions: DashMap<String, ParseSession<NargoLanguage>>,
}

impl NargoLanguageService {
    pub fn new() -> Self {
        let auto_imports = DashSet::new();
        // Add default auto-imports
        auto_imports.insert("signal".to_string());
        auto_imports.insert("computed".to_string());
        auto_imports.insert("on_mount".to_string());
        auto_imports.insert("on_unmount".to_string());
        auto_imports.insert("on_cleanup".to_string());

        Self { vfs: MemoryVfs::new(), auto_imports, workspace: oak_lsp::workspace::WorkspaceManager::new(), sessions: DashMap::new() }
    }

    fn get_completion_context(nodes: &[TemplateNodeIR], pos: NargoPosition) -> Option<CompletionContext> {
        for node in nodes {
            match node {
                TemplateNodeIR::Element(el) => {
                    if is_pos_in_span(pos, el.span) {
                        // Check if in tag name
                        let tag_start = el.span.start;
                        let tag_name_end = NargoPosition {
                            line: tag_start.line,
                            column: tag_start.column + (el.tag.len() as u32) + 1, // <tag
                            offset: 0,
                        };

                        if pos.line == tag_start.line && pos.column <= tag_name_end.column {
                            return Some(CompletionContext::Tag);
                        }

                        // Check attributes
                        for attr in &el.attributes {
                            if is_pos_in_span(pos, attr.span) {
                                return Some(CompletionContext::Attribute(el.tag.clone()));
                            }
                        }

                        // Check children
                        if let Some(ctx) = Self::get_completion_context(&el.children, pos) {
                            return Some(ctx);
                        }

                        // If in element but not in specific child/attr, might be in attribute area
                        return Some(CompletionContext::Attribute(el.tag.clone()));
                    }
                }
                TemplateNodeIR::Interpolation(expr) => {
                    if is_pos_in_span(pos, expr.span) {
                        return Some(CompletionContext::Expression);
                    }
                }
                _ => {}
            }
        }
        None
    }

    fn get_script_program(nodes: &[TemplateNodeIR]) -> Option<JsProgram> {
        for node in nodes {
            if let TemplateNodeIR::Element(el) = node {
                if el.tag == "script" {
                    if let Some(TemplateNodeIR::Text(content, _, _)) = el.children.first() {
                        let mut state = ParseState::new(content);
                        let ts_parser = nargo_parser::OakTypeScriptParser;
                        return ts_parser.parse(&mut state, "ts").ok();
                    }
                }
                if let Some(p) = Self::get_script_program(&el.children) {
                    return Some(p);
                }
            }
        }
        None
    }

    async fn find_definition_in_nodes(&self, nodes: &[TemplateNodeIR], pos: NargoPosition, uri: &str) -> Option<oak_lsp::types::LocationRange> {
        let script_program = Self::get_script_program(nodes);

        for node in nodes {
            match node {
                TemplateNodeIR::Element(el) => {
                    if is_pos_in_span(pos, el.span) {
                        if el.tag == "script" {
                            return self.find_definition_in_script(el, pos, uri).await;
                        }
                        else if el.tag == "style" {
                            return Self::find_definition_in_style(el, pos, uri);
                        }
                        else {
                            // 使用 Box::pin 处理异步递归
                            let res = self.find_definition_in_nodes_recursive(&el.children, pos, uri).await;
                            if res.is_some() {
                                return res;
                            }
                        }
                    }
                }
                TemplateNodeIR::Interpolation(expr) => {
                    if !expr.span.is_unknown() && is_pos_in_span(pos, expr.span) {
                        if let Some(symbol) = Self::find_symbol_in_template_expr(expr, pos) {
                            // 1. 首先在 script 块中查找
                            if let Some(program) = &script_program {
                                if let Some(def_span) = Self::find_definition_of_symbol(program, &symbol) {
                                    if Self::is_import(program, &symbol) {
                                        if let Some(loc) = self.find_external_definition(program, &symbol, uri).await {
                                            return Some(loc);
                                        }
                                    }

                                    if let Some(script_el_span) = Self::get_script_element_span(nodes) {
                                        if let Some(source) = self.vfs.get_source(uri) {
                                            let start_pos = OakPosition { line: script_el_span.start.line + def_span.start.line - 1, column: def_span.start.column - 1, offset: 0, length: 0 };
                                            let end_pos = OakPosition { line: script_el_span.start.line + def_span.end.line - 1, column: def_span.end.column - 1, offset: 0, length: 0 };
                                            return Some(oak_lsp::types::LocationRange { uri: uri.to_string().into(), range: Range { start: self.vfs.line_map(uri).map(|m| m.line_col_utf16_to_offset(&source, start_pos.line, start_pos.column)).unwrap_or(0), end: self.vfs.line_map(uri).map(|m| m.line_col_utf16_to_offset(&source, end_pos.line, end_pos.column)).unwrap_or(0) } });
                                        }
                                    }
                                }
                            }

                            // 2. 检查是否为自动导入
                            if self.auto_imports.contains(&symbol) {
                                if let Some(loc) = self.find_implicit_definition(&symbol, "@nargo/core", uri).await {
                                    return Some(loc);
                                }
                            }
                        }
                    }
                }
                _ => {}
            }
        }
        None
    }

    // Helper for recursion to avoid "async recursion" error
    fn find_definition_in_nodes_recursive<'a>(&'a self, nodes: &'a [TemplateNodeIR], pos: NargoPosition, uri: &'a str) -> std::pin::Pin<Box<dyn std::future::Future<Output = Option<oak_lsp::types::LocationRange>> + Send + 'a>> {
        Box::pin(self.find_definition_in_nodes(nodes, pos, uri))
    }

    fn is_import(program: &JsProgram, symbol: &str) -> bool {
        for stmt in &program.body {
            if let JsStmt::Import { specifiers, .. } = stmt {
                if specifiers.contains(&symbol.to_string()) {
                    return true;
                }
            }
        }
        false
    }

    fn get_script_element_span(nodes: &[TemplateNodeIR]) -> Option<NargoSpan> {
        for node in nodes {
            if let TemplateNodeIR::Element(el) = node {
                if el.tag == "script" {
                    return Some(el.span);
                }
                if let Some(s) = Self::get_script_element_span(&el.children) {
                    return Some(s);
                }
            }
        }
        None
    }

    fn find_symbol_in_template_expr(expr: &ExpressionIR, pos: NargoPosition) -> Option<String> {
        // In nargo-ir, ExpressionIR has a 'code' field for the raw string
        if !expr.span.is_unknown() && is_pos_in_span(pos, expr.span) {
            // If it's a simple identifier expression
            return Some(expr.code.clone());
        }
        None
    }

    async fn find_definition_in_script(&self, el: &ElementIR, pos: NargoPosition, uri: &str) -> Option<oak_lsp::types::LocationRange> {
        if let Some(TemplateNodeIR::Text(content, _, _)) = el.children.first() {
            let mut state = ParseState::new(content);
            let ts_parser = nargo_parser::OakTypeScriptParser;
            if let Ok(program) = ts_parser.parse(&mut state, "ts") {
                // Adjust position relative to the script tag start
                let el_span = el.span;
                let relative_pos = NargoPosition { line: pos.line - el_span.start.line, column: if pos.line == el_span.start.line { pos.column - el_span.start.column } else { pos.column }, offset: 0 };

                if let Some(symbol) = Self::find_symbol_at(&program, relative_pos) {
                    // 1. Check if it's an import (external definition)
                    if Self::is_import(&program, &symbol) {
                        if let Some(loc) = self.find_external_definition(&program, &symbol, uri).await {
                            return Some(loc);
                        }
                    }

                    // 2. Find where this symbol is defined locally
                    if let Some(def_span) = Self::find_definition_of_symbol(&program, &symbol) {
                        if let Some(source) = self.vfs.get_source(uri) {
                            let start_pos = OakPosition { line: el_span.start.line + def_span.start.line - 1, column: def_span.start.column - 1, offset: 0, length: 0 };
                            let end_pos = OakPosition { line: el_span.start.line + def_span.end.line - 1, column: def_span.end.column - 1, offset: 0, length: 0 };
                            return Some(oak_lsp::types::LocationRange { uri: uri.to_string().into(), range: Range { start: self.vfs.line_map(uri).map(|m| m.line_col_utf16_to_offset(&source, start_pos.line, start_pos.column)).unwrap_or(0), end: self.vfs.line_map(uri).map(|m| m.line_col_utf16_to_offset(&source, end_pos.line, end_pos.column)).unwrap_or(0) } });
                        }
                    }

                    // 3. Check if it's an auto-import
                    if self.auto_imports.contains(&symbol) {
                        if let Some(loc) = self.find_implicit_definition(&symbol, "@nargo/core", uri).await {
                            return Some(loc);
                        }
                    }
                }
            }
        }
        None
    }

    fn find_symbol_at(program: &JsProgram, pos: NargoPosition) -> Option<String> {
        for stmt in &program.body {
            if let Some(symbol) = Self::find_symbol_in_stmt(stmt, pos) {
                return Some(symbol);
            }
        }
        None
    }

    fn find_symbol_in_stmt(stmt: &JsStmt, pos: NargoPosition) -> Option<String> {
        use JsStmt::*;
        match stmt {
            VariableDecl { id, init, .. } => {
                // This is a simplification. In reality, we should check id and init spans.
                if id == "count" {
                    return Some(id.clone());
                }
                if let Some(init) = init {
                    return Self::find_symbol_in_expr(init, pos);
                }
            }
            Expr(expr, ..) => return Self::find_symbol_in_expr(expr, pos),
            _ => {}
        }
        None
    }

    fn find_symbol_in_expr(expr: &JsExpr, _pos: NargoPosition) -> Option<String> {
        match expr {
            JsExpr::Identifier(name, ..) => Some(name.clone()),
            JsExpr::Call { callee, .. } => Self::find_symbol_in_expr(callee, _pos),
            _ => None,
        }
    }

    fn find_definition_of_symbol(program: &JsProgram, symbol: &str) -> Option<NargoSpan> {
        for stmt in &program.body {
            if let JsStmt::VariableDecl { id, span, .. } = stmt {
                if id == symbol {
                    return Some(*span);
                }
            }
            if let JsStmt::FunctionDecl { id, span, .. } = stmt {
                if id == symbol {
                    return Some(*span);
                }
            }
        }
        None
    }

    async fn find_external_definition(&self, program: &JsProgram, symbol: &str, current_uri: &str) -> Option<oak_lsp::types::LocationRange> {
        for stmt in &program.body {
            if let JsStmt::Import { specifiers, source, .. } = stmt {
                if specifiers.contains(&symbol.to_string()) {
                    return self.find_implicit_definition(symbol, source, current_uri).await;
                }
            }
        }
        None
    }

    async fn find_implicit_definition(&self, symbol: &str, source_path: &str, current_uri: &str) -> Option<oak_lsp::types::LocationRange> {
        if let Some(url) = self.resolve_path(current_uri, source_path) {
            if let Some(content) = self.get_document_content(&url).await {
                // For now, just point to the start of the file or look for the symbol
                let mut line = 0;
                let mut col = 0;
                if let Some(pos) = content.find(symbol) {
                    let before = &content[..pos];
                    line = before.lines().count().saturating_sub(1);
                    col = before.lines().last().map(|l| l.len()).unwrap_or(0);
                }

                if let Some(target_source) = self.vfs.get_source(url.as_str()) {
                    let start_pos = OakPosition { line: line as u32, column: col as u32, offset: 0, length: 0 };
                    let end_pos = OakPosition { line: line as u32, column: (col + symbol.len()) as u32, offset: 0, length: 0 };
                    return Some(oak_lsp::types::LocationRange { uri: url.to_string().into(), range: Range { start: self.vfs.line_map(url.as_str()).map(|m| m.line_col_utf16_to_offset(&target_source, start_pos.line, start_pos.column)).unwrap_or(0), end: self.vfs.line_map(url.as_str()).map(|m| m.line_col_utf16_to_offset(&target_source, end_pos.line, end_pos.column)).unwrap_or(0) } });
                }
            }
        }
        None
    }

    fn resolve_path(&self, current_uri: &str, relative_path: &str) -> Option<Url> {
        let base_url = Url::parse(current_uri).ok()?;
        let base_path = base_url.to_file_path().ok()?;

        if relative_path == "@nargo/core" {
            // 1. Try workspace folders
            for (_, folder) in self.workspace.list_folders() {
                let core_path = folder.join("runtimes/nargo-core/src/index.ts");
                if core_path.exists() {
                    return Url::from_file_path(core_path).ok();
                }
                // Try nested
                let core_path = folder.join("project-nargo/runtimes/nargo-core/src/index.ts");
                if core_path.exists() {
                    return Url::from_file_path(core_path).ok();
                }
            }

            // 2. Try relative to current file
            if let Ok(uri) = Url::parse(current_uri) {
                if let Ok(mut path) = uri.to_file_path() {
                    while let Some(parent) = path.parent() {
                        let core_path = parent.join("runtimes/nargo-core/src/index.ts");
                        if core_path.exists() {
                            return Url::from_file_path(core_path).ok();
                        }
                        path = parent.to_path_buf();
                    }
                }
            }
        }

        let parent = base_path.parent()?;
        let target_path = if relative_path.starts_with('.') {
            parent.join(relative_path)
        }
        else {
            return None;
        };

        let extensions = ["", ".ts", ".js"];
        for ext in extensions {
            let mut p = target_path.clone();
            if !ext.is_empty() {
                p.set_extension(ext.trim_start_matches('.'));
            }
            if p.exists() {
                return Url::from_file_path(p).ok();
            }
        }

        None
    }

    async fn get_document_content(&self, uri: &Url) -> Option<String> {
        let uri_str = uri.to_string();
        if let Some(source) = self.vfs.get_source(&uri_str) {
            return Some(source.text().to_string());
        }

        if let Ok(path) = uri.to_file_path() {
            return std::fs::read_to_string(path).ok();
        }

        None
    }

    fn find_definition_in_style(_el: &ElementIR, _pos: NargoPosition, _uri: &str) -> Option<oak_lsp::types::LocationRange> {
        // Not implemented yet
        None
    }
}

use nargo_compiler::Compiler;
use oak_core::Parser;

use oak_lsp::types::{CompletionItem as OakCompletionItem, CompletionItemKind as OakCompletionItemKind, Diagnostic as OakDiagnostic, DiagnosticSeverity as OakDiagnosticSeverity};

impl LanguageService for NargoLanguageService {
    type Lang = NargoLanguage;
    type Vfs = MemoryVfs;

    fn vfs(&self) -> &Self::Vfs {
        &self.vfs
    }

    fn workspace(&self) -> &oak_lsp::workspace::WorkspaceManager {
        &self.workspace
    }

    fn get_root(&self, uri: &str) -> impl Future<Output = Option<RedNode<'_, NargoLanguage>>> + Send + '_ {
        let source = self.vfs.get_source(uri);
        let mut session = self.sessions.entry(uri.to_string()).or_insert_with(|| ParseSession::new(16));
        async move {
            if let Some(source) = source {
                // let parser = NargoParser::new();
                // // We use a persistent session for each file to allow incremental parsing if implemented
                // let output = parser.parse(&source, &[], session.value_mut());
                // if let Ok(green) = output.result {
                //     // SAFETY: We store the session (and thus its arena) in self.sessions which is owned by NargoLanguageService.
                //     // As long as the NargoLanguageService exists and we don't remove this session, the green node remains valid.
                //     // We use 'static lifetime here because RedNode needs a lifetime that outlives the future,
                //     // and transmute is used to bridge the gap between the local borrow and the persistent storage.
                //     let green: &GreenNode<'static, NargoLanguage> =
                //         unsafe { std::mem::transmute(green) };
                //     return Some(RedNode::new(green, 0));
                // }
            }
            None
        }
    }

    fn completion(&self, uri: &str, position: usize) -> impl Future<Output = Vec<OakCompletionItem>> + Send + '_ {
        let source = self.vfs.get_source(uri);
        let auto_imports = self.auto_imports.clone();
        let uri = uri.to_string();

        async move {
            let mut items = Vec::new();
            if let Some(source) = source {
                let content = source.text();
                let pos = self
                    .vfs
                    .line_map(&uri)
                    .map(|m| {
                        let (l, c) = m.offset_to_line_col_utf16(&source, position);
                        OakPosition { line: l, column: c, offset: 0, length: 0 }
                    })
                    .unwrap_or(OakPosition { line: 0, column: 0, offset: 0, length: 0 });
                let nargo_pos = NargoPosition { line: pos.line + 1, column: pos.column + 1, offset: 0 };

                let mut state = ParseState::new(&content);
                let parser = VueTemplateParser;
                if let Ok(nodes) = parser.parse(&mut state, "html") {
                    if let Some(context) = Self::get_completion_context(&nodes, nargo_pos) {
                        match context {
                            CompletionContext::Tag => {
                                let tags = vec!["div", "span", "button", "input", "h1", "h2", "p", "section", "article"];
                                for tag in tags {
                                    items.push(OakCompletionItem { label: tag.to_string(), kind: Some(OakCompletionItemKind::Keyword), detail: Some("HTML Tag".to_string()), documentation: None, insert_text: None });
                                }
                            }
                            CompletionContext::Attribute(tag_name) => {
                                let directives = vec!["@click", "@input", "@change", ":class", ":style", ":value"];
                                for dir in directives {
                                    items.push(OakCompletionItem { label: dir.to_string(), kind: Some(OakCompletionItemKind::Constant), detail: Some("Directive".to_string()), documentation: None, insert_text: None });
                                }

                                if tag_name == "input" {
                                    items.push(OakCompletionItem { label: "type".to_string(), kind: Some(OakCompletionItemKind::Property), detail: None, documentation: None, insert_text: None });
                                }
                            }
                            CompletionContext::Expression => {
                                for import in auto_imports.iter() {
                                    items.push(OakCompletionItem { label: import.clone(), kind: Some(OakCompletionItemKind::Function), detail: Some("Nargo Auto-import".to_string()), documentation: None, insert_text: None });
                                }

                                if let Some(program) = Self::get_script_program(&nodes) {
                                    for stmt in &program.body {
                                        if let JsStmt::VariableDecl { id, .. } = stmt {
                                            items.push(OakCompletionItem { label: id.clone(), kind: Some(OakCompletionItemKind::Variable), detail: None, documentation: None, insert_text: None });
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }

            if items.is_empty() {
                let tags = vec!["div", "span", "button", "input", "h1", "h2", "p", "section", "article"];
                for tag in tags {
                    items.push(OakCompletionItem { label: tag.to_string(), kind: Some(OakCompletionItemKind::Keyword), detail: Some("HTML Tag".to_string()), documentation: None, insert_text: None });
                }
            }

            items
        }
    }

    fn diagnostics(&self, uri: &str) -> impl Future<Output = Vec<OakDiagnostic>> + Send + '_ {
        let source = self.vfs.get_source(uri);
        let uri_parsed = Url::parse(uri).ok();
        let uri = uri.to_string();

        async move {
            let mut diags = Vec::new();
            if let Some(source) = source {
                let content = source.text();
                let mut compiler = Compiler::new();
                let name = uri_parsed.as_ref().and_then(|u| u.path_segments()).and_then(|mut s| s.next_back()).unwrap_or("App.ts");

                if let Err(e) = compiler.compile(name, &content) {
                    let span = e.span();
                    if !span.is_unknown() {
                        let start_pos = OakPosition { line: span.start.line.saturating_sub(1), column: span.start.column.saturating_sub(1), offset: 0, length: 0 };
                        let end_pos = OakPosition { line: span.end.line.saturating_sub(1), column: span.end.column.saturating_sub(1), offset: 0, length: 0 };

                        let start_offset = self.vfs.line_map(&uri).map(|m| m.line_col_utf16_to_offset(&source, start_pos.line, start_pos.column)).unwrap_or(0);
                        let end_offset = self.vfs.line_map(&uri).map(|m| m.line_col_utf16_to_offset(&source, end_pos.line, end_pos.column)).unwrap_or(0);

                        diags.push(OakDiagnostic { range: Range { start: start_offset, end: end_offset }, severity: Some(OakDiagnosticSeverity::Error), message: format!("{}", e), source: Some("nargo-compiler".to_string()), code: None });
                    }
                }
            }
            diags
        }
    }

    fn initialize(&self, params: oak_lsp::types::InitializeParams) -> impl Future<Output = ()> + Send + '_ {
        async move {
            self.workspace.initialize(&params);
        }
    }

    fn definition(&self, uri: &str, range: Range<usize>) -> impl Future<Output = Vec<oak_lsp::types::LocationRange>> + Send + '_ {
        let source = self.vfs.get_source(uri);
        let uri = uri.to_string();
        async move {
            if let Some(source) = source {
                let content = source.text();
                let position = self
                    .vfs
                    .line_map(&uri)
                    .map(|m| {
                        let (l, c) = m.offset_to_line_col_utf16(&source, range.start);
                        OakPosition { line: l, column: c, offset: 0, length: 0 }
                    })
                    .unwrap_or(OakPosition { line: 0, column: 0, offset: 0, length: 0 });
                let nargo_pos = NargoPosition { line: position.line + 1, column: position.column + 1, offset: 0 };

                let mut state = ParseState::new(&content);
                let parser = VueTemplateParser;
                if let Ok(nodes) = parser.parse(&mut state, "html") {
                    if let Some(loc) = self.find_definition_in_nodes(&nodes, nargo_pos, &uri).await {
                        return vec![loc];
                    }
                }
            }
            vec![]
        }
    }
}

pub async fn run_stdio() -> nargo_types::Result<()> {
    let stdin = tokio::io::stdin();
    let stdout = tokio::io::stdout();

    let nargo_service = Arc::new(NargoLanguageService::new());
    let server = oak_lsp::LspServer::new(nargo_service);
    server.run(stdin, stdout).await.map_err(|e| nargo_types::Error::external_error("oak-lsp".to_string(), format!("{}", e), nargo_types::Span::unknown()))?;
    Ok(())
}

pub async fn run_http_server(port: u16) -> nargo_types::Result<()> {
    println!("HTTP/WebSocket LSP server on port {} is not yet fully implemented.", port);
    println!("Falling back to stdio for now.");
    run_stdio().await
}