//! Searching the code of public repositories, with Zoekt
//! (github.com/sourcegraph/zoekt): a trigram index, its shards on disk,
//! that answers substrings and regular expressions across every
//! repository at once.
//!
//! knotless-serve runs `zoekt-webserver` over the shards' directory, on a
//! port of its own on loopback, and asks it in JSON
//! ([`Code::search`]); `zoekt-git-index` writes a repository's shards
//! from its mirror's default branch ([`Code::index`]). The webserver
//! loads the shards as they are written, so a push is searched once its
//! repository is indexed again. Neither is reached from outside.
//!
//! Which repositories are indexed follows the index ([`follow`]): each
//! public one with no shards yet once at start, then each again as its
//! refs or its record move. One gone from the index, or taken down, has
//! its shards deleted as it is indexed again, and is never answered
//! meanwhile: what Zoekt finds is shown only for repositories the index
//! lists ([`crate::index::Index::listed`]).

use crate::{App, AppError};
use anyhow::{Context, Result};
use base64::Engine;
use futures::StreamExt;
use knotless::status::CodeIndex;
use knotless::telemetry::{self, Fields, Level};
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;

/// How long a search may take in Zoekt.
const WALL: Duration = Duration::from_secs(10);

/// The most files counted by language for a search.
const FACETED: usize = 1000;

/// How long `zoekt-webserver` has to start answering.
const START: Duration = Duration::from_secs(30);

/// Zoekt, over the shards in `dir`.
pub struct Code {
    dir: PathBuf,
    url: String,
    http: reqwest::Client,
    _server: tokio::process::Child,
}

/// A file Zoekt found, in the repository with `key` of `did`.
#[derive(Debug)]
pub struct Hit {
    pub did: String,
    pub key: String,
    pub path: String,
    pub language: String,
    /// The commit of its default branch indexed, which its lines are of.
    pub commit: String,
    pub chunks: Vec<Chunk>,
}

/// Lines of a file around what was found in them.
#[derive(Serialize, Debug, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct Chunk {
    /// The first line's number, from 1.
    pub start_line: u32,
    /// Its lines, joined by newlines, with none at the end.
    pub content: String,
    pub ranges: Vec<Found>,
    /// What its highlighting paints over `content`, as a file's is
    /// encoded ([`crate::highlight::Ranges::encode`]): the lines read
    /// alone, without the rest of their file.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub highlights: Option<String>,
}

/// What was found on one line, by UTF-8 bytes into it.
#[derive(Serialize, Debug, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct Found {
    pub line: u32,
    pub byte_start: usize,
    pub byte_end: usize,
}

/// What a search found: the files shown, and how many files and matches
/// there were in all.
#[derive(Debug, Default)]
pub struct Results {
    pub hits: Vec<Hit>,
    pub file_count: u64,
    pub match_count: u64,
}

// What Zoekt's JSON API answers, as its Go types are named.
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct Reply {
    result: Option<ZResult>,
    error: Option<String>,
}

#[derive(Deserialize, Default)]
#[serde(rename_all = "PascalCase")]
struct ZResult {
    file_count: u64,
    match_count: u64,
    #[serde(default)]
    files: Option<Vec<ZFile>>,
}

#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct ZFile {
    #[serde(default)]
    score: f64,
    file_name: String,
    repository: String,
    language: String,
    #[serde(default)]
    version: String,
    #[serde(default)]
    chunk_matches: Option<Vec<ZChunk>>,
}

#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct ZChunk {
    content: String,
    content_start: ZLocation,
    ranges: Vec<ZRange>,
    file_name: bool,
}

#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct ZLocation {
    byte_offset: usize,
    line_number: u32,
}

#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct ZRange {
    start: ZLocation,
    end: ZLocation,
}

impl Code {
    /// Starts `zoekt-webserver` over `dir`, made if it is missing, and
    /// waits for it to answer.
    pub async fn start(dir: &Path) -> Result<Code> {
        std::fs::create_dir_all(dir.join("meta"))?;
        let port = std::net::TcpListener::bind("127.0.0.1:0")?
            .local_addr()?
            .port();
        let listen = format!("127.0.0.1:{port}");
        let server = tokio::process::Command::new("zoekt-webserver")
            .arg("-index")
            .arg(dir)
            .args(["-listen", &listen, "-rpc", "-html=false"])
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .kill_on_drop(true)
            .spawn()
            .context("zoekt-webserver, which code search runs, is not on PATH")?;
        let url = format!("http://{listen}");
        let http = reqwest::Client::new();
        let health = format!("{url}/healthz");
        let up = async {
            while http.get(&health).send().await.is_err() {
                tokio::time::sleep(Duration::from_millis(50)).await;
            }
        };
        tokio::time::timeout(START, up)
            .await
            .context("zoekt-webserver did not answer")?;
        Ok(Code {
            dir: dir.to_path_buf(),
            url,
            http,
            _server: server,
        })
    }

    /// Writes the shards of `did`'s repository with `key` from its mirror
    /// at `git`, of its default branch, over any it had: a commit indexed
    /// already is left as it is.
    pub async fn index(&self, did: &str, key: &str, git: &Path) -> Result<()> {
        let prefix = prefix(did, key);
        let meta = self.dir.join("meta").join(format!("{prefix}.json"));
        std::fs::write(&meta, json!({ "Name": format!("{did}/{key}") }).to_string())?;
        let out = tokio::process::Command::new("zoekt-git-index")
            .arg("-index")
            .arg(&self.dir)
            .args(["-shard_prefix_override", &prefix, "-meta"])
            .arg(&meta)
            .args(["-branches", "HEAD", "-submodules=false", "-disable_ctags"])
            .arg(git)
            .output()
            .await
            .context("zoekt-git-index, which code search runs, is not on PATH")?;
        let said = String::from_utf8_lossy(&out.stderr);
        anyhow::ensure!(out.status.success(), "zoekt-git-index: {said}");
        Ok(())
    }

    /// The shards of `did`'s repository with `key`.
    fn shards(&self, did: &str, key: &str) -> Vec<PathBuf> {
        let start = format!("{}_v", prefix(did, key));
        let entries = std::fs::read_dir(&self.dir).into_iter().flatten().flatten();
        let named = |p: &PathBuf| {
            p.file_name()
                .is_some_and(|n| n.to_string_lossy().starts_with(&start))
        };
        entries.map(|e| e.path()).filter(named).collect()
    }

    /// The index as `/_status` tells of it: whether Zoekt answers, and
    /// the repositories, shards and bytes it holds.
    pub async fn check(&self) -> CodeIndex {
        let mut repos = std::collections::BTreeSet::new();
        let (mut shards, mut bytes) = (0, 0);
        for entry in std::fs::read_dir(&self.dir).into_iter().flatten().flatten() {
            let name = entry.file_name().to_string_lossy().into_owned();
            // Its `meta` directory is no shard.
            let Some((repo, _)) = name
                .strip_suffix(".zoekt")
                .and_then(|n| n.rsplit_once("_v"))
            else {
                continue;
            };
            repos.insert(repo.to_string());
            shards += 1;
            bytes += entry.metadata().map_or(0, |m| m.len());
        }
        let health = format!("{}/healthz", self.url);
        let answered = self.http.get(health).send().await;
        let error = answered.and_then(|r| r.error_for_status()).err();
        CodeIndex {
            ok: error.is_none(),
            error: error.as_ref().map(ToString::to_string),
            repos: repos.len() as u64,
            shards,
            bytes,
        }
    }

    /// Whether `did`'s repository with `key` has shards.
    pub fn has(&self, did: &str, key: &str) -> bool {
        !self.shards(did, key).is_empty()
    }

    /// Deletes the shards of `did`'s repository with `key`.
    pub fn forget(&self, did: &str, key: &str) {
        for shard in self.shards(did, key) {
            let _ = std::fs::remove_file(shard);
        }
        let _ = std::fs::remove_file(
            self.dir
                .join("meta")
                .join(format!("{}.json", prefix(did, key))),
        );
    }

    /// The files at most `limit` that `q`, in Zoekt's query language,
    /// finds, in `did`'s repository with `key` alone if one is named: 400
    /// for a query Zoekt cannot read.
    pub async fn search(
        &self,
        q: &str,
        repo: Option<(&str, &str)>,
        limit: usize,
    ) -> Result<Results, AppError> {
        let opts = json!({
            "ChunkMatches": true,
            "NumContextLines": 1,
            "MaxDocDisplayCount": limit,
            "MaxWallTime": WALL.as_nanos() as u64,
        });
        let found = self.ask(within(q, repo), opts).await?;
        let mut files = found.files.unwrap_or_default();
        // The best first, and files Zoekt scores alike, as its shards
        // answered in parallel, by repository and path: so each page
        // after the first goes on from it rather than showing it again.
        files.sort_by(|a, b| {
            let named = a
                .repository
                .cmp(&b.repository)
                .then(a.file_name.cmp(&b.file_name));
            b.score.total_cmp(&a.score).then(named)
        });
        let hits = files.into_iter().filter_map(hit).collect();
        Ok(Results {
            hits,
            file_count: found.file_count,
            match_count: found.match_count,
        })
    }

    /// The files `q` finds in any language, [`FACETED`] at most, with no
    /// lines: what its `lang:` filters leave out counted too, so each
    /// language can be picked from them.
    pub async fn faceted(&self, q: &str, repo: Option<(&str, &str)>) -> Result<Vec<Hit>, AppError> {
        let opts = json!({
            "MaxDocDisplayCount": FACETED,
            "MaxWallTime": WALL.as_nanos() as u64,
        });
        let q = format!("type:file ({})", unlanguaged(q));
        let found = self.ask(within(&q, repo), opts).await?;
        let files = found.files.unwrap_or_default().into_iter();
        let bare = files.map(|f| ZFile {
            chunk_matches: None,
            ..f
        });
        Ok(bare.filter_map(hit).collect())
    }

    /// What Zoekt answers `q` with `opts`.
    async fn ask(&self, q: String, opts: serde_json::Value) -> Result<ZResult, AppError> {
        let asked = json!({ "Q": q, "Opts": opts });
        let url = format!("{}/api/search", self.url);
        let res = self.http.post(url).json(&asked).send().await;
        let res = res
            .map_err(anyhow::Error::from)
            .map_err(AppError::internal)?;
        // Zoekt's status, as a query it cannot read is 400.
        let status = res.status();
        let reply = res.json::<Reply>().await;
        let reply = reply
            .map_err(anyhow::Error::from)
            .map_err(AppError::internal)?;
        if let Some(error) = reply.error {
            return Err(AppError(status, format!("code search: {error}")));
        }
        Ok(reply.result.unwrap_or_default())
    }
}

/// A file as Zoekt found it, if its repository is named as [`Code::index`]
/// names one.
fn hit(file: ZFile) -> Option<Hit> {
    let (did, key) = file.repository.rsplit_once('/')?;
    let chunks = file.chunk_matches.unwrap_or_default();
    let mut chunks: Vec<Chunk> = chunks
        .into_iter()
        .filter(|c| !c.file_name)
        .map(|c| chunk(c, &file.file_name))
        .collect();
    // Zoekt answers them best first: shown as the file reads.
    chunks.sort_by_key(|c| c.start_line);
    Some(Hit {
        did: did.to_string(),
        key: key.to_string(),
        path: file.file_name,
        language: file.language,
        commit: file.version,
        chunks,
    })
}

/// A chunk's lines, and what was found, by line, out of Zoekt's offsets
/// into the file.
fn chunk(c: ZChunk, path: &str) -> Chunk {
    let bytes = base64::engine::general_purpose::STANDARD
        .decode(&c.content)
        .unwrap_or_default();
    let bytes = bytes.strip_suffix(b"\n").unwrap_or(&bytes);
    let base = c.content_start.byte_offset;
    let mut starts = Vec::new();
    let mut at = 0;
    let lines: Vec<&[u8]> = bytes.split(|b| *b == b'\n').collect();
    for line in &lines {
        starts.push(at);
        at += line.len() + 1;
    }
    let first = c.content_start.line_number;
    let mut ranges = Vec::new();
    for r in &c.ranges {
        let (s, e) = (r.start.byte_offset - base, r.end.byte_offset - base);
        for (i, (start, line)) in starts.iter().zip(&lines).enumerate() {
            let (a, b) = (s.max(*start), e.min(start + line.len()));
            if a < b {
                ranges.push(Found {
                    line: first + i as u32,
                    byte_start: a - start,
                    byte_end: b - start,
                });
            }
        }
    }
    let content = String::from_utf8_lossy(bytes).into_owned();
    let painted = crate::highlight::code(path, &content);
    Chunk {
        start_line: first,
        highlights: painted.as_deref().map(crate::highlight::Ranges::encode),
        content,
        ranges,
    }
}

/// `q` kept to `did`'s repository with `key`, if one is named.
fn within(q: &str, repo: Option<(&str, &str)>) -> String {
    match repo {
        Some((did, key)) => format!("repo:^{}$ ({q})", escape(&format!("{did}/{key}"))),
        None => q.to_string(),
    }
}

/// `q` without its `lang:` and `-lang:` filters, its words joined by
/// single spaces, or `q` itself when nothing else is left of it. An
/// operator joining a filter to the rest goes with it, and a `)` it ends
/// with stays, so what is left still reads as a query.
fn unlanguaged(q: &str) -> String {
    let operator = |w: &str| matches!(w, "or" | "and" | "not");
    let mut kept: Vec<&str> = Vec::new();
    // A filter was just taken out with no operator before it: one after
    // it goes instead.
    let mut after = false;
    let mut words = q.split_whitespace();
    while let Some(word) = words.next() {
        let filter = word.strip_prefix('-').unwrap_or(word).strip_prefix("lang:");
        let Some(name) = filter else {
            if !(after && operator(word)) {
                kept.push(word);
            }
            after = false;
            continue;
        };
        // A quoted name runs on to the word its quote closes in.
        let mut end = name;
        if name.starts_with('"') && !name[1..].contains('"') {
            end = words.by_ref().find(|w| w.contains('"')).unwrap_or("");
        }
        let closing = end.len() - end.trim_end_matches(')').len();
        after = !kept.last().is_some_and(|w| operator(w));
        if !after {
            kept.pop();
        }
        kept.extend(std::iter::repeat_n(")", closing));
    }
    if kept.is_empty() {
        return q.to_string();
    }
    kept.join(" ")
}

/// The prefix of the shards of `did`'s repository with `key`: `did/key`,
/// each byte but letters, digits, `.` and `-` written `_` and its hex, so
/// that it is a file name Zoekt keeps as it is, and no two are one.
fn prefix(did: &str, key: &str) -> String {
    format!("{did}/{key}")
        .bytes()
        .map(|b| match b {
            b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'.' | b'-' => (b as char).to_string(),
            _ => format!("_{b:02x}"),
        })
        .collect()
}

/// `s` as a regular expression matching it alone.
fn escape(s: &str) -> String {
    s.chars()
        .map(|c| match c.is_ascii_alphanumeric() {
            true => c.to_string(),
            false => format!("\\{c}"),
        })
        .collect()
}

/// What indexes the public repositories as the index asks
/// ([`Index::reindex`]), when serve keeps a code index: it asks first for
/// each with no shards yet, then answers each ask, for as long as serve
/// runs. Set up here, before it is spawned, so serve without a code
/// index starts nothing.
///
/// [`Index::reindex`]: crate::index::Index::reindex
pub fn follow(app: Arc<App>) -> Option<impl Future<Output = ()>> {
    let (code, index) = app.code.clone().zip(app.index.clone())?;
    let asked = index.indexing()?;
    for (did, key) in index.code_repos().unwrap_or_default() {
        if !code.has(&did, &key) {
            index.reindex(&did, &key);
        }
    }
    let asked = futures::stream::unfold(asked, |mut a| async { a.recv().await.map(|r| (r, a)) });
    Some(asked.for_each(move |(did, key)| indexed(app.clone(), code.clone(), did, key)))
}

/// Indexes `did`'s repository with `key` again, as [`follow`] is asked,
/// and notes how it went.
async fn indexed(app: Arc<App>, code: Arc<Code>, did: String, key: String) {
    let error = reindex(&app, &code, &did, &key).await.err();
    let error = error.as_ref().map(ToString::to_string);
    let level = [Level::INFO, Level::WARN][error.is_some() as usize];
    let fields = Fields {
        did: Some(&did),
        key: Some(&key),
        error,
        ..Fields::default()
    };
    telemetry::note(level, "code indexed", &fields);
}

/// Indexes `did`'s repository with `key` again, or forgets it if the
/// index no longer lists it or its default branch has no commit to index.
async fn reindex(app: &App, code: &Code, did: &str, key: &str) -> Result<()> {
    let index = app.index.as_ref().context("no index")?;
    index.reindexed(did, key);
    let Some(listed) = index.listed(did, key)? else {
        code.forget(did, key);
        return Ok(());
    };
    let git = app.mirror(did, &listed.name).await.ok();
    let git = git.context("its mirror could not be read")?;
    if git.rev_parse("HEAD").await.is_err() {
        code.forget(did, key);
        return Ok(());
    }
    code.index(did, key, &git.dir).await
}

/// Zoekt over `dir`, as if its webserver had died: nothing answers.
#[cfg(test)]
pub(crate) fn dead(dir: &Path) -> Code {
    Code {
        dir: dir.to_path_buf(),
        url: "http://127.0.0.1:1".into(),
        http: reqwest::Client::new(),
        _server: tokio::process::Command::new("true").spawn().unwrap(),
    }
}

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

    #[test]
    fn languages_are_left_out_to_count_them() {
        assert_eq!(unlanguaged("count lang:Rust"), "count");
        assert_eq!(unlanguaged("lang:\"Vim Script\" a -lang:toml b"), "a b");
        assert_eq!(unlanguaged("lang:\"Go\" x"), "x");
        assert_eq!(unlanguaged("lang:Rust"), "lang:Rust");
        // An operator joining a filter goes with it, before or after.
        assert_eq!(unlanguaged("foo or lang:Rust"), "foo");
        assert_eq!(unlanguaged("lang:Rust or foo"), "foo");
        assert_eq!(unlanguaged("a or lang:Rust or b"), "a or b");
        // A group the filter closed stays closed.
        assert_eq!(unlanguaged("(foo or lang:Rust) bar"), "(foo ) bar");
        assert_eq!(unlanguaged("(a or lang:\"Vim Script\") b"), "(a ) b");
        assert_eq!(within("a", Some(("d", "k"))), "repo:^d\\/k$ (a)");
    }

    #[test]
    fn prefixes_are_file_names_of_their_own() {
        assert_eq!(prefix("did:plc:ab", "3k.x-y"), "did_3aplc_3aab_2f3k.x-y");
        assert_ne!(prefix("did:plc:a", "b_c"), prefix("did:plc:a", "b:c"));
        assert_eq!(escape("did:plc:a/k.1"), "did\\:plc\\:a\\/k\\.1");
    }

    #[test]
    fn chunks_are_cut_into_lines() {
        let content =
            base64::engine::general_purpose::STANDARD.encode("fn a() {}\nlet abc = a();\n");
        let at = |byte_offset, line_number| ZLocation {
            byte_offset,
            line_number,
        };
        let zc = ZChunk {
            content: content.clone(),
            // The chunk starts at byte 100, line 7, of its file.
            content_start: at(100, 7),
            // `a` in line 7, and from its `{}` to the end of the next line's `abc`.
            ranges: vec![
                ZRange {
                    start: at(103, 7),
                    end: at(104, 7),
                },
                ZRange {
                    start: at(107, 7),
                    end: at(117, 8),
                },
            ],
            file_name: false,
        };
        let c = chunk(zc, "src/a.rs");
        assert_eq!(c.start_line, 7);
        assert_eq!(c.content, "fn a() {}\nlet abc = a();");
        let found = |line, byte_start, byte_end| Found {
            line,
            byte_start,
            byte_end,
        };
        assert_eq!(c.ranges, [found(7, 3, 4), found(7, 7, 9), found(8, 0, 7)]);
        // Painted as Rust reads, by its path: `fn` and `let` keywords.
        assert!(c.highlights.unwrap().contains("keyword:0,2,"));
        // A file no grammar here reads is painted not at all.
        let plain = ZChunk {
            content,
            content_start: at(0, 1),
            ranges: vec![],
            file_name: false,
        };
        assert_eq!(chunk(plain, "notes.txt").highlights, None);
        // A file's chunks, in the order of its lines.
        let at_line = |line| ZChunk {
            content: base64::engine::general_purpose::STANDARD.encode("x\n"),
            content_start: at(0, line),
            ranges: vec![],
            file_name: false,
        };
        let found = hit(ZFile {
            score: 0.0,
            file_name: "a".into(),
            repository: "did:plc:a/k".into(),
            language: String::new(),
            version: String::new(),
            chunk_matches: Some(vec![at_line(30), at_line(2)]),
        })
        .unwrap();
        let lines: Vec<u32> = found.chunks.iter().map(|c| c.start_line).collect();
        assert_eq!(lines, [2, 30]);
        assert!(
            hit(ZFile {
                score: 0.0,
                file_name: "a".into(),
                repository: "no-slash".into(),
                language: String::new(),
                version: String::new(),
                chunk_matches: None,
            })
            .is_none()
        );
    }
}