//! Pi-compatible Claude subscription OAuth. Secrets are never logged or serialized into settings.
use serde::{Deserialize, Serialize};

#[derive(Clone, Serialize, Deserialize)]
pub(crate) struct Credentials {
    pub(crate) access_token: String,
    refresh_token: String,
    expires_at: u64,
}

#[cfg(not(target_family = "wasm"))]
pub(crate) use desktop::*;

#[cfg(not(target_family = "wasm"))]
mod desktop {
    use std::collections::HashMap;
    use std::io::{BufRead, BufReader, Write};
    use std::net::{TcpListener, TcpStream};
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::{Arc, Mutex, OnceLock};
    use std::time::{Duration, SystemTime, UNIX_EPOCH};

    use base64::engine::general_purpose::URL_SAFE_NO_PAD;
    use base64::Engine;
    use futures::channel::oneshot;
    use futures::StreamExt;
    use galaxyui::{Entity, ModelContext, SingletonEntity};
    use galaxyui_extras::secure_storage::AppContextExt as _;
    use instant::Instant;
    use rand::RngCore;
    use settings::Setting;
    use sha2::{Digest, Sha256};
    use url::Url;

    use super::Credentials;
    use crate::ai::claude_subscription::{is_stored_credential, OAUTH_REFERENCE_PREFIX};
    use crate::settings::AISettings;

    // Pi's browser flow, checked 2026-10-04:
    // https://github.com/earendil-works/pi/blob/main/packages/ai/src/auth/oauth/anthropic.ts
    const CLIENT_ID: &str = "9d1c250a-e61b-44d9-88ed-5944d1962f5e";
    const AUTHORIZE_URL: &str = "https://claude.ai/oauth/authorize";
    const TOKEN_URL: &str = "https://platform.claude.com/v1/oauth/token";
    const REDIRECT_URI: &str = "http://localhost:53692/callback";
    const SCOPES: &str = "org:create_api_key user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload";
    const SIGN_IN_TIMEOUT: Duration = Duration::from_secs(5 * 60);
    type Reply = oneshot::Sender<Result<String, String>>;
    enum Command {
        Resolve(String, Reply),
        Invalidate(String),
        DiscardDraft(String),
    }
    static SENDER: OnceLock<Mutex<Option<async_channel::Sender<Command>>>> = OnceLock::new();

    fn send(command: Command) -> Result<(), String> {
        SENDER
            .get_or_init(|| Mutex::new(None))
            .lock()
            .unwrap()
            .as_ref()
            .ok_or("Claude authorization service is unavailable.")?
            .try_send(command)
            .map_err(|_| "Claude authorization service is busy or unavailable.".to_string())
    }

    pub(crate) fn invalidate(reference: &str) {
        let _ = send(Command::Invalidate(reference.to_string()));
    }

    pub(crate) fn discard_draft(reference: &str) {
        if reference.starts_with(OAUTH_REFERENCE_PREFIX) {
            let _ = send(Command::DiscardDraft(reference.to_string()));
        }
    }

    pub(crate) async fn load_credential(reference: String) -> Result<String, String> {
        if !is_stored_credential(&reference) {
            return Ok(reference);
        }
        let (tx, rx) = oneshot::channel();
        send(Command::Resolve(reference, tx))?;
        let result = futures::future::select(
            Box::pin(rx),
            Box::pin(async_io::Timer::after(Duration::from_secs(35))),
        )
        .await;
        match result {
            futures::future::Either::Left((result, _)) => {
                result.map_err(|_| "Claude authorization service stopped.".to_string())?
            }
            futures::future::Either::Right(_) => {
                Err("Claude token refresh timed out. Try again or sign in again.".to_string())
            }
        }
    }

    struct Refresh {
        id: uuid::Uuid,
        replies: Vec<Reply>,
    }

    /// Marshals secure-storage access to the UI thread and coalesces concurrent refreshes.
    pub(crate) struct ClaudeOAuthModel {
        refreshing: HashMap<String, Refresh>,
    }
    impl Entity for ClaudeOAuthModel {
        type Event = ();
    }
    impl SingletonEntity for ClaudeOAuthModel {}

    impl ClaudeOAuthModel {
        pub(crate) fn new(ctx: &mut ModelContext<Self>) -> Self {
            let (tx, rx) = async_channel::bounded(64);
            *SENDER.get_or_init(|| Mutex::new(None)).lock().unwrap() = Some(tx);
            ctx.spawn_stream_local(
                rx,
                |me, command, ctx| match command {
                    Command::Resolve(reference, reply) => me.resolve(reference, reply, ctx),
                    Command::Invalidate(reference) => me.invalidate(&reference),
                    Command::DiscardDraft(reference) => {
                        let saved = AISettings::as_ref(ctx)
                            .openai_providers
                            .value()
                            .iter()
                            .any(|provider| provider.api_key.as_deref() == Some(&reference));
                        if !saved {
                            me.invalidate(&reference);
                            let _ = ctx.secure_storage().remove_value(&reference);
                        }
                    }
                },
                |_, _| {},
            );
            Self {
                refreshing: HashMap::new(),
            }
        }

        fn invalidate(&mut self, reference: &str) {
            if let Some(refresh) = self.refreshing.remove(reference) {
                for reply in refresh.replies {
                    let _ = reply.send(Err("Claude subscription was signed out.".to_string()));
                }
            }
        }

        fn resolve(&mut self, reference: String, reply: Reply, ctx: &mut ModelContext<Self>) {
            if reply.is_canceled() {
                return;
            }
            if let Some(refresh) = self.refreshing.get_mut(&reference) {
                refresh.replies.push(reply);
                return;
            }
            let value = match ctx.secure_storage().read_value(&reference) {
                Ok(value) => value,
                Err(_) => {
                    let _ = reply.send(Err(
                        "Claude subscription is signed out. Sign in again.".to_string()
                    ));
                    return;
                }
            };
            if !reference.starts_with(OAUTH_REFERENCE_PREFIX) {
                let _ = reply.send(Ok(value));
                return;
            }
            let credentials: Credentials = match serde_json::from_str(&value) {
                Ok(credentials) => credentials,
                Err(_) => {
                    let _ = reply.send(Err(
                        "Claude credentials are invalid. Sign in again.".to_string()
                    ));
                    return;
                }
            };
            if credentials.expires_at > now() {
                let _ = reply.send(Ok(credentials.access_token));
                return;
            }
            let id = uuid::Uuid::new_v4();
            self.refreshing.insert(
                reference.clone(),
                Refresh {
                    id,
                    replies: vec![reply],
                },
            );
            ctx.spawn(
                async move {
                    token_request(serde_json::json!({
                        "grant_type": "refresh_token", "client_id": CLIENT_ID,
                        "refresh_token": credentials.refresh_token,
                    }))
                    .await
                },
                move |me, result, ctx| {
                    if me
                        .refreshing
                        .get(&reference)
                        .is_none_or(|refresh| refresh.id != id)
                    {
                        return;
                    }
                    let refresh = me.refreshing.remove(&reference).unwrap();
                    let result = result.and_then(|credentials| {
                        // Signing out deletes the store entry before this callback can commit.
                        if ctx.secure_storage().read_value(&reference).is_err() {
                            return Err("Claude subscription was signed out.".to_string());
                        }
                        save_credentials(&reference, &credentials, ctx)?;
                        Ok(credentials.access_token)
                    });
                    for reply in refresh.replies {
                        let _ = reply.send(result.clone());
                    }
                },
            );
        }
    }

    fn now() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs()
    }

    pub(crate) fn save_credentials(
        reference: &str,
        credentials: &Credentials,
        ctx: &galaxyui::AppContext,
    ) -> Result<(), String> {
        let json = serde_json::to_string(credentials)
            .map_err(|_| "Could not encode Claude credentials.".to_string())?;
        ctx.secure_storage()
            .write_value_with_owner_only_fallback(reference, &json)
            .map_err(|_| {
                "Could not securely save Claude credentials. Sign-in was not saved.".to_string()
            })
    }

    #[derive(serde::Deserialize)]
    struct TokenResponse {
        access_token: String,
        refresh_token: String,
        expires_in: u64,
    }

    async fn token_request(body: serde_json::Value) -> Result<Credentials, String> {
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(30))
            .redirect(reqwest::redirect::Policy::none())
            .build()
            .map_err(|_| "Could not create Claude authorization client.".to_string())?;
        let response = client
            .post(TOKEN_URL)
            .header("Accept", "application/json")
            .json(&body)
            .send()
            .await
            .map_err(|_| {
                "Claude authorization request failed. Try signing in again.".to_string()
            })?;
        if !response.status().is_success() {
            // Never expose response bodies: providers may echo authorization codes or tokens.
            return Err(format!(
                "Claude authorization failed (HTTP {}). Sign in again.",
                response.status().as_u16()
            ));
        }
        if response
            .content_length()
            .is_some_and(|length| length > 64 * 1024)
        {
            return Err("Claude authorization returned an oversized response.".to_string());
        }
        let mut bytes = Vec::new();
        let mut stream = response.bytes_stream();
        while let Some(chunk) = stream.next().await {
            let chunk = chunk.map_err(|_| "Claude authorization response failed.".to_string())?;
            if bytes.len().saturating_add(chunk.len()) > 64 * 1024 {
                return Err("Claude authorization returned an oversized response.".to_string());
            }
            bytes.extend_from_slice(&chunk);
        }
        let tokens: TokenResponse = serde_json::from_slice(&bytes)
            .map_err(|_| "Claude authorization returned invalid credentials.".to_string())?;
        credentials_from_response(tokens)
    }

    fn credentials_from_response(tokens: TokenResponse) -> Result<Credentials, String> {
        if !tokens.access_token.contains("sk-ant-oat")
            || tokens.refresh_token.is_empty()
            || tokens.expires_in == 0
        {
            return Err("Claude authorization returned incomplete credentials.".to_string());
        }
        Ok(Credentials {
            access_token: tokens.access_token,
            refresh_token: tokens.refresh_token,
            expires_at: now()
                .saturating_add(tokens.expires_in)
                .saturating_sub(5 * 60),
        })
    }

    pub(crate) struct BrowserLogin {
        pub(crate) url: Url,
        verifier: String,
        callback: oneshot::Receiver<Result<String, String>>,
        cancelled: Arc<AtomicBool>,
    }
    impl Drop for BrowserLogin {
        fn drop(&mut self) {
            self.cancelled.store(true, Ordering::Relaxed);
        }
    }
    impl BrowserLogin {
        pub(crate) fn start() -> Result<Self, String> {
            let listener = TcpListener::bind(("127.0.0.1", 53692))
                .map_err(|_| "Claude sign-in needs localhost port 53692. Close another sign-in attempt and try again.".to_string())?;
            listener
                .set_nonblocking(true)
                .map_err(|_| "Could not configure Claude sign-in callback.".to_string())?;
            let mut bytes = [0; 32];
            rand::thread_rng().fill_bytes(&mut bytes);
            let verifier = URL_SAFE_NO_PAD.encode(bytes);
            let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()));
            let url = authorize_url(&verifier, &challenge)?;
            let cancelled = Arc::new(AtomicBool::new(false));
            let stop = cancelled.clone();
            let state = verifier.clone();
            let (tx, callback) = oneshot::channel();
            std::thread::Builder::new()
                .name("claude-oauth-callback".to_string())
                .spawn(move || {
                    let deadline = Instant::now() + SIGN_IN_TIMEOUT;
                    while !stop.load(Ordering::Relaxed) && Instant::now() < deadline {
                        match listener.accept() {
                            Ok((stream, _)) => {
                                if let Some(code) = read_callback(stream, &state) {
                                    let _ = tx.send(Ok(code));
                                    return;
                                }
                            }
                            Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                                std::thread::sleep(Duration::from_millis(50))
                            }
                            Err(_) => {
                                let _ = tx.send(Err("Claude sign-in callback failed.".to_string()));
                                return;
                            }
                        }
                    }
                    let _ = tx.send(Err(
                        "Claude browser sign-in expired or was cancelled.".to_string()
                    ));
                })
                .map_err(|_| "Could not start Claude sign-in callback.".to_string())?;
            Ok(Self {
                url,
                verifier,
                callback,
                cancelled,
            })
        }

        pub(crate) async fn finish(mut self) -> Result<Credentials, String> {
            let code = (&mut self.callback)
                .await
                .map_err(|_| "Claude sign-in callback stopped.".to_string())??;
            token_request(serde_json::json!({
                "grant_type": "authorization_code", "client_id": CLIENT_ID, "code": code,
                "state": self.verifier, "redirect_uri": REDIRECT_URI, "code_verifier": self.verifier,
            })).await
        }
    }

    pub(super) fn authorize_url(verifier: &str, challenge: &str) -> Result<Url, String> {
        let mut url = Url::parse(AUTHORIZE_URL)
            .map_err(|_| "Invalid Claude authorization URL.".to_string())?;
        url.query_pairs_mut()
            .append_pair("code", "true")
            .append_pair("client_id", CLIENT_ID)
            .append_pair("response_type", "code")
            .append_pair("redirect_uri", REDIRECT_URI)
            .append_pair("scope", SCOPES)
            .append_pair("code_challenge", challenge)
            .append_pair("code_challenge_method", "S256")
            .append_pair("state", verifier);
        Ok(url)
    }

    pub(super) fn callback_code(request: &str, state: &str) -> Option<String> {
        let mut parts = request.split_whitespace();
        if parts.next()? != "GET" {
            return None;
        }
        let path = parts.next()?;
        if !path.starts_with('/') || path.starts_with("//") {
            return None;
        }
        let url = Url::parse(&format!("http://localhost:53692{path}")).ok()?;
        if url.path() != "/callback" {
            return None;
        }
        let params: HashMap<_, _> = url.query_pairs().into_owned().collect();
        if params.get("state")?.as_str() != state {
            return None;
        }
        params.get("code").filter(|code| !code.is_empty()).cloned()
    }

    fn read_callback(mut stream: TcpStream, state: &str) -> Option<String> {
        stream.set_read_timeout(Some(Duration::from_secs(1))).ok()?;
        stream
            .set_write_timeout(Some(Duration::from_secs(1)))
            .ok()?;
        let mut request = String::new();
        // Bound the request line independently of whether a newline is supplied.
        std::io::Read::take(BufReader::new(&mut stream), 8192)
            .read_line(&mut request)
            .ok()?;
        let code = callback_code(&request, state);
        let (status, message) = if code.is_some() {
            (
                "200 OK",
                "Signed in to Claude. You can close this tab and return to Galaxy.",
            )
        } else {
            ("400 Bad Request", "This callback did not match the active Galaxy sign-in. Return to Galaxy and try again.")
        };
        let body = format!("<!doctype html><title>Galaxy — Claude sign-in</title><p>{message}</p>");
        let response = format!("HTTP/1.1 {status}\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nCache-Control: no-store\r\nConnection: close\r\n\r\n{body}", body.len());
        let _ = stream.write_all(response.as_bytes());
        code
    }
}

#[cfg(all(test, not(target_family = "wasm")))]
#[path = "claude_oauth_tests.rs"]
mod tests;