//! The workspace's side of remote connections: serving this device's
//! projects and threads to other devices, and showing other devices'
//! projects and threads in the sidebar.

use std::{
    collections::{HashMap, HashSet},
    hash::{Hash as _, Hasher as _},
    path::PathBuf,
    rc::Rc,
    time::{Duration, SystemTime},
};

use gpui_kit::component::{WindowExt as _, notification::Notification};
use gpui_kit::*;

use super::{Project, ThreadSettings, Workspace, hostname, initials, tint_for};
use crate::claude::PermissionMode;
use crate::git::GitInfo;
use crate::models::Provider;
use crate::remote::{
    ClientMessage, ProjectInfo, ServerMessage, ThreadAction, ThreadSummary,
    client::{ClientEvent, ConnectionStatus, RemoteClient},
    new_key,
    server::{Peer, RemoteServer, ServerEvent},
};
use crate::session::{
    NewSession, PermissionDecision, Reasoning, RemoteLink, Session, from_millis, store_uploads,
};
use crate::settings::{self, RemoteHost, WorkspaceMode};

/// How often changed threads are sent to subscribed clients.
pub(super) const FLUSH: Duration = Duration::from_millis(150);

/// A connection to another device's app.
pub(super) struct Environment {
    pub(super) host: RemoteHost,
    pub(super) status: ConnectionStatus,
    pub(super) client: Rc<RemoteClient>,
    _events: Task<()>,
}

/// One client's interest in a thread: the hashes of the transcript items it
/// was last sent, or `None` before the first send.
type Subscription = Option<Vec<u64>>;

#[derive(Default)]
pub(super) struct RemoteState {
    server: Option<RemoteServer>,
    /// The port and key the running server was started with.
    server_config: Option<(u16, String)>,
    server_error: Option<String>,
    peers: Vec<Peer>,
    subscriptions: HashMap<u64, HashMap<String, Subscription>>,
    dirty_threads: HashSet<String>,
    state_dirty: bool,
    last_state: String,
    pub(super) environments: Vec<Environment>,
    _server_events: Option<Task<()>>,
}

/// What the settings show about remote connections.
pub struct RemoteOverview {
    /// The port this device listens on, when it does.
    pub serving: Option<u16>,
    pub error: Option<String>,
    pub peers: Vec<Peer>,
    /// Each connection to another device, by its id in the settings.
    pub environments: Vec<(String, ConnectionStatus)>,
}

impl Workspace {
    /// Starts, stops or restarts the server and the connections to match the
    /// settings.
    pub(super) fn sync_remote(&mut self, window: &mut Window, cx: &mut Context<Self>) {
        let (serve, port, mut key, hosts) = settings::get(|settings| {
            (settings.remote_server, settings.remote_port, settings.remote_key.clone(), settings.remotes.clone())
        });
        if serve && key.is_empty() {
            key = new_key();
            let saved = key.clone();
            settings::update(|settings| settings.remote_key = saved);
        }
        let wanted = serve.then(|| (port, key));
        if wanted != self.remote.server_config {
            self.remote.server = None;
            self.remote._server_events = None;
            self.remote.peers.clear();
            self.remote.subscriptions.clear();
            self.remote.server_error = None;
            self.remote.server_config = wanted.clone();
            if let Some((port, key)) = wanted {
                match RemoteServer::start(port, key, hostname()) {
                    Ok(server) => {
                        let events = server.events.clone();
                        self.remote._server_events = Some(cx.spawn_in(window, async move |this, cx| {
                            while let Ok(event) = events.recv().await {
                                if this.update_in(cx, |this, window, cx| this.on_server_event(event, window, cx)).is_err() {
                                    break;
                                }
                            }
                        }));
                        self.remote.server = Some(server);
                        self.remote.state_dirty = true;
                    }
                    Err(err) => self.remote.server_error = Some(format!("Couldn't listen on port {port}: {err}")),
                }
            }
        }

        let wanted: Vec<RemoteHost> = hosts.into_iter().filter(|host| host.enabled && !host.address.trim().is_empty()).collect();
        let stale: Vec<String> = self
            .remote
            .environments
            .iter()
            .filter(|env| !wanted.iter().any(|host| *host == env.host))
            .map(|env| env.host.id.clone())
            .collect();
        for id in stale {
            self.drop_environment(&id, window, cx);
        }
        for host in wanted {
            if !self.remote.environments.iter().any(|env| env.host.id == host.id) {
                self.connect_environment(host, window, cx);
            }
        }
        cx.notify();
    }

    fn connect_environment(&mut self, host: RemoteHost, window: &mut Window, cx: &mut Context<Self>) {
        let client = Rc::new(RemoteClient::connect(&host.address, host.key.clone(), hostname()));
        let events = client.events.clone();
        let id = host.id.clone();
        let task = cx.spawn_in(window, async move |this, cx| {
            while let Ok(event) = events.recv().await {
                let id = id.clone();
                if this.update_in(cx, |this, window, cx| this.on_client_event(&id, event, window, cx)).is_err() {
                    break;
                }
            }
        });
        self.remote.environments.push(Environment { host, status: ConnectionStatus::Connecting, client, _events: task });
    }

    /// Disconnects from a device and takes its projects and threads away.
    fn drop_environment(&mut self, id: &str, window: &mut Window, cx: &mut Context<Self>) {
        self.remote.environments.retain(|env| env.host.id != id);
        let removed: Vec<Entity<Session>> = self
            .projects
            .iter()
            .filter(|project| project.environment.as_deref() == Some(id))
            .flat_map(|project| project.sessions.clone())
            .collect();
        for session in &removed {
            if let Some(ix) = self.project_of(session) {
                self.remove_thread(ix, session, cx);
            }
        }
        let active = self.active.clone();
        self.projects.retain(|project| project.environment.as_deref() != Some(id));
        self.active = active.filter(|session| self.project_of(session).is_some());
        match &self.active {
            Some(session) => self.active_project = self.project_of(session).unwrap_or(0),
            None => {
                self.active_project = 0;
                if let Some((ix, next)) = self.threads(cx).0.into_iter().next() {
                    self.select(ix, next, window, cx);
                }
            }
        }
    }

    pub fn remote_overview(&self) -> RemoteOverview {
        RemoteOverview {
            serving: self.remote.server.as_ref().map(|server| server.port),
            error: self.remote.server_error.clone(),
            peers: self.remote.peers.clone(),
            environments: self.remote.environments.iter().map(|env| (env.host.id.clone(), env.status.clone())).collect(),
        }
    }

    pub fn disconnect_peer(&mut self, conn: u64, cx: &mut Context<Self>) {
        if let Some(server) = &self.remote.server {
            server.disconnect(conn);
        }
        self.remote.peers.retain(|peer| peer.conn != conn);
        self.remote.subscriptions.remove(&conn);
        cx.notify();
    }

    /// Notes a change to a local thread, for the clients watching.
    pub(super) fn remote_thread_changed(&mut self, session: &Entity<Session>, cx: &App) {
        if self.remote.server.is_none() {
            return;
        }
        let session = session.read(cx);
        if !session.is_remote() {
            self.remote.dirty_threads.insert(session.id.clone());
            self.remote.state_dirty = true;
        }
    }

    /// The local thread with a given id.
    fn local_thread(&self, id: &str, cx: &App) -> Option<(usize, Entity<Session>)> {
        self.projects.iter().enumerate().filter(|(_, project)| project.environment.is_none()).find_map(|(ix, project)| {
            project.sessions.iter().find(|session| session.read(cx).id == id).map(|session| (ix, session.clone()))
        })
    }

    // ---- serving ------------------------------------------------------------

    fn host_state(&self, cx: &App) -> ServerMessage {
        let local = self.projects.iter().filter(|project| project.environment.is_none());
        let projects = local
            .clone()
            .map(|project| {
                let git = project.git.clone();
                ProjectInfo {
                    path: project.path.clone(),
                    name: project.name.to_string(),
                    git: git.is_some(),
                    branch: git.as_ref().and_then(|git| git.branch.clone()),
                    branches: git.as_ref().map(|git| git.branches.clone()).unwrap_or_default(),
                    added: git.as_ref().map_or(0, |git| git.added),
                    removed: git.as_ref().map_or(0, |git| git.removed),
                }
            })
            .collect();
        let threads = local
            .flat_map(|project| &project.sessions)
            .map(|session| session.read(cx))
            .filter(|session| !session.is_blank())
            .map(Session::summary)
            .collect();
        ServerMessage::State { projects, threads }
    }

    fn on_server_event(&mut self, event: ServerEvent, window: &mut Window, cx: &mut Context<Self>) {
        let Some(server) = &self.remote.server else { return };
        match event {
            ServerEvent::Connected { conn, name, address } => {
                server.send(conn, &self.host_state(cx));
                self.remote.peers.push(Peer { conn, name: name.clone(), address });
                self.remote.subscriptions.insert(conn, HashMap::new());
                if settings::get(|settings| settings.in_app_notifications) {
                    window.push_notification(Notification::info(format!("{name} connected.")), cx);
                }
                cx.notify();
            }
            ServerEvent::Disconnected { conn } => {
                self.remote.peers.retain(|peer| peer.conn != conn);
                self.remote.subscriptions.remove(&conn);
                cx.notify();
            }
            ServerEvent::Message { conn, message } => self.on_client_message(conn, message, window, cx),
        }
    }

    fn on_client_message(&mut self, conn: u64, message: ClientMessage, window: &mut Window, cx: &mut Context<Self>) {
        match message {
            ClientMessage::Hello { .. } => {}
            ClientMessage::Subscribe { thread } => {
                self.remote.subscriptions.entry(conn).or_default().insert(thread.clone(), None);
                self.remote.dirty_threads.insert(thread);
            }
            ClientMessage::Unsubscribe { thread } => {
                if let Some(threads) = self.remote.subscriptions.get_mut(&conn) {
                    threads.remove(&thread);
                }
            }
            ClientMessage::Action { thread, action } => {
                let Some((_, session)) = self.local_thread(&thread, cx) else {
                    self.reply(conn, ServerMessage::Error { message: "That thread is gone.".into() });
                    return;
                };
                self.run_action(session, action, window, cx);
            }
            ClientMessage::Create {
                reference,
                project,
                instance,
                provider,
                model,
                permission_mode,
                reasoning,
                fast_mode,
                worktree,
                base,
                text,
                attachments,
            } => {
                let Some(project_ix) =
                    self.projects.iter().position(|candidate| candidate.environment.is_none() && candidate.path == project)
                else {
                    self.reply(conn, ServerMessage::Error { message: format!("{} isn't a project here.", project.display()) });
                    return;
                };
                let thread_settings = ThreadSettings {
                    instance,
                    provider,
                    model,
                    mode: PermissionMode::from_cli(&permission_mode).unwrap_or_default(),
                    reasoning: Reasoning::from_key(&reasoning).unwrap_or_default(),
                    fast_mode,
                };
                let mode = if worktree { WorkspaceMode::Worktree } else { WorkspaceMode::Local };
                let plan = self.worktree_plan(project_ix, mode).map(|mut plan| {
                    if base.is_some() {
                        plan.base = base;
                    }
                    plan
                });
                let session = self.spawn_thread(project_ix, thread_settings, plan, window, cx);
                let paths = store_uploads(attachments);
                session.update(cx, |session, cx| session.send_message(text, paths, cx));
                let id = session.read(cx).id.clone();
                self.remote.subscriptions.entry(conn).or_default().insert(id.clone(), None);
                self.remote.dirty_threads.insert(id.clone());
                self.remote.state_dirty = true;
                self.reply(conn, ServerMessage::Created { reference, thread: id });
            }
        }
    }

    fn reply(&self, conn: u64, message: ServerMessage) {
        if let Some(server) = &self.remote.server {
            server.send(conn, &message);
        }
    }

    /// Carries out what a client asked of a local thread.
    fn run_action(&mut self, session: Entity<Session>, action: ThreadAction, window: &mut Window, cx: &mut Context<Self>) {
        match action {
            ThreadAction::Send { text, attachments } => {
                let paths = store_uploads(attachments);
                session.update(cx, |session, cx| session.send_message(text, paths, cx));
            }
            ThreadAction::Queue { text, attachments } => {
                let paths = store_uploads(attachments);
                session.update(cx, |session, cx| session.queue_message(text, paths, cx));
            }
            ThreadAction::SteerQueued { ix } => session.update(cx, |session, cx| session.steer_queued(ix, cx)),
            ThreadAction::RemoveQueued { ix } => session.update(cx, |session, cx| session.remove_queued(ix, cx)),
            ThreadAction::Interrupt => session.update(cx, |session, cx| session.interrupt(cx)),
            ThreadAction::Permission { ix, decision } => {
                let decision = match decision.as_str() {
                    "always" => PermissionDecision::AllowAlways,
                    "deny" => PermissionDecision::Deny,
                    _ => PermissionDecision::Allow,
                };
                session.update(cx, |session, cx| session.respond_permission(ix, decision, cx));
            }
            ThreadAction::Answer { ix, option } => session.update(cx, |session, cx| session.answer_question(ix, option, cx)),
            ThreadAction::SetModel { model } => session.update(cx, |session, cx| session.set_model(model, cx)),
            ThreadAction::SetPermissionMode { mode } => {
                if let Some(mode) = PermissionMode::from_cli(&mode) {
                    session.update(cx, |session, cx| session.set_permission_mode(mode, cx));
                }
            }
            ThreadAction::SetReasoning { reasoning } => {
                if let Some(reasoning) = Reasoning::from_key(&reasoning) {
                    session.update(cx, |session, cx| session.set_reasoning(reasoning, cx));
                }
            }
            ThreadAction::SetFastMode { on } => session.update(cx, |session, cx| session.set_fast_mode(on, cx)),
            ThreadAction::Settle => self.settle_thread(session, cx),
            ThreadAction::Unsettle => session.update(cx, |session, cx| session.unsettle(cx)),
            ThreadAction::Pin { pinned } => session.update(cx, |session, cx| session.set_pinned(pinned, cx)),
            ThreadAction::Snooze { until } => session.update(cx, |session, cx| session.snooze(until.map(from_millis), cx)),
            ThreadAction::Rename { title } => session.update(cx, |session, cx| session.rename(title, cx)),
            ThreadAction::CancelResume => session.update(cx, |session, cx| session.cancel_resume(cx)),
            ThreadAction::ScheduleResume { at } => {
                session.update(cx, |session, cx| session.schedule_resume(from_millis(at), cx))
            }
            ThreadAction::Delete => self.forget_thread(session, window, cx),
        }
    }

    /// Sends clients the thread list when it changed, and each subscribed
    /// thread that changed from the first item that did.
    pub(super) fn flush_remote(&mut self, cx: &mut Context<Self>) {
        let Some(server) = &self.remote.server else { return };
        if std::mem::take(&mut self.remote.state_dirty) {
            let state = self.host_state(cx);
            let text = serde_json::to_string(&state).unwrap_or_default();
            if text != self.remote.last_state {
                self.remote.last_state = text;
                server.broadcast(&state);
            }
        }
        let dirty = std::mem::take(&mut self.remote.dirty_threads);
        if dirty.is_empty() {
            return;
        }
        let mut snapshots = HashMap::new();
        for (conn, threads) in &mut self.remote.subscriptions {
            for (id, sent) in threads.iter_mut() {
                if !dirty.contains(id) && sent.is_some() {
                    continue;
                }
                let Some(session) = self
                    .projects
                    .iter()
                    .filter(|project| project.environment.is_none())
                    .flat_map(|project| &project.sessions)
                    .find(|session| session.read(cx).id == *id)
                else {
                    continue;
                };
                let (summary, detail, items, hashes) = snapshots
                    .entry(id.clone())
                    .or_insert_with(|| {
                        let session = session.read(cx);
                        let items = session.saved_items(0);
                        let hashes: Vec<u64> = items.iter().map(item_hash).collect();
                        (session.summary(), session.detail(), items, hashes)
                    })
                    .clone();
                let previous = sent.take().unwrap_or_default();
                let from = previous.iter().zip(&hashes).take_while(|(a, b)| a == b).count();
                server.send(
                    *conn,
                    &ServerMessage::Thread { summary, detail, from, len: items.len(), items: items[from..].to_vec() },
                );
                *sent = Some(hashes);
            }
        }
    }

    // ---- connecting to other devices -------------------------------------------

    fn on_client_event(&mut self, env: &str, event: ClientEvent, window: &mut Window, cx: &mut Context<Self>) {
        match event {
            ClientEvent::Status(status) => {
                let connected = matches!(status, ConnectionStatus::Connected(_));
                if let Some(environment) = self.remote.environments.iter_mut().find(|candidate| candidate.host.id == env) {
                    environment.status = status;
                }
                if connected {
                    // A new connection knows nothing of what this one watched.
                    for session in self.environment_sessions(env) {
                        session.update(cx, |session, _| {
                            if let Some(remote) = session.remote.as_mut()
                                && remote.subscribed
                            {
                                remote.subscribed = false;
                                session.subscribe_remote();
                            }
                        });
                    }
                }
                cx.notify();
            }
            ClientEvent::Message(ServerMessage::State { projects, threads }) => {
                self.reconcile_environment(env, projects, threads, window, cx)
            }
            ClientEvent::Message(ServerMessage::Thread { summary, detail, from, len, items }) => {
                if let Some(session) = self.remote_session(env, &summary.id, cx) {
                    session.update(cx, |session, cx| session.apply_remote(&summary, detail, from, len, items, cx));
                }
            }
            ClientEvent::Message(ServerMessage::Created { reference, thread }) => {
                let session = self.environment_sessions(env).into_iter().find(|session| {
                    session.read(cx).remote.as_ref().is_some_and(|remote| remote.creating == Some(reference))
                });
                if let Some(session) = session {
                    session.update(cx, |session, cx| session.remote_created(thread, cx));
                }
            }
            ClientEvent::Message(ServerMessage::Error { message }) => {
                let name = self.environment_name(env);
                window.push_notification(Notification::error(message).title(name), cx);
            }
            ClientEvent::Message(ServerMessage::Welcome { .. } | ServerMessage::Denied { .. }) => {}
        }
    }

    fn environment_sessions(&self, env: &str) -> Vec<Entity<Session>> {
        self.projects
            .iter()
            .filter(|project| project.environment.as_deref() == Some(env))
            .flat_map(|project| project.sessions.clone())
            .collect()
    }

    fn remote_session(&self, env: &str, id: &str, cx: &App) -> Option<Entity<Session>> {
        self.environment_sessions(env).into_iter().find(|session| {
            session.read(cx).remote.as_ref().and_then(|remote| remote.thread.as_deref()) == Some(id)
        })
    }

    pub(super) fn environment_name(&self, env: &str) -> String {
        self.remote
            .environments
            .iter()
            .find(|candidate| candidate.host.id == env)
            .map(|env| env.host.label())
            .unwrap_or_else(|| env.to_string())
    }

    /// The link a new mirrored thread on `env` uses.
    pub(super) fn remote_link(&self, env: &str, thread: Option<String>) -> Option<RemoteLink> {
        let environment = self.remote.environments.iter().find(|candidate| candidate.host.id == env)?;
        Some(RemoteLink {
            env: env.to_string(),
            thread,
            client: environment.client.clone(),
            subscribed: false,
            loaded: false,
            attention: false,
            creating: None,
        })
    }

    /// The index of a device's project, adding it if needed.
    fn ensure_remote_project(&mut self, env: &str, path: PathBuf, name: Option<&str>) -> usize {
        if let Some(ix) =
            self.projects.iter().position(|project| project.environment.as_deref() == Some(env) && project.path == path)
        {
            return ix;
        }
        let name = name.map(str::to_string).unwrap_or_else(|| {
            path.file_name().map(|name| name.to_string_lossy().to_string()).unwrap_or_else(|| path.display().to_string())
        });
        self.projects.push(Project {
            path,
            initials: initials(&name).into(),
            tint: tint_for(&name),
            name: name.into(),
            added: SystemTime::now(),
            sessions: Vec::new(),
            git: None,
            environment: Some(env.to_string()),
        });
        self.projects.len() - 1
    }

    /// Brings a device's projects and threads in line with its latest list.
    fn reconcile_environment(
        &mut self,
        env: &str,
        projects: Vec<ProjectInfo>,
        threads: Vec<ThreadSummary>,
        window: &mut Window,
        cx: &mut Context<Self>,
    ) {
        for info in &projects {
            let ix = self.ensure_remote_project(env, info.path.clone(), Some(&info.name));
            self.projects[ix].git = info.git.then(|| GitInfo {
                branch: info.branch.clone(),
                branches: info.branches.clone(),
                added: info.added,
                removed: info.removed,
            });
        }
        let known: HashSet<&str> = threads.iter().map(|summary| summary.id.as_str()).collect();
        for session in self.environment_sessions(env) {
            let gone = session
                .read(cx)
                .remote
                .as_ref()
                .and_then(|remote| remote.thread.as_deref())
                .is_some_and(|id| !known.contains(id));
            if gone && let Some(ix) = self.project_of(&session) {
                let was_active = self.active.as_ref() == Some(&session);
                self.remove_thread(ix, &session, cx);
                if was_active {
                    self.active = None;
                    if let Some((ix, next)) = self.threads(cx).0.into_iter().next() {
                        self.select(ix, next, window, cx);
                    }
                }
            }
        }
        for summary in threads {
            match self.remote_session(env, &summary.id, cx) {
                Some(session) => session.update(cx, |session, cx| {
                    session.apply_summary(&summary);
                    cx.notify();
                }),
                None => {
                    let Some(link) = self.remote_link(env, Some(summary.id.clone())) else { continue };
                    let project_ix = self.ensure_remote_project(env, summary.project.clone(), None);
                    let session = cx.new(|cx| Session::remote(&summary, link, cx));
                    self.attach(project_ix, session, window, cx);
                }
            }
        }
        cx.notify();
    }

    /// A new, empty thread for one of a device's projects.
    pub(super) fn new_remote_session(
        &mut self,
        env: &str,
        project: PathBuf,
        settings: ThreadSettings,
        worktree: bool,
        cx: &mut Context<Self>,
    ) -> Option<Entity<Session>> {
        let link = self.remote_link(env, None)?;
        let plan = worktree.then(|| crate::worktree::WorktreePlan {
            project: project.clone(),
            base: None,
            from_origin: false,
            submodules: settings::Submodules::Skip,
        });
        let new = NewSession {
            instance: settings.instance,
            provider: if settings.provider == Provider::Pi { Provider::Pi } else { Provider::Claude },
            project,
            model: settings.model,
            permission_mode: settings.mode,
            reasoning: settings.reasoning,
            fast_mode: settings.fast_mode,
            worktree: plan,
        };
        Some(cx.new(|_| Session::new_remote(new, link)))
    }
}

fn item_hash<T: serde::Serialize>(item: &T) -> u64 {
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    serde_json::to_string(item).unwrap_or_default().hash(&mut hasher);
    hasher.finish()
}