//! Interactive provider login carried between a CLI on a configured machine
//! and the phone. The CLI remains the only credential writer: this layer keeps
//! its short-lived process and relays only the authorization URL and the code
//! a person copies back from the browser.

use std::collections::HashMap;
use std::io::{BufRead, BufReader, Read, Write};
use std::process::Stdio;
use std::sync::{Arc, Mutex, mpsc};
use std::time::{Duration, Instant};

use anyhow::{Context, Result};
use rand::Rng;
use serde::Serialize;

use crate::config::{MachineConfig, ProviderConfig};
use crate::session::transport::Transport;
use crate::usage::UsageMonitor;

const LOGIN_TIMEOUT: Duration = Duration::from_secs(10 * 60);
const AUTHORIZATION_URL_TIMEOUT: Duration = Duration::from_secs(15);
const OUTPUT_POLL: Duration = Duration::from_millis(50);
const START_WAIT: Duration = Duration::from_secs(16);

type Key = (String, String);

#[derive(Debug, Clone, Serialize)]
#[serde(tag = "state", rename_all = "camelCase")]
pub enum LoginState {
    Starting,
    WaitingForCode {
        #[serde(rename = "authorizationUrl")]
        authorization_url: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        detail: Option<String>,
    },
    Submitting,
    Succeeded,
    Failed {
        detail: String,
    },
    Cancelled,
}

impl LoginState {
    fn terminal(&self) -> bool {
        matches!(
            self,
            Self::Succeeded | Self::Failed { .. } | Self::Cancelled
        )
    }
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LoginInfo {
    pub attempt: String,
    #[serde(flatten)]
    pub state: LoginState,
}

#[derive(Clone)]
struct Attempt {
    id: String,
    state: Arc<Mutex<LoginState>>,
    input: mpsc::Sender<Input>,
}

impl Attempt {
    fn info(&self) -> LoginInfo {
        LoginInfo {
            attempt: self.id.clone(),
            state: self.state.lock().unwrap().clone(),
        }
    }
}

enum Input {
    Code(String),
    Cancel,
}

/// The active login per machine and provider. Completed attempts stay until a
/// new one replaces them, so a phone that briefly loses its connection can ask
/// how the operation ended rather than being handed an ambiguous 404.
pub struct LoginManager {
    attempts: Mutex<HashMap<Key, Attempt>>,
    usage: Arc<UsageMonitor>,
}

impl LoginManager {
    pub fn new(usage: Arc<UsageMonitor>) -> Self {
        Self {
            attempts: Mutex::new(HashMap::new()),
            usage,
        }
    }

    pub fn start(&self, machine: MachineConfig, provider: ProviderConfig) -> LoginInfo {
        let provider_key = provider
            .kind
            .usage_provider()
            .expect("a login route only accepts a metered provider")
            .to_string();
        let key = (machine.id.clone(), provider_key);
        let mut attempts = self.attempts.lock().unwrap();
        if let Some(attempt) = attempts.get(&key)
            && !attempt.state.lock().unwrap().terminal()
        {
            return attempt.info();
        }

        let id = attempt_id();
        let state = Arc::new(Mutex::new(LoginState::Starting));
        let (input, commands) = mpsc::channel();
        let attempt = Attempt {
            id: id.clone(),
            state: Arc::clone(&state),
            input,
        };
        attempts.insert(key, attempt.clone());
        drop(attempts);

        // Seed the cache before the worker takes the gate. A usage request in
        // that small handoff window sees a fresh, truthful state and cannot
        // start a second CLI against the same credential.
        let gate = self.usage.claude_authentication_started(&machine);
        let usage = Arc::clone(&self.usage);
        std::thread::spawn(move || {
            let _guard = gate.lock().unwrap();
            run_login(&machine, &provider, commands, &state);
            usage.claude_authentication_finished(&machine.id);
        });

        attempt.info()
    }

    pub fn wait_until_ready(&self, machine: &str, provider: &str, attempt: &str) -> LoginInfo {
        let started = Instant::now();
        loop {
            let info = self.read(machine, provider, attempt).unwrap_or(LoginInfo {
                attempt: attempt.to_string(),
                state: LoginState::Failed {
                    detail: "the sign-in attempt disappeared".to_string(),
                },
            });
            if !matches!(info.state, LoginState::Starting) || started.elapsed() >= START_WAIT {
                return info;
            }
            std::thread::sleep(OUTPUT_POLL);
        }
    }

    pub fn read(&self, machine: &str, provider: &str, attempt: &str) -> Option<LoginInfo> {
        let attempts = self.attempts.lock().unwrap();
        let found = attempts.get(&(machine.to_string(), provider.to_string()))?;
        (found.id == attempt).then(|| found.info())
    }

    pub fn submit(
        &self,
        machine: &str,
        provider: &str,
        attempt: &str,
        code: &str,
    ) -> Result<LoginInfo> {
        let code = valid_code(code)?;
        let attempts = self.attempts.lock().unwrap();
        let found = attempts
            .get(&(machine.to_string(), provider.to_string()))
            .filter(|found| found.id == attempt)
            .context("no such sign-in attempt")?;
        if found.state.lock().unwrap().terminal() {
            return Ok(found.info());
        }
        *found.state.lock().unwrap() = LoginState::Submitting;
        if found.input.send(Input::Code(code.to_string())).is_err() {
            *found.state.lock().unwrap() = LoginState::Failed {
                detail: "the sign-in process has stopped".to_string(),
            };
            anyhow::bail!("the sign-in process has stopped");
        }
        Ok(found.info())
    }

    pub fn cancel(&self, machine: &str, provider: &str, attempt: &str) -> Result<LoginInfo> {
        let attempts = self.attempts.lock().unwrap();
        let found = attempts
            .get(&(machine.to_string(), provider.to_string()))
            .filter(|found| found.id == attempt)
            .context("no such sign-in attempt")?;
        if !found.state.lock().unwrap().terminal() {
            let _ = found.input.send(Input::Cancel);
        }
        Ok(found.info())
    }

    /// Interactive helpers are unlike sessions: nothing adopts them after a
    /// server restart. End every one while the process is still here to reap
    /// the child it launched.
    pub fn cancel_all(&self) {
        let attempts = self.attempts.lock().unwrap();
        for attempt in attempts.values() {
            if !attempt.state.lock().unwrap().terminal() {
                let _ = attempt.input.send(Input::Cancel);
            }
        }
        let pending: Vec<_> = attempts
            .values()
            .map(|attempt| Arc::clone(&attempt.state))
            .collect();
        drop(attempts);
        let deadline = Instant::now() + Duration::from_secs(2);
        while Instant::now() < deadline
            && pending
                .iter()
                .any(|state| !state.lock().unwrap().terminal())
        {
            std::thread::sleep(OUTPUT_POLL);
        }
    }
}

fn run_login(
    machine: &MachineConfig,
    provider: &ProviderConfig,
    commands: mpsc::Receiver<Input>,
    state: &Arc<Mutex<LoginState>>,
) {
    if let Err(err) = run_login_inner(machine, provider, commands, state) {
        *state.lock().unwrap() = LoginState::Failed {
            detail: format!("couldn't sign in to Claude on {}: {err:#}", machine.name),
        };
    }
}

fn run_login_inner(
    machine: &MachineConfig,
    provider: &ProviderConfig,
    commands: mpsc::Receiver<Input>,
    state: &Arc<Mutex<LoginState>>,
) -> Result<()> {
    let transport = Transport::for_machine(machine);
    let args = vec![
        "BROWSER=/bin/false".to_string(),
        provider.program().to_string(),
        "auth".to_string(),
        "login".to_string(),
        "--claudeai".to_string(),
    ];
    let host = match &transport {
        Transport::Here => None,
        Transport::Ssh { ssh, .. } => Some(ssh),
    };
    let mut command = crate::ssh::command(host, "env", &args, None, None);
    command
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());
    let mut child = command
        .spawn()
        .with_context(|| format!("couldn't run {} auth login", provider.program()))?;
    let mut stdin = child
        .stdin
        .take()
        .context("the login process has no stdin")?;
    let stdout = child
        .stdout
        .take()
        .context("the login process has no stdout")?;
    let stderr = child
        .stderr
        .take()
        .context("the login process has no stderr")?;
    let (output, lines) = mpsc::channel();
    read_lines(stdout, output.clone());
    read_lines(stderr, output);

    let started = Instant::now();
    let mut authorization_url = None;
    let mut last_line = None;
    loop {
        while let Ok(line) = lines.try_recv() {
            if let Some(url) = authorization_url_in(&line) {
                authorization_url = Some(url.to_string());
                *state.lock().unwrap() = LoginState::WaitingForCode {
                    authorization_url: url.to_string(),
                    detail: None,
                };
            } else if line.to_ascii_lowercase().contains("invalid code") {
                if let Some(url) = &authorization_url {
                    *state.lock().unwrap() = LoginState::WaitingForCode {
                        authorization_url: url.clone(),
                        detail: Some(
                            "That code was not accepted. Copy the complete code and try again."
                                .to_string(),
                        ),
                    };
                }
            } else if !line.trim().is_empty() {
                last_line = Some(line.trim().chars().take(500).collect::<String>());
            }
        }

        match commands.recv_timeout(OUTPUT_POLL) {
            Ok(Input::Code(code)) => {
                *state.lock().unwrap() = LoginState::Submitting;
                writeln!(stdin, "{code}").context("couldn't send the login code")?;
                stdin.flush().context("couldn't send the login code")?;
            }
            Ok(Input::Cancel) => {
                let _ = child.kill();
                let _ = child.wait();
                *state.lock().unwrap() = LoginState::Cancelled;
                return Ok(());
            }
            Err(mpsc::RecvTimeoutError::Disconnected) => {
                let _ = child.kill();
                let _ = child.wait();
                anyhow::bail!("the phone disconnected from the sign-in attempt");
            }
            Err(mpsc::RecvTimeoutError::Timeout) => {}
        }

        if let Some(status) = child
            .try_wait()
            .context("couldn't check the login process")?
        {
            *state.lock().unwrap() = if status.success() {
                LoginState::Succeeded
            } else {
                LoginState::Failed {
                    detail: last_line
                        .unwrap_or_else(|| format!("Claude's login process exited with {status}")),
                }
            };
            return Ok(());
        }
        if authorization_url.is_none() && started.elapsed() >= AUTHORIZATION_URL_TIMEOUT {
            let _ = child.kill();
            let _ = child.wait();
            anyhow::bail!("the Claude CLI did not provide an authorization URL");
        }
        if started.elapsed() >= LOGIN_TIMEOUT {
            let _ = child.kill();
            let _ = child.wait();
            anyhow::bail!("the sign-in attempt expired; start it again");
        }
    }
}

fn read_lines(reader: impl Read + Send + 'static, output: mpsc::Sender<String>) {
    std::thread::spawn(move || {
        for line in BufReader::new(reader).lines().map_while(Result::ok) {
            let _ = output.send(line);
        }
    });
}

fn authorization_url_in(line: &str) -> Option<&str> {
    let start = line.find("https://")?;
    let tail = &line[start..];
    let end = tail
        .find(|character: char| character.is_whitespace() || character == '\u{1b}')
        .unwrap_or(tail.len());
    let url = &tail[..end];
    (url.starts_with("https://claude.com/") || url.starts_with("https://platform.claude.com/"))
        .then_some(url)
}

fn valid_code(code: &str) -> Result<&str> {
    let code = code.trim();
    anyhow::ensure!(!code.is_empty(), "the login code is empty");
    anyhow::ensure!(code.len() <= 4096, "the login code is too long");
    anyhow::ensure!(
        !code.chars().any(char::is_control),
        "the login code contains a line break or control character"
    );
    Ok(code)
}

fn attempt_id() -> String {
    let mut bytes = [0u8; 16];
    rand::rng().fill_bytes(&mut bytes);
    bytes.iter().map(|byte| format!("{byte:02x}")).collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::DriverKind;

    #[test]
    fn extracts_only_anthropics_https_login_url() {
        assert_eq!(
            authorization_url_in("visit: https://claude.com/cai/oauth/authorize?state=x"),
            Some("https://claude.com/cai/oauth/authorize?state=x")
        );
        assert!(authorization_url_in("visit: http://claude.com/nope").is_none());
        assert!(authorization_url_in("visit: https://example.com/nope").is_none());
    }

    #[test]
    fn login_code_is_one_bounded_line() {
        assert_eq!(valid_code("  abc#state  ").unwrap(), "abc#state");
        assert!(valid_code("\n").is_err());
        assert!(valid_code("a\nb").is_err());
        assert!(valid_code(&"x".repeat(4097)).is_err());
    }

    #[cfg(unix)]
    #[test]
    fn relays_a_headless_cli_login_without_taking_over_its_credentials() {
        use std::os::unix::fs::PermissionsExt;

        let dir = tempfile::tempdir().expect("tempdir");
        let cli = dir.path().join("fake-claude");
        std::fs::write(
            &cli,
            "#!/bin/sh\necho 'https://claude.com/cai/oauth/authorize?state=test'\nIFS= read -r code\n[ \"$code\" = 'the-code' ]\n",
        )
        .expect("write fake CLI");
        std::fs::set_permissions(&cli, std::fs::Permissions::from_mode(0o700))
            .expect("make fake CLI executable");

        let monitor = Arc::new(UsageMonitor::new(Default::default()));
        let logins = LoginManager::new(monitor);
        let machine = MachineConfig {
            id: "vm".to_string(),
            name: "test vm".to_string(),
            ssh: None,
            providers: Vec::new(),
        };
        let provider = ProviderConfig {
            name: "claude-cli".to_string(),
            kind: DriverKind::ClaudeCli,
            command: Some(cli.display().to_string()),
            models: Vec::new(),
            mcp_servers: Vec::new(),
            model_settings: Default::default(),
            max_loaded: None,
        };

        let started = logins.start(machine, provider);
        let ready = logins.wait_until_ready("vm", "claude", &started.attempt);
        assert!(matches!(ready.state, LoginState::WaitingForCode { .. }));
        let wire = serde_json::to_value(&ready).expect("serialize login state");
        assert!(wire.get("authorizationUrl").is_some(), "{wire}");
        assert!(wire.get("authorization_url").is_none(), "{wire}");
        let submitted = logins
            .submit("vm", "claude", &started.attempt, "the-code")
            .expect("submit code");
        assert!(matches!(submitted.state, LoginState::Submitting));

        let deadline = Instant::now() + Duration::from_secs(2);
        loop {
            let finished = logins
                .read("vm", "claude", &started.attempt)
                .expect("attempt remains readable");
            if matches!(finished.state, LoginState::Succeeded) {
                break;
            }
            assert!(
                Instant::now() < deadline,
                "login did not finish: {finished:?}"
            );
            std::thread::sleep(OUTPUT_POLL);
        }
    }
}