1use serde::Serialize;
2use std::io::Write;
3use topo_core::ScoredFile;
4
5pub struct JsonlWriter {
7 query: String,
8 preset: String,
9 max_bytes: Option<u64>,
10 min_score: f64,
11}
12
13#[derive(Serialize)]
14#[serde(rename_all = "PascalCase")]
15struct Header {
16 version: String,
17 query: String,
18 preset: String,
19 budget: Budget,
20 min_score: f64,
21}
22
23#[derive(Serialize)]
24#[serde(rename_all = "PascalCase")]
25struct Budget {
26 #[serde(skip_serializing_if = "Option::is_none")]
27 max_bytes: Option<u64>,
28}
29
30#[derive(Serialize)]
31#[serde(rename_all = "PascalCase")]
32struct FileEntry {
33 path: String,
34 score: f64,
35 tokens: u64,
36 language: String,
37 role: String,
38}
39
40#[derive(Serialize)]
41#[serde(rename_all = "PascalCase")]
42struct Footer {
43 total_files: usize,
44 total_tokens: u64,
45 scanned_files: usize,
46}
47
48impl JsonlWriter {
49 pub fn new(query: &str, preset: &str) -> Self {
50 Self {
51 query: query.to_string(),
52 preset: preset.to_string(),
53 max_bytes: None,
54 min_score: 0.0,
55 }
56 }
57
58 pub fn max_bytes(mut self, max_bytes: Option<u64>) -> Self {
59 self.max_bytes = max_bytes;
60 self
61 }
62
63 pub fn min_score(mut self, min_score: f64) -> Self {
64 self.min_score = min_score;
65 self
66 }
67
68 pub fn render(&self, files: &[ScoredFile], scanned_count: usize) -> anyhow::Result<String> {
70 let mut buf = Vec::new();
71 self.write_to(&mut buf, files, scanned_count)?;
72 Ok(String::from_utf8(buf)?)
73 }
74
75 pub fn write_to(
77 &self,
78 writer: &mut dyn Write,
79 files: &[ScoredFile],
80 scanned_count: usize,
81 ) -> anyhow::Result<()> {
82 let header = Header {
84 version: "0.3".to_string(),
85 query: self.query.clone(),
86 preset: self.preset.clone(),
87 budget: Budget {
88 max_bytes: self.max_bytes,
89 },
90 min_score: self.min_score,
91 };
92 serde_json::to_writer(&mut *writer, &header)?;
93 writeln!(writer)?;
94
95 let mut total_tokens = 0u64;
97 for file in files {
98 let entry = FileEntry {
99 path: file.path.clone(),
100 score: file.score,
101 tokens: file.tokens,
102 language: file.language.as_str().to_string(),
103 role: file.role.as_str().to_string(),
104 };
105 serde_json::to_writer(&mut *writer, &entry)?;
106 writeln!(writer)?;
107 total_tokens += file.tokens;
108 }
109
110 let footer = Footer {
112 total_files: files.len(),
113 total_tokens,
114 scanned_files: scanned_count,
115 };
116 serde_json::to_writer(&mut *writer, &footer)?;
117 writeln!(writer)?;
118
119 Ok(())
120 }
121}