// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

//! The russh server: publickey auth, exec dispatch into pack processes.

use std::sync::Arc;
use std::time::Duration;

use config::Config;
use model::KeyKind;
use russh::keys::{HashAlg, PublicKey};
use russh::server::{Auth, ChannelOpenHandle, Handler, Msg, Server as _, Session};
use russh::{Channel, ChannelId, MethodKind, MethodSet};
use store::Store;
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
use tokio::process::ChildStdin;
use tokio::sync::Mutex;

use crate::SshError;
use crate::{command, resolve};

/// State shared across every connection.
struct Shared {
    store: Store,
    config: Arc<Config>,
    /// The `--config` path to hand hooks via `FABRICA_CONFIG`.
    config_path: Option<String>,
    /// The instance signing secret, for minting short-lived LFS bearer tokens.
    secret: String,
}

/// The russh `Server` factory.
struct GitServer {
    shared: Arc<Shared>,
}

impl russh::server::Server for GitServer {
    type Handler = GitHandler;

    fn new_client(&mut self, _peer: Option<std::net::SocketAddr>) -> GitHandler {
        GitHandler {
            shared: Arc::clone(&self.shared),
            identity: None,
            stdin: Arc::new(Mutex::new(None)),
        }
    }
}

/// The authenticated identity, bound after publickey auth.
#[derive(Clone)]
struct Identity {
    user_id: String,
    username: String,
    is_admin: bool,
    key_id: String,
}

/// One connection's handler.
struct GitHandler {
    shared: Arc<Shared>,
    identity: Option<Identity>,
    /// The running pack subprocess's stdin, written to as client data arrives.
    stdin: Arc<Mutex<Option<ChildStdin>>>,
}

impl Handler for GitHandler {
    type Error = russh::Error;

    async fn auth_publickey(
        &mut self,
        _user: &str,
        public_key: &PublicKey,
    ) -> Result<Auth, Self::Error> {
        // Identify by key fingerprint, not the SSH username (everyone is `git@`).
        let fingerprint = public_key.fingerprint(HashAlg::Sha256).to_string();
        let Ok(Some(key)) = self
            .shared
            .store
            .key_by_fingerprint(KeyKind::Ssh, &fingerprint)
            .await
        else {
            return Ok(Auth::reject());
        };
        let Ok(Some(user)) = self.shared.store.user_by_id(&key.user_id).await else {
            return Ok(Auth::reject());
        };
        if user.disabled_at.is_some() {
            return Ok(Auth::reject());
        }
        let _ = self.shared.store.touch_key_used(&key.id).await;
        // A successful public-key auth proves possession of the private key, so
        // the key is now verified.
        let _ = self.shared.store.set_key_verified(&key.id).await;
        self.identity = Some(Identity {
            user_id: user.id,
            username: user.username,
            is_admin: user.is_admin,
            key_id: key.id,
        });
        Ok(Auth::Accept)
    }

    async fn channel_open_session(
        &mut self,
        _channel: Channel<Msg>,
        reply: ChannelOpenHandle,
        _session: &mut Session,
    ) -> Result<(), Self::Error> {
        // Dropping the handle rejects the channel, so accept it explicitly.
        reply.accept().await;
        Ok(())
    }

    async fn exec_request(
        &mut self,
        channel: ChannelId,
        data: &[u8],
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        session.channel_success(channel)?;
        let command = String::from_utf8_lossy(data).into_owned();
        if let Err(banner) = self.spawn_exec(channel, &command, session).await {
            let _ = session.extended_data(channel, 1, banner.into_bytes());
            let _ = session.exit_status_request(channel, 1);
            let _ = session.eof(channel);
            let _ = session.close(channel);
        }
        Ok(())
    }

    async fn shell_request(
        &mut self,
        channel: ChannelId,
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        session.channel_success(channel)?;
        // Use the async handle so the banner/close actually flush before return.
        let handle = session.handle();
        let _ = handle
            .extended_data(channel, 1, self.no_shell_banner().into_bytes())
            .await;
        let _ = handle.exit_status_request(channel, 1).await;
        let _ = handle.eof(channel).await;
        let _ = handle.close(channel).await;
        Ok(())
    }

    async fn data(
        &mut self,
        _channel: ChannelId,
        data: &[u8],
        _session: &mut Session,
    ) -> Result<(), Self::Error> {
        if let Some(stdin) = self.stdin.lock().await.as_mut() {
            let _ = stdin.write_all(data).await;
        }
        Ok(())
    }

    async fn channel_eof(
        &mut self,
        _channel: ChannelId,
        _session: &mut Session,
    ) -> Result<(), Self::Error> {
        // Close the child's stdin so upload-pack/receive-pack sees EOF.
        *self.stdin.lock().await = None;
        Ok(())
    }
}

impl GitHandler {
    /// The friendly shell-access banner.
    fn no_shell_banner(&self) -> String {
        let who = self
            .identity
            .as_ref()
            .map_or("there", |i| i.username.as_str());
        format!("Hi {who}! fabrica does not provide shell access.\n")
    }

    /// Authorize the exec command and, on success, spawn the pack process and start
    /// streaming. On failure returns the banner to write to stderr.
    async fn spawn_exec(
        &mut self,
        channel: ChannelId,
        command: &str,
        session: &mut Session,
    ) -> Result<(), String> {
        let Some(identity) = self.identity.clone() else {
            return Err("Authentication required.".to_string());
        };
        // git-lfs-authenticate: hand back an HTTP endpoint + bearer token, then exit.
        if let Some(lfs) = command::parse_lfs(command) {
            return self
                .lfs_authenticate(channel, &identity, &lfs, session)
                .await;
        }
        let parsed = command::parse(command).ok_or_else(|| self.no_shell_banner())?;
        let (owner, repo_path) = resolve::normalize(&parsed.path)
            .ok_or_else(|| "Invalid repository path.".to_string())?;

        let repo = self
            .resolve_repo(&owner, &repo_path, &identity, parsed.service)
            .await?;

        let repo_dir = git::repo_path(&self.shared.config.storage.repo_dir, &repo.id)
            .map_err(|_| "Repository is unavailable.".to_string())?;

        let env = git::pack::PackEnv {
            home: self.shared.config.storage.data_dir.join("tmp"),
            config_path: self.shared.config_path.clone(),
            user_id: Some(identity.user_id),
            user_name: Some(identity.username),
            repo_id: Some(repo.id),
            key_id: Some(identity.key_id),
            protocol: None,
        };
        let mut child = git::pack::spawn(
            &self.shared.config.git.binary,
            parsed.service,
            &[],
            &repo_dir,
            &env,
        )
        .map_err(|e| format!("Failed to start git: {e}"))?;

        let (Some(stdin), Some(mut stdout), Some(mut stderr)) =
            (child.stdin.take(), child.stdout.take(), child.stderr.take())
        else {
            return Err("Internal error.".to_string());
        };
        *self.stdin.lock().await = Some(stdin);

        let handle = session.handle();
        tokio::spawn(async move {
            let mut obuf = vec![0u8; 32 * 1024];
            let mut ebuf = vec![0u8; 32 * 1024];
            let (mut out_done, mut err_done) = (false, false);
            while !(out_done && err_done) {
                tokio::select! {
                    r = stdout.read(&mut obuf), if !out_done => match r {
                        Ok(0) | Err(_) => out_done = true,
                        Ok(n) => { let _ = handle.data(channel, obuf[..n].to_vec()).await; }
                    },
                    r = stderr.read(&mut ebuf), if !err_done => match r {
                        Ok(0) | Err(_) => err_done = true,
                        Ok(n) => { let _ = handle.extended_data(channel, 1, ebuf[..n].to_vec()).await; }
                    },
                }
            }
            let code = child
                .wait()
                .await
                .ok()
                .and_then(|s| s.code())
                .and_then(|c| u32::try_from(c).ok())
                .unwrap_or(0);
            let _ = handle.exit_status_request(channel, code).await;
            let _ = handle.eof(channel).await;
            let _ = handle.close(channel).await;
        });
        Ok(())
    }

    /// Handle `git-lfs-authenticate`: authorize the repo for the operation, mint a
    /// short-lived bearer token, and write the LFS endpoint JSON to stdout. The
    /// actual object transfer then happens over HTTP against that endpoint.
    async fn lfs_authenticate(
        &mut self,
        channel: ChannelId,
        identity: &Identity,
        lfs: &command::LfsAuth,
        session: &mut Session,
    ) -> Result<(), String> {
        let (owner, repo_path) =
            resolve::normalize(&lfs.path).ok_or_else(|| "Invalid repository path.".to_string())?;
        let not_found = || "Repository not found.".to_string();
        let owner_user = self
            .shared
            .store
            .user_by_username(&owner)
            .await
            .ok()
            .flatten()
            .ok_or_else(not_found)?;
        let repo = self
            .shared
            .store
            .repo_by_owner_path(&owner_user.id, &repo_path)
            .await
            .ok()
            .flatten()
            .ok_or_else(not_found)?;
        let collaborator = self
            .shared
            .store
            .effective_permission(&repo.id, &identity.user_id)
            .await
            .ok()
            .flatten()
            .and_then(|p| auth::Permission::parse(&p));
        let viewer = auth::Viewer {
            id: identity.user_id.clone(),
            is_admin: identity.is_admin,
        };
        let access = auth::access(
            Some(&viewer),
            &repo,
            collaborator,
            self.shared.config.instance.allow_anonymous,
        );
        let required = if lfs.operation == "upload" {
            auth::Access::Write
        } else {
            auth::Access::Read
        };
        if access < required {
            return Err(if access >= auth::Access::Read {
                "You do not have write access to this repository.".to_string()
            } else {
                not_found()
            });
        }

        // A one-hour bearer token the HTTP LFS endpoints accept (verified by
        // signature + expiry, not the API-token revocation path).
        let now_s = i64::try_from(
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map_or(0, |d| d.as_secs()),
        )
        .unwrap_or(0);
        let claims = auth::Claims {
            sub: identity.user_id.clone(),
            jti: format!("lfs-{}", repo.id),
            scopes: "lfs".to_string(),
            iat: now_s,
            exp: now_s + 3600,
        };
        let token = auth::mint_jwt(&self.shared.secret, &claims)
            .map_err(|_| "Could not issue an LFS token.".to_string())?;
        let href = format!(
            "{}/{}/{}.git/info/lfs",
            self.shared.config.instance.url.trim_end_matches('/'),
            owner,
            repo_path
        );
        // git-lfs reads this JSON from stdout. href/token contain only URL- and
        // JWT-safe characters, so no escaping is needed.
        let json = format!(
            "{{\"href\":\"{href}\",\"header\":{{\"Authorization\":\"Bearer {token}\"}},\"expires_in\":3600}}\n"
        );
        let handle = session.handle();
        let _ = handle.data(channel, json.into_bytes()).await;
        let _ = handle.exit_status_request(channel, 0).await;
        let _ = handle.eof(channel).await;
        let _ = handle.close(channel).await;
        Ok(())
    }

    /// Resolve and authorize the repo for the requested service.
    async fn resolve_repo(
        &self,
        owner: &str,
        repo_path: &str,
        identity: &Identity,
        service: git::pack::Service,
    ) -> Result<model::Repo, String> {
        let not_found = || "Repository not found.".to_string();
        let owner_user = self
            .shared
            .store
            .user_by_username(owner)
            .await
            .ok()
            .flatten()
            .ok_or_else(not_found)?;
        let existing = self
            .shared
            .store
            .repo_by_owner_path(&owner_user.id, repo_path)
            .await
            .ok()
            .flatten();
        // Push-to-create: a push (receive-pack) to a non-existent repo owned by
        // the pushing user (or by an admin) auto-creates it, at the instance
        // default visibility. Any other case is a 404 so a stranger cannot probe
        // for repositories.
        let repo = if let Some(repo) = existing {
            repo
        } else if matches!(service, git::pack::Service::ReceivePack)
            && (identity.user_id == owner_user.id || identity.is_admin)
        {
            self.create_on_push(&owner_user, repo_path).await?
        } else {
            return Err(not_found());
        };

        // Archived repositories are read-only: refuse any push.
        if matches!(service, git::pack::Service::ReceivePack) && repo.archived_at.is_some() {
            return Err("This repository is archived and is read-only.".to_string());
        }

        let collaborator = self
            .shared
            .store
            .effective_permission(&repo.id, &identity.user_id)
            .await
            .ok()
            .flatten()
            .and_then(|p| auth::Permission::parse(&p));
        let viewer = auth::Viewer {
            id: identity.user_id.clone(),
            is_admin: identity.is_admin,
        };
        let access = auth::access(
            Some(&viewer),
            &repo,
            collaborator,
            self.shared.config.instance.allow_anonymous,
        );

        let required = match service {
            git::pack::Service::ReceivePack => auth::Access::Write,
            _ => auth::Access::Read,
        };
        if access >= required {
            Ok(repo)
        } else if access >= auth::Access::Read {
            Err("You do not have write access to this repository.".to_string())
        } else {
            // No read access at all: do not confirm the repo exists.
            Err(not_found())
        }
    }

    /// Create a repository on first push: any missing group prefix, the database
    /// row (visibility = instance default), and the bare repo on disk. Rolls the
    /// row back if the on-disk creation fails.
    async fn create_on_push(
        &self,
        owner: &model::User,
        repo_path: &str,
    ) -> Result<model::Repo, String> {
        let (group_path, leaf) = match repo_path.rsplit_once('/') {
            Some((group, leaf)) => (Some(group), leaf),
            None => (None, repo_path),
        };
        let group_id = match group_path {
            Some(gp) => Some(
                self.shared
                    .store
                    .ensure_group_path(&owner.id, gp)
                    .await
                    .map_err(|e| e.to_string())?
                    .id,
            ),
            None => None,
        };
        let branch = self.shared.config.instance.default_branch.clone();
        let visibility = model::Visibility::from_token(
            self.shared.config.instance.default_visibility.as_token(),
        )
        .unwrap_or(model::Visibility::Private);
        let repo = self
            .shared
            .store
            .create_repo(store::NewRepo {
                owner_id: owner.id.clone(),
                group_id,
                name: leaf.to_string(),
                path: repo_path.to_string(),
                description: None,
                visibility,
                default_branch: branch.clone(),
            })
            .await
            .map_err(|e| e.to_string())?;

        let hook = std::env::current_exe().unwrap_or_else(|_| std::path::PathBuf::from("fabrica"));
        if let Err(err) = git::create_bare(
            &self.shared.config.storage.repo_dir,
            &repo.id,
            &branch,
            &hook,
        ) {
            let _ = self.shared.store.delete_repo(&repo.id).await;
            return Err(format!("Could not create repository: {err}"));
        }
        tracing::info!(owner = %owner.username, path = %repo_path, "push-to-create: repository created");
        Ok(repo)
    }
}

/// Run the SSH server until the process exits. A no-op when `ssh.enabled` is false.
///
/// # Errors
///
/// Returns [`SshError`] if the host key cannot be prepared or the server fails to
/// bind or run.
pub async fn serve(
    config: Arc<Config>,
    store: Store,
    config_path: Option<String>,
    secret: String,
) -> Result<(), SshError> {
    if !config.ssh.enabled {
        tracing::info!("ssh disabled (ssh.enabled = false)");
        return Ok(());
    }

    let host_pem = crate::hostkey::load_or_generate(&config.ssh.host_key_path)?;
    let host_key = russh::keys::PrivateKey::from_openssh(host_pem.as_str())
        .map_err(|e| SshError::HostKey(format!("parsing host key: {e}")))?;
    let russh_config = Arc::new(russh::server::Config {
        methods: MethodSet::from(&[MethodKind::PublicKey][..]),
        keys: vec![host_key],
        inactivity_timeout: Some(Duration::from_secs(3600)),
        auth_rejection_time: Duration::from_secs(2),
        ..russh::server::Config::default()
    });

    let shared = Arc::new(Shared {
        store,
        config: Arc::clone(&config),
        config_path,
        secret,
    });
    let addr = (config.ssh.address.clone(), config.ssh.port);
    tracing::info!(address = %config.ssh.address, port = config.ssh.port, "ssh server listening");

    let mut server = GitServer { shared };
    server
        .run_on_address(russh_config, addr)
        .await
        .map_err(|e| SshError::Server(e.to_string()))
}