paper-gap 0.1.4

use crate::Result;
use crate::arxiv::Paper;
use serde::{Deserialize, Serialize};
use std::time::Duration;

const BASE_URL: &str = "https://paperswithcode.com/api/v1";

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct CodeLink {
    pub source_provider: String,
    pub paper_id: Option<String>,
    pub repository_url: String,
    pub repository_name: Option<String>,
    pub is_official: Option<bool>,
}

#[derive(Debug, Deserialize)]
struct List<T> {
    results: Vec<T>,
}

#[derive(Debug, Deserialize)]
struct PwcPaper {
    id: String,
}

#[derive(Debug, Deserialize)]
struct PwcRepository {
    url: Option<String>,
    name: Option<String>,
    is_official: Option<bool>,
}

pub fn paper_search_url(paper: &Paper) -> String {
    if let Some(arxiv_id) = paper.arxiv_id.as_deref() {
        format!(
            "{BASE_URL}/papers/?arxiv_id={}",
            urlencoding::encode(trim_arxiv_version(arxiv_id))
        )
    } else {
        format!("{BASE_URL}/papers/?q={}", urlencoding::encode(&paper.title))
    }
}

pub fn paper_repositories_url(pwc_paper_id: &str) -> String {
    format!(
        "{BASE_URL}/papers/{}/repositories/",
        urlencoding::encode(pwc_paper_id)
    )
}

pub async fn code_links(paper: &Paper) -> Result<Vec<CodeLink>> {
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(8))
        .build()?;
    let paper_response = client.get(paper_search_url(paper)).send().await?;
    if paper_response.status() == reqwest::StatusCode::NOT_FOUND {
        return Ok(Vec::new());
    }
    let paper_response = paper_response.error_for_status()?;
    let papers = parse_paper_search(&paper_response.text().await?)?;
    let Some(pwc_paper) = papers.first() else {
        return Ok(Vec::new());
    };

    let repo_response = client
        .get(paper_repositories_url(&pwc_paper.id))
        .send()
        .await?;
    if repo_response.status() == reqwest::StatusCode::NOT_FOUND {
        return Ok(Vec::new());
    }
    let repo_response = repo_response.error_for_status()?;
    parse_repositories(&pwc_paper.id, &repo_response.text().await?)
}

fn parse_paper_search(json: &str) -> Result<Vec<PwcPaper>> {
    Ok(serde_json::from_str::<List<PwcPaper>>(json)?.results)
}

pub fn parse_repositories(paper_id: &str, json: &str) -> Result<Vec<CodeLink>> {
    let repos = serde_json::from_str::<List<PwcRepository>>(json)?;
    Ok(repos
        .results
        .into_iter()
        .filter_map(|repo| {
            Some(CodeLink {
                source_provider: "papers-with-code".into(),
                paper_id: Some(paper_id.into()),
                repository_url: repo.url?,
                repository_name: repo.name,
                is_official: repo.is_official,
            })
        })
        .collect())
}

fn trim_arxiv_version(arxiv_id: &str) -> &str {
    arxiv_id.rsplit_once('v').map_or(arxiv_id, |(id, version)| {
        if version.chars().all(|c| c.is_ascii_digit()) {
            id
        } else {
            arxiv_id
        }
    })
}