/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::BTreeSet;
use std::collections::HashMap;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_events::dispatch::EventDispatcher;
use buck2_interpreter::starlark_debug::StarlarkDebugController;
use debugserver_types as dap;
use dupe::Dupe;
use futures::StreamExt;
use gazebo::prelude::*;
use itertools::Itertools;
use starlark::debug::prepare_dap_adapter;
use starlark::debug::resolve_breakpoints;
use starlark::debug::DapAdapter;
use starlark::debug::DapAdapterClient;
use starlark::debug::DapAdapterEvalHook;
use starlark::debug::ResolvedBreakpoints;
use starlark::debug::StepKind;
use starlark::syntax::AstModule;
use starlark::syntax::Dialect;
use starlark::syntax::DialectTypes;
use thiserror::Error;
use tokio::select;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::Semaphore;
use tokio::time::Instant;
use tokio_stream::wrappers::UnboundedReceiverStream;
use tracing::debug;

use crate::controller::BuckStarlarkDebugController;
use crate::dap_api::dap_event;
use crate::dap_api::dispatch;
use crate::dap_api::err_response;
use crate::dap_api::ContinueArguments;
use crate::dap_api::DebugServer;
use crate::error::StarlarkDebuggerError;
use crate::run::ToClientMessage;
use crate::BuckStarlarkDebuggerHandle;
use crate::HandleData;
use crate::HandleId;
use crate::HookId;

/// The DAP capabilities that our debugserver supports.
///
/// See https://microsoft.github.io/debug-adapter-protocol/specification#Types_Capabilities
fn capabilities() -> serde_json::Value {
    // debugserver_types is out of date and missing fields on Capabilities and so we just construct
    // a little json map explicitly ourselves.
    serde_json::json!({
        "supports_configuration_done_request": true,
        "supports_evaluate_for_hovers": true,
        "supports_set_variable": true,
        "supports_step_in_targets_request": true,
        "supports_conditional_breakpoints": true,

        // This is different from starlark's `dap_capabilities`. The buck starlark debugger treats
        // each ongoing starlark Evaluation as a separate thread and handles requests appropriately.
        "supports_single_thread_execution_requests": true,
    })
}

#[derive(Debug, Error)]
enum DebuggerError {
    #[error("SetBreakpointsArguments invalid: {0:?}")]
    InvalidSetBreakpoints(dap::SetBreakpointsArguments),
}

/// The buck starlark debugger server. Most of the work is managed by the single-threaded server state.
///
/// There will be several references to the BuckStarlarkDebuggerServer instance and it will forward messages
/// along to the state.
#[derive(Debug)]
pub(crate) struct BuckStarlarkDebuggerServer {
    to_state: mpsc::UnboundedSender<ServerMessage>,
    next_handle_id: AtomicU32,
    /// When debugging a starlark evaluation, we wrap it in tokio::task::block_in_place (so that when it is paused
    /// it doesn't block a tokio worker thread), but we still want to ensure that we aren't over-saturating the
    /// local resources so we use this semaphore to limit how many evaluation can currently run.
    /// TODO(cjhopman): It probably actually makes sense for this to be a more general mechanism in buck so that
    /// long-running things aren't holding tokio workers busy without yielding. That'd then also better integrate
    /// with user-requested resource limits (e.g. `-j 4`).
    eval_semaphore: Arc<Semaphore>,
}

impl BuckStarlarkDebuggerServer {
    pub(crate) fn new(
        to_client: mpsc::UnboundedSender<ToClientMessage>,
        project_root: ProjectRoot,
    ) -> Self {
        let (to_state, state_recv) = mpsc::unbounded_channel();
        tokio::task::spawn(async move {
            let mut server = ServerState::new(to_client, project_root);
            let res = server.run(state_recv).await;
            // We always send the ::Shutdown message when the state thread finishes. It may be normal
            // shutdown (on detach()) or indicate an internal state error or that the client has already
            // exited (and dropped its side of the channel).
            let _ignored = server.to_client.send(ToClientMessage::Shutdown(res));
        });
        // TODO(cjhopman): figure out a better value for the limit on evaluations. This should probably be
        // passed in as we should at the least respect any `-j` flag.
        Self {
            to_state,
            eval_semaphore: Arc::new(Semaphore::new(num_cpus::get())),
            next_handle_id: AtomicU32::new(0),
        }
    }

    /// Send a message to the ServerState.
    ///
    /// Errors from the send are ignored, callers should assume that the state could encounter an
    /// error and shutdown early (but will not otherwise lose messages).
    ///
    /// Even if we were to return an error from this, that could only be used as an indication that the
    /// state thread has finished but the `Ok()` would not necessarily mean that the state thread would
    /// ever see the sent message.
    fn maybe_to_state(&self, msg: ServerMessage) {
        let _ignored = self.to_state.send(msg);
    }

    pub(crate) fn new_handle(
        self: &Arc<Self>,
        events: EventDispatcher,
    ) -> Option<BuckStarlarkDebuggerHandle> {
        let handle_id = HandleId(self.next_handle_id.fetch_add(1, Ordering::Relaxed));
        self.maybe_to_state(ServerMessage::NewHandle {
            id: handle_id,
            events,
        });
        Some(BuckStarlarkDebuggerHandle(Arc::new(HandleData {
            id: handle_id,
            server: self.clone(),
        })))
    }

    /// Acquires a StarlarkDebugController for starlark evaluation.
    ///
    /// The debugger server limits how many concurrent evaluations there can be. This is important as
    /// `with_starlark_eval_provider` will send starlark evaluations to tokio blocking threads when a
    /// debugger is attached and so we need to otherwise limit the concurrent ones.
    pub(crate) async fn start_eval(
        self: &Arc<Self>,
        handle: &BuckStarlarkDebuggerHandle,
        description: &str,
    ) -> anyhow::Result<Box<dyn StarlarkDebugController>> {
        debug!("starting debug-hooked eval {}", description);
        let permit = self.eval_semaphore.dupe().acquire_owned().await?;
        let (send, recv) = oneshot::channel();

        // Here we actually need a response from the state. We still ignore any send
        // errors, but if the state thread has hit some internal error and is shutting down
        // we'll notice when we poll the recv channel.
        self.maybe_to_state(ServerMessage::NewHook {
            handle: handle.dupe(),
            description: description.to_owned(),
            response_channel: send,
        });

        let (hook_id, eval_wrapper) = match recv.await {
            Ok(v) => v,
            Err(..) => {
                // This indicates the state thread is shutting down (or hit an internal error).
                // That could be due to the debugger detaching from buck2. This does not indicate
                // an error for other on-going buck commands, and so we'll allow starlark execution
                // to continue as normal. In this case, the hook_id doesn't matter.
                (HookId(u32::MAX), None)
            }
        };
        Ok(Box::new(BuckStarlarkDebugController::new(
            eval_wrapper,
            hook_id,
            description,
            self,
            permit,
        )))
    }

    /// Indicates that a hook has been dropped and so an evaluation has finished.
    pub(crate) fn drop_hook(&self, hook_id: HookId) {
        self.maybe_to_state(ServerMessage::DropHook { id: hook_id });
    }

    /// Indicates that a handle has been dropped and so a command (and its associated
    /// dice computation) has finished.
    pub(crate) fn drop_handle(&self, handle_id: HandleId) {
        self.maybe_to_state(ServerMessage::DropHandle { id: handle_id });
    }

    /// Called when a starlark evaluation is paused (e.g. at a breakpoint).
    pub(crate) fn event_stopped(&self, hook_id: HookId) {
        self.maybe_to_state(ServerMessage::EvalStopped { hook_id });
    }

    /// Called to forward along requests from the DAP client.
    pub(crate) fn send_request(&self, req: dap::Request) -> anyhow::Result<()> {
        // If the state encountered an error or is shutting down, it may never see this
        // request. But that's okay since in that case it will send back the Shutdown message
        // and that'll make its way back to the client.
        self.maybe_to_state(ServerMessage::DapRequest { req });
        Ok(())
    }

    /// Called when the DAP client has disconnected.
    pub(crate) fn detach(&self) -> anyhow::Result<()> {
        self.maybe_to_state(ServerMessage::Detach);
        Ok(())
    }
}

/// Messages to the debugger server state
enum ServerMessage {
    NewHook {
        handle: BuckStarlarkDebuggerHandle,
        description: String,
        response_channel: oneshot::Sender<(HookId, Option<Box<dyn DapAdapterEvalHook>>)>,
    },
    NewHandle {
        id: HandleId,
        events: EventDispatcher,
    },
    DropHook {
        id: HookId,
    },
    DropHandle {
        id: HandleId,
    },
    DapRequest {
        req: dap::Request,
    },
    EvalStopped {
        hook_id: HookId,
    },
    Detach,
}

/// The ServerState is the main thing implementing the debug adapter protocol.
///
/// It runs on a single thread to more easily handle the concurrent requests and events
/// from the DAP client and buck's multithreaded starlark evaluation.
#[derive(Debug)]
struct ServerState {
    /// Sends messages back to the DAP client. Errors on this channel indicate the
    /// `starlark debug-attach` command has ended and so we treat them as state
    /// errors to stop the state thread (similar to a detach() call).
    to_client: mpsc::UnboundedSender<ToClientMessage>,

    /// The currently set breakpoints. New hooks will be initialized with these.
    set_breakpoints: HashMap<String, ResolvedBreakpoints>,

    /// The project root is used to get the current source code to resolve breakpoints.
    project_root: ProjectRoot,

    /// Currently executing buck commands, this is primarily used to send debugger snapshots.
    current_commands: HashMap<HandleId, CommandState>,

    /// Current starlark evaluation hooks.
    current_hooks: HashMap<HookId, HookState>,

    /// HookIds are simply incrementing.
    next_hook_id: HookId,

    /// We keep a pool of "pseudo" thread ids to appear to use only as many threads as we ever
    /// have concurrent starlark evaluations. Some DAP clients may provide slightly nicer ui
    /// when we do this instead of providing the real thread ids (which might range over some
    /// 100s of ids)
    free_pseudo_threads: BTreeSet<u32>,
    next_pseudo_thread: u32,
}

static TOP_FRAME_LOCALS_ID: i64 = 2000;

impl DebugServer for ServerState {
    fn initialize(
        &mut self,
        _x: dap::InitializeRequestArguments,
    ) -> anyhow::Result<Option<serde_json::Value>> {
        Ok(Some(capabilities()))
    }

    fn set_breakpoints(
        &mut self,
        mut x: dap::SetBreakpointsArguments,
    ) -> anyhow::Result<dap::SetBreakpointsResponseBody> {
        // buck will use the project-relative paths when parsing asts with the starlark interpreter. We need to match that.
        let source = x
            .source
            .path
            .take()
            .ok_or_else(|| DebuggerError::InvalidSetBreakpoints(x.clone()))?;
        // vscode sends absolute paths, so need to relativize them
        let project_relative = self.project_root.relativize(AbsNormPath::new(&source)?)?;
        let source = project_relative.to_string();
        x.source.path = Some(source.clone());

        // We currently just resolve new breakpoints against the current state of the file. This isn't quite correct, but oh well.
        let resolved = resolve_breakpoints(&x, &self.get_ast(&project_relative)?)?;
        for hook_state in self.current_hooks.values() {
            hook_state.adapter.set_breakpoints(&source, &resolved)?;
        }
        let response = resolved.to_response();
        self.set_breakpoints.insert(source, resolved);
        Ok(response)
    }

    fn set_exception_breakpoints(
        &mut self,
        _x: dap::SetExceptionBreakpointsArguments,
    ) -> anyhow::Result<()> {
        // TODO(cjhopman): This may not make sense in starlark and could be a more informative error. Or possibly we
        // could use it to break on `fail()`.
        Err(StarlarkDebuggerError::Unimplemented.into())
    }

    fn attach(&mut self, _x: dap::AttachRequestArguments) -> anyhow::Result<()> {
        Ok(())
    }

    fn threads(&mut self) -> anyhow::Result<dap::ThreadsResponseBody> {
        let mut threads = Vec::with_capacity(self.current_hooks.len());
        for hook_state in self
            .current_hooks
            .values()
            .sorted_unstable_by_key(|v| v.pseudo_thread_id)
        {
            threads.push(dap::Thread {
                id: hook_state.pseudo_thread_id as i64,
                name: hook_state.pseudo_thread_name.clone(),
            });
        }
        Ok(dap::ThreadsResponseBody { threads })
    }

    fn configuration_done(&mut self) -> anyhow::Result<()> {
        Ok(())
    }

    fn stack_trace(
        &mut self,
        x: dap::StackTraceArguments,
    ) -> anyhow::Result<dap::StackTraceResponseBody> {
        let hook = self.find_hook_by_pseudo_thread(x.thread_id)?;
        let mut trace_response = hook.adapter.stack_trace(x)?;
        for frame in &mut trace_response.stack_frames {
            // rewrite the sources to be absolute (like vscode sent us)
            if let Some(source) = &mut frame.source {
                if let Some(path) = &mut source.path {
                    let abs_path = self.project_root.resolve(ProjectRelativePath::new(path)?);
                    *path = abs_path.to_string();
                }
            }

            // rewrite the frame ids to be {tid}{frame_id}. starlark doesn't write frame ids > 10000, our thread ids should be <1000.
            frame.id |= (hook.pseudo_thread_id as i64) << 16;
        }
        Ok(trace_response)
    }

    fn scopes(&mut self, x: dap::ScopesArguments) -> anyhow::Result<dap::ScopesResponseBody> {
        let thread_id = x.frame_id >> 16;
        let frame_id = x.frame_id & 0xFFFF;
        if frame_id != 0 {
            return Ok(dap::ScopesResponseBody { scopes: Vec::new() });
        }

        let hook = self.find_hook_by_pseudo_thread(thread_id)?;
        let scopes_info = hook.adapter.scopes()?;
        Ok(dap::ScopesResponseBody {
            scopes: vec![dap::Scope {
                name: "Locals".to_owned(),
                named_variables: Some(scopes_info.num_locals as i64),
                // rewrite variables reference to include our threadid. we don't currently send back any
                // variables other than locals so the TOP_FRAME_LOCALS_ID doesnt really matter (though we
                // do check against it below).
                variables_reference: (thread_id << 16) | TOP_FRAME_LOCALS_ID,
                expensive: false,
                column: None,
                end_column: None,
                end_line: None,
                indexed_variables: None,
                line: None,
                source: None,
            }],
        })
    }

    fn variables(
        &mut self,
        x: dap::VariablesArguments,
    ) -> anyhow::Result<dap::VariablesResponseBody> {
        let thread_id = x.variables_reference >> 16;
        let variables_id = x.variables_reference & 0xFFFF;
        // We only understand the TOP_FRAME_LOCALS_ID id
        if variables_id != TOP_FRAME_LOCALS_ID {
            return Ok(dap::VariablesResponseBody {
                variables: Vec::new(),
            });
        }

        let hook = self.find_hook_by_pseudo_thread(thread_id)?;
        let vars_info = hook.adapter.variables()?;
        Ok(dap::VariablesResponseBody {
            variables: vars_info.locals.into_map(|var| var.to_dap()),
        })
    }

    fn source(&mut self, _x: dap::SourceArguments) -> anyhow::Result<dap::SourceResponseBody> {
        Err(StarlarkDebuggerError::Unimplemented.into())
    }

    fn continue_(&mut self, x: ContinueArguments) -> anyhow::Result<dap::ContinueResponseBody> {
        let hook = self.find_hook_by_pseudo_thread(x.thread_id)?;
        hook.adapter.continue_()?;

        Ok(dap::ContinueResponseBody {
            all_threads_continued: Some(false),
        })
    }

    fn next(&mut self, x: dap::NextArguments) -> anyhow::Result<()> {
        let hook = self.find_hook_by_pseudo_thread(x.thread_id)?;
        hook.adapter.step(StepKind::Over)?;
        Ok(())
    }

    fn step_in(&mut self, x: dap::StepInArguments) -> anyhow::Result<()> {
        let hook = self.find_hook_by_pseudo_thread(x.thread_id)?;
        hook.adapter.step(StepKind::Into)?;
        Ok(())
    }

    fn step_out(&mut self, x: dap::StepOutArguments) -> anyhow::Result<()> {
        let hook = self.find_hook_by_pseudo_thread(x.thread_id)?;
        hook.adapter.step(StepKind::Out)?;
        Ok(())
    }

    fn evaluate(&mut self, x: dap::EvaluateArguments) -> anyhow::Result<dap::EvaluateResponseBody> {
        let frame_id = match x.frame_id {
            Some(v) => v,
            None => {
                return Err(StarlarkDebuggerError::Unimplemented.into());
            }
        };
        let thread_id = frame_id >> 16;
        let frame_id = frame_id & 0xFFFF;
        if frame_id != 0 {
            return Err(StarlarkDebuggerError::Unimplemented.into());
        }

        let hook = self.find_hook_by_pseudo_thread(thread_id)?;
        hook.adapter.evaluate(&x.expression)
    }

    fn disconnect(&mut self, _x: dap::DisconnectArguments) -> anyhow::Result<()> {
        Ok(())
    }
}

impl ServerState {
    fn new(to_client: mpsc::UnboundedSender<ToClientMessage>, project_root: ProjectRoot) -> Self {
        Self {
            to_client,
            project_root,
            current_commands: HashMap::new(),
            current_hooks: HashMap::new(),
            free_pseudo_threads: BTreeSet::new(),
            next_pseudo_thread: 0,
            next_hook_id: HookId(0),
            set_breakpoints: HashMap::new(),
        }
    }

    // The main run loop for the ServerState.
    async fn run(&mut self, recv: mpsc::UnboundedReceiver<ServerMessage>) -> anyhow::Result<()> {
        let mut recv = UnboundedReceiverStream::new(recv);

        self.to_client
            .send(ToClientMessage::Event(dap_event::<()>("initialized", None)))?;

        // We send regular debugger snapshots to all current commands. When we aren't
        // receiving events we send them only once a minute, if there are events they
        // will be reflected in a snapshot within 1sec.
        // As new commands requesting handles are also events, a new command should get
        // its first debugger snapshot within 1s of starting (and can then show UI
        // indicating that a debugger is attached).
        let mut next_snapshot = Instant::now() + Duration::from_secs(1);

        // This loop continues until we encounter an error or get a detach() message.
        loop {
            let timer = tokio::time::sleep_until(next_snapshot);
            select! {
                _ = timer => {
                    let mut snapshot = self.get_snapshot();
                    for (id, state) in &self.current_commands {
                        snapshot.this_handle = id.0;
                        state.events.instant_event(snapshot.clone());
                    }
                    next_snapshot = Instant::now() + Duration::from_secs(60);
                }
                msg = recv.next() => {
                    match msg {
                        None => break,
                        Some(msg) => {
                            if !self.handle_message(msg)? {
                                break;
                            }
                        }
                    }
                    next_snapshot = std::cmp::min(next_snapshot, Instant::now() + Duration::from_secs(1));
                }
            }
        }
        Ok(())
    }

    /// Gets a snapshot with information about all currently stopped starlark evaluations.
    ///
    /// This is sent even when nothing is currently stopped but won't ever be sent if a
    /// debugger isn't attached, and so when a command receives the snapshot it can know
    /// that a debugger is attached to the buck daemon.
    fn get_snapshot(&self) -> buck2_data::DebugAdapterSnapshot {
        let mut current_handles = HashMap::new();

        for hook_state in self.current_hooks.values() {
            if let Some(v) = &hook_state.stopped_at {
                let handle_snapshot = current_handles
                    .entry(hook_state.handle_id.0)
                    .or_insert_with(|| buck2_data::DebugAdapterCommandSnapshot {
                        stopped_evals: Vec::new(),
                    });
                handle_snapshot
                    .stopped_evals
                    .push(buck2_data::DebugAdapterStoppedEval {
                        description: hook_state.pseudo_thread_name.clone(),
                        stopped_at: v.clone(),
                    })
            }
        }

        buck2_data::DebugAdapterSnapshot {
            // this will be changed as appropriate before sending the snapshot.
            this_handle: 0,
            current_handles,
        }
    }

    /// Returns `false` on detach to indicate the state thread should stop running.
    fn handle_message(&mut self, msg: ServerMessage) -> anyhow::Result<bool> {
        match msg {
            ServerMessage::NewHook {
                handle,
                description,
                response_channel,
            } => {
                let resp = self.new_hook(handle, description)?;
                response_channel
                    .send(resp)
                    .map_err(|_| anyhow::anyhow!("channel closed"))?;
            }
            ServerMessage::NewHandle { id, events } => self.new_handle(id, events),
            ServerMessage::DropHook { id } => self.drop_hook(id)?,
            ServerMessage::DropHandle { id } => self.drop_handle(id),
            ServerMessage::DapRequest { req } => {
                let response = match dispatch(self, &req) {
                    Ok(resp) => resp,
                    Err(err) => err_response(&req, &err),
                };
                self.to_client.send(ToClientMessage::Response(response))?;
            }
            ServerMessage::EvalStopped { hook_id } => self.eval_stopped(hook_id)?,
            ServerMessage::Detach => {
                self.detach();
                return Ok(false);
            }
        }
        Ok(true)
    }

    fn next_hook_id(&mut self) -> (HookId, u32) {
        let next = self.next_hook_id;
        self.next_hook_id = HookId(self.next_hook_id.0 + 1);
        let pseudo_thread_id = match self.free_pseudo_threads.pop_first() {
            Some(v) => v,
            None => {
                // if only u32 had fetch_add.
                self.next_pseudo_thread += 1;
                self.next_pseudo_thread - 1
            }
        };

        (next, pseudo_thread_id)
    }

    fn new_hook(
        &mut self,
        handle: BuckStarlarkDebuggerHandle,
        description: String,
    ) -> anyhow::Result<(HookId, Option<Box<dyn DapAdapterEvalHook>>)> {
        let (hook_id, pseudo_thread_id) = self.next_hook_id();

        let client = Box::new(BuckStarlarkDapAdapterClient {
            handle: handle.dupe(),
            hook_id,
        });
        let (hook, eval_wrapper) = prepare_dap_adapter(client);
        let hook_state = HookState {
            adapter: Box::new(hook),
            pseudo_thread_id,
            pseudo_thread_name: description,
            stopped_at: None,
            handle_id: handle.0.id,
        };

        for (source, breakpoints) in &self.set_breakpoints {
            hook_state.adapter.set_breakpoints(source, breakpoints)?;
        }
        self.current_hooks.insert(hook_id, hook_state);

        self.to_client.send(ToClientMessage::Event(dap_event(
            "thread",
            Some(&dap::ThreadEventBody {
                reason: "started".to_owned(),
                thread_id: pseudo_thread_id as i64,
            }),
        )))?;

        Ok((hook_id, Some(Box::new(eval_wrapper))))
    }

    fn new_handle(&mut self, id: HandleId, events: EventDispatcher) {
        self.current_commands.insert(id, CommandState { events });
    }

    fn drop_hook(&mut self, hook_id: HookId) -> anyhow::Result<()> {
        if let Some(state) = self.current_hooks.remove(&hook_id) {
            self.to_client.send(ToClientMessage::Event(dap_event(
                "thread",
                Some(&dap::ThreadEventBody {
                    reason: "exited".to_owned(),
                    thread_id: state.pseudo_thread_id as i64,
                }),
            )))?;
            self.free_pseudo_threads.insert(state.pseudo_thread_id);
        }
        Ok(())
    }

    fn drop_handle(&mut self, handle_id: HandleId) {
        self.current_commands.remove(&handle_id);
    }

    fn eval_stopped(&mut self, hook_id: HookId) -> anyhow::Result<()> {
        debug!("eval stopped {}", hook_id);
        let state = self.current_hooks.get_mut(&hook_id).unwrap();
        let top_frame = state.adapter.top_frame();
        let description = match top_frame {
            Ok(Some(v)) => describe_frame(v),
            _ => "???".to_owned(),
        };
        state.stopped_at = Some(description);
        let thread_id = state.pseudo_thread_id;

        let msg = dap::StoppedEventBody {
            reason: "breakpoint".to_owned(),
            thread_id: Some(thread_id as i64),
            description: Some("Hello".to_owned()),
            all_threads_stopped: Some(false),
            preserve_focus_hint: None,
            text: None,
        };

        self.to_client
            .send(ToClientMessage::Event(dap_event("stopped", Some(&msg))))?;
        Ok(())
    }

    fn detach(&mut self) {
        // Dropping the DapAdapter should make any hooked Evaluator continue freely.
        self.current_hooks.clear();
    }

    fn find_hook_by_pseudo_thread(&self, thread_id: i64) -> anyhow::Result<&HookState> {
        let thread_id = thread_id as u32;
        for hook_state in self.current_hooks.values() {
            if hook_state.pseudo_thread_id == thread_id {
                return Ok(hook_state);
            }
        }
        Err(anyhow::anyhow!("can't find evaluator thread"))
    }

    fn get_ast(&self, source: &ProjectRelativePath) -> anyhow::Result<AstModule> {
        debug!("tried to get ast `{}`", source);
        let abs_path = self.project_root.resolve(source);
        let content = fs_util::read_to_string_opt(abs_path)?.unwrap();
        AstModule::parse(
            &source.to_string(),
            content,
            &Dialect {
                enable_def: true,
                enable_lambda: true,
                enable_load: true,
                enable_keyword_only_arguments: true,
                enable_types: DialectTypes::ParseOnly,
                enable_load_reexport: false,
                enable_top_level_stmt: true,
                ..Dialect::Standard
            },
        )
    }
}

/// Our implementation of starlark's DapAdapterClient. It basically just needs to
/// forward along events to the server with the hook_id (so the server can tell
/// which evaluation the event came from).
#[derive(Debug)]
struct BuckStarlarkDapAdapterClient {
    handle: BuckStarlarkDebuggerHandle,
    hook_id: HookId,
}

impl DapAdapterClient for BuckStarlarkDapAdapterClient {
    fn event_stopped(&self) {
        self.handle.0.server.event_stopped(self.hook_id)
    }
}

/// Information about ongoing commands held by the debugger server.
#[derive(Debug)]
struct CommandState {
    /// The EventDispatcher for the command. This is used to send debugger snapshots.
    events: EventDispatcher,
}

/// Information about ongoing starlark evaluations held by the debugger server.
#[derive(Debug)]
struct HookState {
    /// The starlark DapAdapter used to send commands into the starlark evaluation.
    adapter: Box<dyn DapAdapter>,
    /// The thread id we report to the DAP client for this evaluation. This is unique
    /// at any point in time (no two concurrent evaluations have the same thread id), but
    /// doesn't correlate to the real thread id.
    pseudo_thread_id: u32,
    /// The thread name that we report to the DAP client. This is a user-friendly description
    /// of the starlark operation on the thread (e.g. "load_buildfile:some_cell//some/package").
    pseudo_thread_name: String,
    /// If the evaluation is stopped (for example, at a breakpoint) this is a description of where
    /// it's stopped at. This is used for the debugger snapshots so that buck commands can provide
    /// UI affordances for the stopped evaluations.
    stopped_at: Option<String>,
    /// The id of the corresponding handle (also used for snapshots so a command can tell if a
    /// stopped evaluation is from itself or another command).
    handle_id: HandleId,
}

/// Provides a simple description of a stack frame, typically "<file>:<line>".
fn describe_frame(frame: dap::StackFrame) -> String {
    match frame {
        dap::StackFrame {
            source: Some(dap::Source {
                path: Some(path), ..
            }),
            line,
            ..
        } => {
            format!("{}:{}", path, line)
        }
        _ => "???".to_owned(),
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::io::Write;

use async_trait::async_trait;
use buck2_cli_proto::DapRequest;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::common::ConsoleType;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::events_ctx::PartialResultCtx;
use buck2_client_ctx::events_ctx::PartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::ide_support::ide_message_stream;
use buck2_client_ctx::stream_util::reborrow_stream_for_static;
use buck2_client_ctx::streaming::StreamingCommand;
use buck2_client_ctx::subscribers::subscriber::EventSubscriber;
use buck2_event_observer::unpack_event::unpack_event;
use buck2_event_observer::unpack_event::UnpackedBuckEvent;
use buck2_events::BuckEvent;
use futures::StreamExt;
use once_cell::sync::Lazy;

/// Run the starlark debug adapter protocol server
///
/// This forwards requests received on stdin to a debug server running in the
/// buck daemon. DAP events and responses are returned from the daemon and sent
/// to this command's stdout.
#[derive(Debug, clap::Parser)]
#[clap(name = "starlark-debug-attach")]
pub struct StarlarkDebugAttachCommand {
    #[clap(flatten)]
    config_opts: CommonBuildConfigurationOptions,

    #[clap(flatten)]
    event_log_opts: CommonDaemonCommandOptions,
}

pub fn write_dap_message(out: &mut impl Write, msg: &[u8]) -> anyhow::Result<()> {
    write!(out, "Content-Length: {}\r\n\r\n", msg.len())?;
    out.write_all(msg)?;
    out.flush()?;
    Ok(())
}

/// All DAP messages are written to stdout.
fn send_message_to_dap_client(msg: &[u8]) -> anyhow::Result<()> {
    let stdout = std::io::stdout();
    let mut stdout = stdout.lock();
    write_dap_message(&mut stdout, msg)?;
    Ok(())
}

#[async_trait]
impl StreamingCommand for StarlarkDebugAttachCommand {
    const COMMAND_NAME: &'static str = "starlark-debug-attach";

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let client_context =
            ctx.client_context(&self.config_opts, matches, ctx.sanitized_argv.argv.clone())?;

        let stream = ide_message_stream::<_, debugserver_types::Request>(ctx.stdin()).filter_map(
            |m| async move {
                match m {
                    Ok(dap_json) => Some(DapRequest { dap_json }),
                    Err(e) => {
                        let _ignored = buck2_client_ctx::eprintln!(
                            "Could not read message from stdin: `{}`",
                            e
                        );
                        // TODO(cjhopman): the client just hangs at this point. We should probably error out (or
                        // distinguish between FramedRead errors and errors of us converting to a Request).
                        None
                    }
                }
            },
        );

        let mut partial_result_handler = DapPartialResultHandler;

        reborrow_stream_for_static(
            stream,
            |stream| async move {
                buckd
                    .with_flushing()
                    .dap(client_context, stream, &mut partial_result_handler)
                    .await
            },
            // The DAP server side does not handle hangups. So, until it does... we never hang up:
            || None,
        )
        .await??;

        ExitResult::success()
    }

    fn console_opts(&self) -> &CommonConsoleOptions {
        // This should only be communicated with by an IDE, so disable anything other
        // than the simple console
        static SIMPLE_CONSOLE: Lazy<CommonConsoleOptions> = Lazy::new(|| CommonConsoleOptions {
            console_type: ConsoleType::Simple,
            ui: vec![],
            no_interactive_console: true,
        });
        &SIMPLE_CONSOLE
    }

    fn event_log_opts(&self) -> &CommonDaemonCommandOptions {
        &self.event_log_opts
    }

    fn common_opts(&self) -> &CommonBuildConfigurationOptions {
        &self.config_opts
    }

    fn should_show_waiting_message(&self) -> bool {
        // If we're running the debugger, do not show "Waiting for daemon..." if we do not get any spans.
        false
    }

    fn extra_subscribers(&self) -> Vec<Box<dyn EventSubscriber>> {
        /// We add an additional subscriber that converts a handful of informative events
        /// to DAP "output" events. Without this, at best these would go to stderr, but vscode's
        /// executable DAP client ignores stderr, so this subscriber allows us to get that information
        /// into somewhere visible to the user.

        struct ConvertToDap;

        impl ConvertToDap {
            fn write_console(&self, msg: &str) -> anyhow::Result<()> {
                let ev = debugserver_types::OutputEvent {
                    type_: "event".to_owned(),
                    event: "output".to_owned(),
                    // All other events are being sent by the debug support in the server and that's
                    // maintaining the sequence numbers. For us to get the correct sequence number
                    // here would be tricky. Instead, we just set it to 0 and hope that nobody notices/cares
                    // that it's out of order/invalid. The alternative would probably be to
                    // deserialize all events from the server and rewrite their sequence numbers (and
                    // potentially references to those sequence numbers coming back from the dap client).
                    seq: 0,
                    body: debugserver_types::OutputEventBody {
                        category: None,
                        column: None,
                        data: None,
                        line: None,
                        output: format!("{}\n", msg),
                        source: None,
                        variables_reference: None,
                    },
                };
                send_message_to_dap_client(&serde_json::to_vec(&ev)?)
            }
        }

        #[async_trait]
        impl EventSubscriber for ConvertToDap {
            async fn handle_output(&mut self, raw_output: &[u8]) -> anyhow::Result<()> {
                self.write_console(&String::from_utf8_lossy(raw_output))
            }

            async fn handle_tailer_stderr(&mut self, stderr: &str) -> anyhow::Result<()> {
                self.write_console(stderr)
            }

            async fn handle_events(
                &mut self,
                events: &[std::sync::Arc<BuckEvent>],
            ) -> anyhow::Result<()> {
                for ev in events {
                    match unpack_event(ev)? {
                        UnpackedBuckEvent::Instant(_, _, data) => match data {
                            buck2_data::instant_event::Data::StructuredError(soft_error) => {
                                if !soft_error.quiet {
                                    self.write_console(&format!(
                                        "soft error: {}",
                                        &soft_error.payload
                                    ))?;
                                }
                            }
                            buck2_data::instant_event::Data::ConsoleMessage(message) => {
                                self.write_console(&message.message)?;
                            }
                            _ => {}
                        },
                        _ => {}
                    }
                }
                Ok(())
            }

            async fn handle_error(&mut self, error: &anyhow::Error) -> anyhow::Result<()> {
                self.write_console(&format!(
                    "buck2 starlark-attach debugserver error: {}",
                    error
                ))
            }
        }

        vec![Box::new(ConvertToDap)]
    }
}

struct DapPartialResultHandler;

#[async_trait]
impl PartialResultHandler for DapPartialResultHandler {
    type PartialResult = buck2_cli_proto::DapMessage;

    async fn handle_partial_result(
        &mut self,
        mut _ctx: PartialResultCtx<'_, '_>,
        partial_res: buck2_cli_proto::DapMessage,
    ) -> anyhow::Result<()> {
        send_message_to_dap_client(&partial_res.dap_json)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

#![feature(async_closure)]
#![feature(try_blocks)]

use async_trait::async_trait;
use buck2_cli_proto::ClientContext;
use buck2_cli_proto::GenericRequest;
use buck2_client_ctx::argv::Argv;
use buck2_client_ctx::argv::SanitizedArgv;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::daemon::client::StdoutPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::BuckSubcommand;
use buck2_client_ctx::streaming::StreamingCommand;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;

use crate::debug::StarlarkDebugAttachCommand;
use crate::lint::StarlarkLintCommand;
use crate::typecheck::StarlarkTypecheckCommand;

mod debug;
mod lint;
mod oracle_buck;
pub mod server;
mod typecheck;
mod util;

#[derive(Debug, clap::Subcommand)]
#[clap(name = "starlark", about = "Run Starlark operations")]
pub enum StarlarkCommand {
    #[clap(flatten)]
    Opaque(StarlarkOpaqueCommand),
    DebugAttach(StarlarkDebugAttachCommand),
}

// Used for subcommands that follow `buck2 audit`'s "opaque" pattern where the command object is serialized
// to the daemon and deserialized there and has a `server_execute()` on the Command object itself (as opposed
// to using structured endpoints in the daemon protocol).
#[derive(Debug, clap::Subcommand, serde::Serialize, serde::Deserialize)]
pub enum StarlarkOpaqueCommand {
    Lint(StarlarkLintCommand),
    Typecheck(StarlarkTypecheckCommand),
}

#[derive(Debug, clap::Parser, serde::Serialize, serde::Deserialize, Default)]
pub struct StarlarkCommandCommonOptions {
    #[clap(flatten)]
    config_opts: CommonBuildConfigurationOptions,

    #[clap(flatten)]
    console_opts: CommonConsoleOptions,

    #[clap(flatten)]
    event_log_opts: CommonDaemonCommandOptions,
}

#[async_trait]
pub trait StarlarkOpaqueSubcommand: Send + Sync + 'static {
    async fn server_execute(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        stdout: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        client_server_ctx: ClientContext,
    ) -> anyhow::Result<()>;

    fn common_opts(&self) -> &StarlarkCommandCommonOptions;
}

impl StarlarkOpaqueCommand {
    pub async fn server_execute(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        stdout: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        client_server_ctx: ClientContext,
    ) -> anyhow::Result<()> {
        self.as_subcommand()
            .server_execute(server_ctx, stdout, client_server_ctx)
            .await
    }
    fn as_subcommand(&self) -> &dyn StarlarkOpaqueSubcommand {
        match self {
            Self::Lint(cmd) => cmd,
            Self::Typecheck(cmd) => cmd,
        }
    }
}

#[async_trait]
impl StreamingCommand for StarlarkOpaqueCommand {
    const COMMAND_NAME: &'static str = "starlark";

    /// Starlark subcommands are all implemented as a generic request to the buckd server that will deserialize the command object.
    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let serialized = serde_json::to_string(&self)?;

        let config_opts = &self.as_subcommand().common_opts().config_opts;
        let context = ctx.client_context(config_opts, matches, ctx.sanitized_argv.argv.clone())?;

        buckd
            .with_flushing()
            .starlark(
                GenericRequest {
                    context: Some(context),
                    serialized_opts: serialized,
                },
                ctx.stdin().console_interaction_stream(self.console_opts()),
                &mut StdoutPartialResultHandler,
            )
            .await??;
        ExitResult::success()
    }

    fn console_opts(&self) -> &CommonConsoleOptions {
        &self.as_subcommand().common_opts().console_opts
    }

    fn event_log_opts(&self) -> &CommonDaemonCommandOptions {
        &self.as_subcommand().common_opts().event_log_opts
    }

    fn common_opts(&self) -> &CommonBuildConfigurationOptions {
        &self.as_subcommand().common_opts().config_opts
    }
}

impl StarlarkCommand {
    pub fn exec(self, matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let matches = matches.subcommand().expect("subcommand not found").1;
        match self {
            StarlarkCommand::Opaque(cmd) => cmd.exec(matches, ctx),
            StarlarkCommand::DebugAttach(cmd) => cmd.exec(matches, ctx),
        }
    }

    pub fn sanitize_argv(&self, argv: Argv) -> SanitizedArgv {
        argv.no_need_to_sanitize()
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::HashMap;
use std::collections::HashSet;
use std::io::Write;
use std::sync::Arc;

use anyhow::Context;
use async_trait::async_trait;
use buck2_cli_proto::ClientContext;
use buck2_client_ctx::path_arg::PathArg;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::dice::data::HasIoProvider;
use buck2_common::dice::file_ops::HasFileOps;
use buck2_common::io::IoProvider;
use buck2_core::cells::name::CellName;
use buck2_core::cells::CellResolver;
use buck2_interpreter::file_type::StarlarkFileType;
use buck2_interpreter::path::StarlarkPath;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::ctx::ServerCommandDiceContext;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use dice::DiceTransaction;
use dupe::Dupe;
use starlark::codemap::FileSpan;
use starlark::errors::Diagnostic;
use starlark::errors::Lint;
use starlark::syntax::AstModule;

use crate::util::environment::Environment;
use crate::util::paths::starlark_files;
use crate::StarlarkCommandCommonOptions;
use crate::StarlarkOpaqueSubcommand;

#[derive(Debug, clap::Parser, serde::Serialize, serde::Deserialize)]
#[clap(name = "starlark-lint", about = "Run the Starlark linter.")]
pub struct StarlarkLintCommand {
    #[clap(flatten)]
    common_opts: StarlarkCommandCommonOptions,

    #[clap(value_name = "PATH", required = true)]
    paths: Vec<PathArg>,
}

/// The cache of names for a path, keyed by its CellName and its path type.
struct Cache<'a> {
    dice: &'a DiceTransaction,
    cached: HashMap<(CellName, StarlarkFileType), Arc<HashSet<String>>>,
}

impl<'a> Cache<'a> {
    pub(crate) fn new(dice: &'a DiceTransaction) -> Cache<'a> {
        Self {
            dice,
            cached: HashMap::new(),
        }
    }

    pub(crate) async fn get_names(
        &mut self,
        path: &StarlarkPath<'_>,
    ) -> anyhow::Result<Arc<HashSet<String>>> {
        let path_type = path.file_type();
        let cell = path.cell();
        if let Some(res) = self.cached.get(&(cell, path_type)) {
            return Ok(res.dupe());
        }
        let env: Environment = Environment::new(cell, path_type, self.dice).await?;
        let res = Arc::new(env.get_names(path_type, self.dice).await?);
        self.cached.insert((cell, path_type), res.dupe());
        Ok(res)
    }
}

async fn lint_file(
    path: &StarlarkPath<'_>,
    cell_resolver: &CellResolver,
    io: &dyn IoProvider,
    cache: &mut Cache<'_>,
) -> anyhow::Result<Vec<Lint>> {
    let dialect = path.file_type().dialect(false);
    let proj_path = cell_resolver.resolve_path(path.path().as_ref().as_ref())?;
    let path_str = proj_path.to_string();
    let content = io
        .read_file_if_exists(proj_path)
        .await?
        .with_context(|| format!("File not found: `{}`", path_str))?;
    match AstModule::parse(&path_str, content.clone(), &dialect) {
        Ok(ast) => Ok(ast.lint(Some(&*cache.get_names(path).await?))),
        Err(err) => {
            // There was a parse error, so we don't want to fail, we want to give a nice error message
            // Do the best we can - it is probably a `Diagnostic`, which gives us more precise info.
            let (span, message) = match err.downcast::<Diagnostic>() {
                Err(err) => (None, err),
                Ok(diag) => (diag.span, diag.message),
            };
            Ok(vec![Lint {
                location: span.unwrap_or_else(|| FileSpan::new(path_str, content)),
                short_name: "parse_error".to_owned(),
                serious: true,
                problem: format!("{:#}", message),
                original: "".to_owned(),
            }])
        }
    }
}

#[async_trait]
impl StarlarkOpaqueSubcommand for StarlarkLintCommand {
    async fn server_execute(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        mut stdout: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        _client_ctx: ClientContext,
    ) -> anyhow::Result<()> {
        server_ctx
            .with_dice_ctx(async move |server_ctx, ctx| {
                let cell_resolver = ctx.get_cell_resolver().await?;
                let fs = ctx.file_ops();
                let io = ctx.global_data().get_io_provider();
                let mut cache = Cache::new(&ctx);

                let mut stdout = stdout.as_writer();
                let mut lint_count = 0;
                let files =
                    starlark_files(&self.paths, server_ctx, &cell_resolver, &fs, &*io).await?;
                for file in &files {
                    let lints = lint_file(&file.borrow(), &cell_resolver, &*io, &mut cache).await?;
                    lint_count += lints.len();
                    for lint in lints {
                        writeln!(stdout, "{}", lint)?;
                    }
                }
                if lint_count > 0 {
                    Err(anyhow::anyhow!("Found {} lints", lint_count))
                } else {
                    writeln!(
                        server_ctx.stderr()?,
                        "Found no lints in {} files",
                        files.len()
                    )?;
                    Ok(())
                }
            })
            .await
    }

    fn common_opts(&self) -> &StarlarkCommandCommonOptions {
        &self.common_opts
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::sync::Arc;

use starlark::docs::get_registered_starlark_docs;
use starlark::environment::Globals;
use starlark::environment::LibraryExtension;
use starlark::typing::*;

pub(crate) fn oracle_buck(globals: &Globals) -> Arc<dyn TypingOracle + Send + Sync> {
    let registered_docs = get_registered_starlark_docs();
    let mut docs = OracleDocs::new();
    docs.add_module(&globals.documentation());
    docs.add_docs(&registered_docs);

    let mut docs2 = OracleDocs::new();
    docs2.add_docs(&registered_docs);

    Arc::new(OracleSeq(vec![
        // TODO(nga): this should only accept globals from enabled extensions, not all extensions.
        Box::new(OracleStandard::new(LibraryExtension::all()))
            as Box<dyn TypingOracle + Send + Sync>,
        Box::new(CustomBuck),
        Box::new(docs),
        Box::new(AddErrors(docs2)),
    ]))
}

struct CustomBuck;

impl TypingOracle for CustomBuck {
    fn attribute(&self, ty: &Ty, attr: TypingAttr) -> Option<Result<Ty, ()>> {
        match ty.as_name()? {
            "dependency" => match attr {
                // We can index by providers, which either appear as functions (for builtin providers)
                // or as Any (for user providers)
                TypingAttr::Index => {
                    Some(Ok(Ty::function(vec![Param::pos_only(Ty::Any)], Ty::Any)))
                }
                TypingAttr::BinOp(TypingBinOp::In) => {
                    Some(Ok(Ty::function(vec![Param::pos_only(Ty::Any)], Ty::bool())))
                }
                _ => None,
            },
            _ => None,
        }
    }

    fn as_function(&self, ty: &TyName) -> Option<Result<TyFunction, ()>> {
        match ty.as_str() {
            "provider_callable" => Some(Ok(TyFunction {
                type_attr: "".to_owned(),
                params: vec![Param::kwargs(Ty::Any)],
                // TODO(nga): this should be more precise.
                result: Box::new(Ty::Any),
            })),
            "rule" => Some(Ok(TyFunction {
                type_attr: "".to_owned(),
                params: vec![Param::kwargs(Ty::Any)],
                result: Box::new(Ty::none()),
            })),
            _ => None,
        }
    }

    fn subtype(&self, require: &TyName, got: &TyName) -> bool {
        match require.as_str() {
            "provider" => got.as_str().ends_with("Info"),
            _ => false,
        }
    }
}

struct AddErrors(OracleDocs);

impl TypingOracle for AddErrors {
    fn attribute(&self, ty: &Ty, _attr: TypingAttr) -> Option<Result<Ty, ()>> {
        if self.0.known_object(ty.as_name()?) {
            Some(Err(()))
        } else {
            None
        }
    }
    fn as_function(&self, _ty: &TyName) -> Option<Result<TyFunction, ()>> {
        Some(Err(()))
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::HashMap;

use buck2_events::dispatch::span_async;
use buck2_server_ctx::command_end::command_end;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;

use crate::StarlarkOpaqueCommand;

pub async fn server_starlark_command(
    ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    req: buck2_cli_proto::GenericRequest,
) -> anyhow::Result<buck2_cli_proto::GenericResponse> {
    let metadata = ctx.request_metadata().await?;
    let start_event = buck2_data::CommandStart {
        metadata: metadata.clone(),
        data: Some(buck2_data::StarlarkCommandStart {}.into()),
    };

    span_async(
        start_event,
        server_starlark_command_inner(metadata, ctx, partial_result_dispatcher, req),
    )
    .await
}

async fn server_starlark_command_inner(
    metadata: HashMap<String, String>,
    context: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    req: buck2_cli_proto::GenericRequest,
) -> (
    anyhow::Result<buck2_cli_proto::GenericResponse>,
    buck2_data::CommandEnd,
) {
    let args = req.serialized_opts.to_owned();
    let dir = context.working_dir().to_string();
    let result = parse_command_and_execute(context, partial_result_dispatcher, req).await;
    let end_event = command_end(
        metadata,
        &result,
        buck2_data::StarlarkCommandEnd {
            status: if result.is_ok() { 0 } else { 1 },
            args,
            dir,
        },
    );

    let result = result.map(|()| buck2_cli_proto::GenericResponse {});

    (result, end_event)
}

async fn parse_command_and_execute(
    context: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    req: buck2_cli_proto::GenericRequest,
) -> anyhow::Result<()> {
    let command: StarlarkOpaqueCommand = serde_json::from_str(&req.serialized_opts)?;
    command
        .server_execute(
            context,
            partial_result_dispatcher,
            req.context.expect("buck cli always sets a client context"),
        )
        .await
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::HashMap;
use std::io::Write;
use std::sync::Arc;

use anyhow::Context;
use async_recursion::async_recursion;
use async_trait::async_trait;
use buck2_cli_proto::ClientContext;
use buck2_client_ctx::path_arg::PathArg;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::dice::data::HasIoProvider;
use buck2_common::dice::file_ops::HasFileOps;
use buck2_common::io::IoProvider;
use buck2_core::cells::name::CellName;
use buck2_core::cells::CellResolver;
use buck2_interpreter::file_type::StarlarkFileType;
use buck2_interpreter::path::OwnedStarlarkPath;
use buck2_interpreter_for_build::interpreter::dice_calculation_delegate::HasCalculationDelegate;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::ctx::ServerCommandDiceContext;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use dice::DiceTransaction;
use dupe::Dupe;
use starlark::environment::Globals;
use starlark::typing::Interface;
use starlark::typing::TypingOracle;

use crate::oracle_buck::oracle_buck;
use crate::util::environment::Environment;
use crate::util::paths::starlark_files;
use crate::StarlarkCommandCommonOptions;
use crate::StarlarkOpaqueSubcommand;

#[derive(Debug, clap::Parser, serde::Serialize, serde::Deserialize)]
#[clap(name = "starlark-typecheck", about = "Run the Starlark typechecker.")]
pub struct StarlarkTypecheckCommand {
    #[clap(flatten)]
    common_opts: StarlarkCommandCommonOptions,

    #[clap(value_name = "PATH", required = true)]
    paths: Vec<PathArg>,
}

struct Cache<'a> {
    // Things we have access to get information
    dice: &'a DiceTransaction,
    io: &'a dyn IoProvider,
    cell_resolver: &'a CellResolver,
    // Things we have access to write information
    stdout: &'a mut (dyn Write + Send + Sync),
    stderr: &'a mut (dyn Write + Send + Sync),
    // Our accumulated state
    oracle: HashMap<(CellName, StarlarkFileType), (Arc<dyn TypingOracle + Send + Sync>, Globals)>,
    cache: HashMap<OwnedStarlarkPath, Interface>,
}

impl<'a> Cache<'a> {
    async fn typecheck(&mut self, path: OwnedStarlarkPath) -> anyhow::Result<()> {
        self.run(path).await?;
        Ok(())
    }

    async fn get_oracle(
        &mut self,
        cell: CellName,
        path_type: StarlarkFileType,
    ) -> anyhow::Result<(Arc<dyn TypingOracle + Send + Sync>, Globals)> {
        match self.oracle.get(&(cell, path_type)) {
            Some((o, g)) => Ok((o.dupe(), g.dupe())),
            None => {
                let globals = Environment::new(cell, path_type, self.dice).await?.globals;
                let res = oracle_buck(&globals);
                self.oracle
                    .insert((cell, path_type), (res.dupe(), globals.dupe()));
                Ok((res, globals))
            }
        }
    }

    async fn get(&mut self, path: OwnedStarlarkPath) -> anyhow::Result<Interface> {
        match self.cache.get(&path) {
            Some(x) => Ok(x.dupe()),
            None => {
                let res = self.run(path.clone()).await?;
                self.cache.insert(path, res.dupe());
                Ok(res)
            }
        }
    }

    #[async_recursion]
    async fn run(&mut self, path: OwnedStarlarkPath) -> anyhow::Result<Interface> {
        let path_ref = path.borrow();
        writeln!(self.stderr, "Type checking: {path_ref}")?;
        let proj_path = self
            .cell_resolver
            .resolve_path(path_ref.path().as_ref().as_ref())?;
        let path_str = proj_path.to_string();
        let src = self
            .io
            .read_file_if_exists(proj_path)
            .await?
            .with_context(|| format!("File not found: `{path_str}`"))?;

        let interp = self
            .dice
            .get_interpreter_calculator(path_ref.cell(), path_ref.build_file_cell())
            .await?;

        let ast = interp.prepare_eval_with_content(path_ref, src)?;
        let mut loads = HashMap::new();
        for x in ast.loads() {
            let y = interp.resolve_load(path_ref, x.module_id).await?;
            let interface = self.get(y.into_starlark_path()).await?;
            loads.insert(x.module_id.to_owned(), interface);
        }
        let (oracle, globals) = self
            .get_oracle(path_ref.cell(), path_ref.file_type())
            .await?;
        let (errors, bindings, interface, approxiomations) =
            ast.typecheck(&*oracle, &globals, &loads);

        if !approxiomations.is_empty() {
            writeln!(self.stderr, "\n\nAPPROXIMATIONS:")?;
            for x in approxiomations {
                writeln!(self.stderr, "{x}")?;
            }
        }

        writeln!(self.stderr, "\n\nBINDINGS:\n{bindings}")?;

        let errors_count = errors.len();
        if errors_count == 0 {
            Ok(interface)
        } else {
            writeln!(self.stdout, "\n\nERRORS:")?;
            for x in errors {
                writeln!(self.stdout, "{x}")?;
            }
            Err(anyhow::anyhow!("Detected {errors_count} errors"))
        }
    }
}

#[async_trait]
impl StarlarkOpaqueSubcommand for StarlarkTypecheckCommand {
    async fn server_execute(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        mut stdout: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        _client_ctx: ClientContext,
    ) -> anyhow::Result<()> {
        server_ctx
            .with_dice_ctx(async move |server_ctx, dice| {
                let cell_resolver = dice.get_cell_resolver().await?;
                let fs = dice.file_ops();
                let io = dice.global_data().get_io_provider();

                let files =
                    starlark_files(&self.paths, server_ctx, &cell_resolver, &fs, &*io).await?;
                let mut stdout = stdout.as_writer();
                let mut stderr = server_ctx.stderr()?;
                let mut cache = Cache {
                    dice: &dice,
                    io: &*io,
                    cell_resolver: &cell_resolver,
                    stdout: &mut stdout,
                    stderr: &mut stderr,
                    oracle: HashMap::new(),
                    cache: HashMap::new(),
                };
                for file in files {
                    cache.typecheck(file).await?;
                }
                let file_count = cache.cache.len();
                writeln!(stderr, "Found no type errors in {file_count} files")?;
                Ok(())
            })
            .await
    }

    fn common_opts(&self) -> &StarlarkCommandCommonOptions {
        &self.common_opts
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::HashSet;

use buck2_core::bzl::ImportPath;
use buck2_core::cells::build_file_cell::BuildFileCell;
use buck2_core::cells::name::CellName;
use buck2_interpreter::file_type::StarlarkFileType;
use buck2_interpreter::import_paths::HasImportPaths;
use buck2_interpreter::load_module::InterpreterCalculation;
use buck2_interpreter::load_module::INTERPRETER_CALCULATION_IMPL;
use buck2_interpreter::prelude_path::PreludePath;
use dice::DiceTransaction;
use starlark::environment::Globals;

/// The environment in which a Starlark file is evaluated.
pub(crate) struct Environment {
    /// The globals that are driven from Rust.
    pub(crate) globals: Globals,
    /// The path to the prelude, if the prelude is loaded in this file.
    /// Note that in a BUCK file the `native` value is also exploded into the top-level.
    prelude: Option<PreludePath>,
    /// A path that is implicitly loaded as additional globals.
    preload: Option<ImportPath>,
}

impl Environment {
    pub(crate) async fn new(
        cell: CellName,
        path_type: StarlarkFileType,
        dice: &DiceTransaction,
    ) -> anyhow::Result<Environment> {
        // Find the information from the globals
        let globals = INTERPRETER_CALCULATION_IMPL
            .get()?
            .global_env_for_file_type(dice, path_type)
            .await?;

        // Next grab the prelude, unless we are in the prelude cell and not a build file
        let prelude = match INTERPRETER_CALCULATION_IMPL
            .get()?
            .prelude_import(dice)
            .await?
        {
            Some(prelude)
                if path_type == StarlarkFileType::Buck || prelude.import_path().cell() != cell =>
            {
                Some(prelude)
            }
            _ => None,
        };

        // Now grab the pre-load things
        let preload = dice
            .import_paths_for_cell(BuildFileCell::new(cell))
            .await?
            .root_import()
            .cloned();

        Ok(Environment {
            globals,
            prelude,
            preload,
        })
    }

    pub(crate) async fn get_names(
        &self,
        path_type: StarlarkFileType,
        dice: &DiceTransaction,
    ) -> anyhow::Result<HashSet<String>> {
        let mut names = HashSet::new();

        for x in self.globals.names() {
            names.insert(x.as_str().to_owned());
        }

        if let Some(prelude) = &self.prelude {
            let m = dice
                .get_loaded_module_from_import_path(prelude.import_path())
                .await?;
            for x in m.env().names() {
                names.insert(x.as_str().to_owned());
            }
            if path_type == StarlarkFileType::Buck {
                if let Some(native) = m.env().get_option("native")? {
                    let native = native.value();
                    for attr in native.dir_attr() {
                        names.insert(attr.to_owned());
                    }
                }
            }
        }

        if let Some(preload) = &self.preload {
            let m = dice.get_loaded_module_from_import_path(preload).await?;
            for x in m.env().names() {
                names.insert(x.as_str().to_owned());
            }
        }

        Ok(names)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

pub(crate) mod environment;
pub(crate) mod paths;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::ops::Deref;

use async_recursion::async_recursion;
use buck2_client_ctx::path_arg::PathArg;
use buck2_common::file_ops::FileOps;
use buck2_common::file_ops::FileType;
use buck2_common::file_ops::RawPathMetadata;
use buck2_common::io::IoProvider;
use buck2_core::build_file_path::BuildFilePath;
use buck2_core::bzl::ImportPath;
use buck2_core::cells::CellResolver;
use buck2_core::fs::paths::file_name::FileName;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_core::package::PackageLabel;
use buck2_interpreter::path::BxlFilePath;
use buck2_interpreter::path::OwnedStarlarkPath;
use buck2_interpreter::path::PackageFilePath;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use dupe::Dupe;
use thiserror::Error;

#[derive(Debug, Error)]
enum StarlarkFilesError {
    #[error("File not found, `{0}`")]
    FileNotFound(ProjectRelativePathBuf),
    #[error("Symlinks and other esoteric files are not supported, `{0}`")]
    UnsupportedFileType(ProjectRelativePathBuf),
}

#[async_recursion]
async fn starlark_file(
    proj_path: ProjectRelativePathBuf,
    // None = this file was given explicitly
    // Some = it was a directory traversal (and we know its type)
    recursive: Option<FileType>,
    cell_resolver: &CellResolver,
    fs: &dyn FileOps,
    io: &dyn IoProvider,
    files: &mut Vec<OwnedStarlarkPath>,
) -> anyhow::Result<()> {
    let cell_path = cell_resolver.get_cell_path(&proj_path)?;
    if recursive.is_some() && fs.is_ignored(cell_path.as_ref()).await? {
        // File is ignored by Buck, give up on it
        return Ok(());
    }

    let typ = match &recursive {
        Some(typ) => typ.dupe(),
        None => match io.read_path_metadata_if_exists(proj_path.clone()).await? {
            None => {
                return Err(StarlarkFilesError::FileNotFound(proj_path).into());
            }
            Some(RawPathMetadata::Directory) => FileType::Directory,
            Some(RawPathMetadata::File(_)) => {
                // It's a shame we throw away the digest we calculated, but not a huge deal (its cheap compared to parsing)
                FileType::File
            }
            Some(RawPathMetadata::Symlink { .. }) => FileType::Symlink,
        },
    };

    match typ {
        FileType::Directory => {
            for x in io.read_dir(proj_path.clone()).await? {
                let Ok(file_name) = FileName::new(&x.file_name) else {
                    // Skip files which buck does not like:
                    // this function works with `CellPath` values,
                    // which cannot be constructed from paths not acceptable by buck.
                    continue;
                };
                let mut child_path = proj_path.clone();
                child_path.push(file_name);
                starlark_file(child_path, Some(x.file_type), cell_resolver, fs, io, files).await?;
            }
        }
        FileType::File => {
            // It's a shame we throw away the digest we calculated, but not a huge deal (its cheap compared to parsing)
            let is_buildfile = match proj_path.file_name() {
                None => false,
                Some(file_name) => cell_resolver
                    .get(cell_path.cell())?
                    .buildfiles()
                    .iter()
                    .any(|x| (*x).deref() == file_name),
            };

            if is_buildfile {
                files.push(OwnedStarlarkPath::BuildFile(BuildFilePath::new(
                    PackageLabel::from_cell_path(cell_path.parent().unwrap()),
                    proj_path.file_name().unwrap().to_owned(),
                )));
            } else if proj_path.as_str().ends_with(".bxl") {
                files.push(OwnedStarlarkPath::BxlFile(BxlFilePath::new(cell_path)?));
            } else if proj_path.ends_with(PackageFilePath::PACKAGE_FILE_NAME) {
                // `parent` must return `Some` if we have a non-empty file, which we do because of above
                files.push(OwnedStarlarkPath::PackageFile(PackageFilePath::for_dir(
                    cell_path.parent().unwrap(),
                )))
            } else if recursive.is_none() || proj_path.as_str().ends_with(".bzl") {
                // If a file was asked for explicitly, and is nothing else, treat it as .bzl file
                // If it's not explicit, just ignore it (probably a source file)
                files.push(OwnedStarlarkPath::LoadFile(ImportPath::new_same_cell(
                    cell_path,
                )?));
            }
        }
        FileType::Symlink | FileType::Unknown => {
            if recursive.is_none() {
                return Err(StarlarkFilesError::UnsupportedFileType(proj_path).into());
            }
        }
    }
    Ok(())
}

/// Find the paths to apply Starlark to (e.g. linter, typecheck)
pub(crate) async fn starlark_files(
    paths: &[PathArg],
    context: &dyn ServerCommandContextTrait,
    cell_resolver: &CellResolver,
    fs: &dyn FileOps,
    io: &dyn IoProvider,
) -> anyhow::Result<Vec<OwnedStarlarkPath>> {
    let mut files = Vec::new();

    for path in paths {
        let path = path.resolve(context.working_dir_abs());
        let cell_path = cell_resolver.get_cell_path_from_abs_path(&path, context.project_root())?;
        let proj_path = cell_resolver.resolve_path(cell_path.as_ref())?;
        starlark_file(proj_path, None, cell_resolver, fs, io, &mut files).await?;
    }
    Ok(files)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::io;

fn main() -> io::Result<()> {
    let proto_files = &["subscription.proto"];

    buck2_protoc_dev::configure()
        .setup_protoc()
        .type_attribute(".", "#[derive(::serde::Serialize, ::serde::Deserialize)]")
        .type_attribute(".", "#[derive(::allocative::Allocative)]")
        .type_attribute(
            "buck.subscription.SubscriptionRequest.request",
            "#[derive(::derive_more::From)]",
        )
        .type_attribute(
            "buck.subscription.SubscriptionResponse.response",
            "#[derive(::derive_more::From)]",
        )
        .compile(proto_files, &["."])
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

tonic::include_proto!("buck.subscription");