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