thought-stream-cli / src / client.rs

use anyhow::{Context, Result};
use chrono::Utc;
use reqwest::{Client as HttpClient, header::{HeaderMap, HeaderValue, AUTHORIZATION}};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::Duration;

use crate::credentials::Credentials;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
    #[serde(rename = "accessJwt")]
    pub access_jwt: String,
    #[serde(rename = "refreshJwt")]
    pub refresh_jwt: String,
    pub handle: String,
    pub did: String,
}

#[derive(Debug, Clone)]
pub struct AtProtoClient {
    http_client: HttpClient,
    base_url: String,
    session: Option<Session>,
}

#[derive(Debug, Serialize)]
struct LoginRequest {
    identifier: String,
    password: String,
}

#[derive(Debug, Serialize)]
struct CreateRecordRequest {
    repo: String,
    collection: String,
    record: Value,
}

#[derive(Debug, Serialize)]
struct BlipRecord {
    #[serde(rename = "$type")]
    record_type: String,
    content: String,
    #[serde(rename = "createdAt")]
    created_at: String,
}

#[derive(Debug, Deserialize)]
struct CreateRecordResponse {
    uri: String,
    cid: String,
}

#[derive(Debug, Deserialize)]
struct RefreshSessionResponse {
    #[serde(rename = "accessJwt")]
    access_jwt: String,
    #[serde(rename = "refreshJwt")]
    refresh_jwt: String,
    handle: String,
    did: String,
}

impl AtProtoClient {
    pub fn new(pds_uri: &str) -> Self {
        Self {
            http_client: HttpClient::new(),
            base_url: pds_uri.to_string(),
            session: None,
        }
    }

    pub async fn login(&mut self, credentials: &Credentials) -> Result<()> {
        let login_url = format!("{}/xrpc/com.atproto.server.createSession", self.base_url);
        
        let request = LoginRequest {
            identifier: credentials.username.clone(),
            password: credentials.password.clone(),
        };

        let response = self.http_client
            .post(&login_url)
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await
            .context("Failed to send login request")?;

        if !response.status().is_success() {
            let status = response.status();
            let error_text = response.text().await.unwrap_or_default();
            anyhow::bail!("Login failed: {} - {}", status, error_text);
        }

        let session: Session = response
            .json()
            .await
            .context("Failed to parse login response")?;

        println!("Authenticated as: {}", session.handle);
        self.session = Some(session);
        Ok(())
    }

    pub async fn publish_blip(&mut self, content: &str) -> Result<String> {
        // Try the request, and if it fails with auth error, refresh and retry once
        match self.try_publish_blip(content).await {
            Ok(uri) => Ok(uri),
            Err(e) => {
                // Check if this is an authentication error by examining the error message
                let error_msg = e.to_string();
                if error_msg.contains("401") || error_msg.contains("Unauthorized") ||
                   error_msg.contains("ExpiredToken") || error_msg.contains("Token has expired") {
                    // Try to refresh the session
                    match self.refresh_session().await {
                        Ok(_) => {
                            // Retry the request with the new token
                            self.try_publish_blip(content).await
                                .context("Failed to publish blip after session refresh")
                        }
                        Err(refresh_err) => {
                            // Session refresh failed - clear any stored session
                            if let Ok(store) = crate::credentials::CredentialStore::new() {
                                let _ = store.clear_session();
                            }
                            Err(anyhow::anyhow!(
                                "Authentication failed and session refresh failed: {}. Please run 'thought login' again.", 
                                refresh_err
                            ))
                        }
                    }
                } else {
                    Err(e)
                }
            }
        }
    }

    async fn try_publish_blip(&self, content: &str) -> Result<String> {
        let session = self.session.as_ref()
            .context("Not authenticated. Please run 'thought login' first.")?;

        let record = BlipRecord {
            record_type: "stream.thought.blip".to_string(),
            content: content.to_string(),
            created_at: Utc::now().to_rfc3339().replace("+00:00", "Z"),
        };

        let request = CreateRecordRequest {
            repo: session.did.clone(),
            collection: "stream.thought.blip".to_string(),
            record: serde_json::to_value(&record)
                .context("Failed to serialize blip record")?,
        };

        let create_url = format!("{}/xrpc/com.atproto.repo.createRecord", self.base_url);
        
        let mut headers = HeaderMap::new();
        headers.insert(
            AUTHORIZATION,
            HeaderValue::from_str(&format!("Bearer {}", session.access_jwt))
                .context("Invalid authorization header")?,
        );
        headers.insert(
            "Content-Type",
            HeaderValue::from_static("application/json"),
        );

        let response = self.http_client
            .post(&create_url)
            .headers(headers)
            .json(&request)
            .timeout(Duration::from_secs(30))
            .send()
            .await
            .context("Failed to send create record request (timeout or network error)")?;

        if !response.status().is_success() {
            let status = response.status();
            let error_text = response.text().await.unwrap_or_default();
            anyhow::bail!("Failed to publish blip: {} - {}", status, error_text);
        }

        let create_response: CreateRecordResponse = response
            .json()
            .await
            .context("Failed to parse create record response")?;

        Ok(create_response.uri)
    }

    pub fn is_authenticated(&self) -> bool {
        self.session.is_some()
    }

    pub fn get_user_did(&self) -> Option<String> {
        self.session.as_ref().map(|s| s.did.clone())
    }

    pub async fn refresh_session(&mut self) -> Result<()> {
        let session = self.session.as_ref()
            .context("No session to refresh. Please login first.")?;

        let refresh_url = format!("{}/xrpc/com.atproto.server.refreshSession", self.base_url);
        
        let mut headers = HeaderMap::new();
        headers.insert(
            AUTHORIZATION,
            HeaderValue::from_str(&format!("Bearer {}", session.refresh_jwt))
                .context("Invalid refresh token for authorization header")?,
        );

        let response = self.http_client
            .post(&refresh_url)
            .headers(headers)
            .timeout(Duration::from_secs(30))
            .send()
            .await
            .context("Failed to send refresh session request (timeout or network error)")?;

        if !response.status().is_success() {
            let status = response.status();
            let error_text = response.text().await.unwrap_or_default();
            anyhow::bail!("Session refresh failed: {} - {}", status, error_text);
        }

        let refresh_response: RefreshSessionResponse = response
            .json()
            .await
            .context("Failed to parse refresh session response")?;

        let new_session = Session {
            access_jwt: refresh_response.access_jwt,
            refresh_jwt: refresh_response.refresh_jwt,
            handle: refresh_response.handle,
            did: refresh_response.did,
        };

        self.session = Some(new_session);
        Ok(())
    }

    pub fn get_session(&self) -> Option<&Session> {
        self.session.as_ref()
    }

    pub fn set_session(&mut self, session: Session) {
        self.session = Some(session);
    }
}