3use std::time::{Duration, Instant};

5use base64::Engine as _;
6use log::{debug, error, info, trace, warn};
7use serde::Deserialize;
8use serde_json::json;
9use whiskers_core::{Model, ModelError, Speaker, Turn};
10
11pub struct Gateway {

13    base_url: String,
14    model: String,
15    max_tokens: u32,
16    agent: ureq::Agent,
17}

19impl Gateway {
20    pub fn new(base_url: impl Into<String>, model: impl Into<String>, max_tokens: u32) -> Self {
21        let (base_url, model) = (base_url.into(), model.into());
22        debug!("gateway client: model {model}, max {max_tokens} tokens, timeout 30 s");
23        let agent = ureq::Agent::config_builder()
24            .timeout_global(Some(Duration::from_secs(30)))
25            .http_status_as_error(false)
26            .build()
27            .into();
28        Self { base_url: base_url.trim_end_matches('/').to_owned(), model, max_tokens, agent }
29    }
30}
31
32#[derive(Deserialize)]
33struct Response {
34    content: Vec<Block>,
35}
36
37#[derive(Deserialize)]
38#[serde(tag = "type", rename_all = "snake_case")]
39enum Block {
40    Text { text: String },
41    #[serde(other)]
42    Other,
43}

46fn content(t: &Turn) -> serde_json::Value {
47    if t.pictures.is_empty() {
48        return json!(t.text);
49    }
50    trace!("turn carries {} picture(s)", t.pictures.len());
51    let mut blocks: Vec<_> = t
52        .pictures
53        .iter()
54        .map(|p| {
55            json!({ "type": "image", "source": {
56                "type": "base64",
57                "media_type": p.media_type,
58                "data": base64::engine::general_purpose::STANDARD.encode(&p.bytes),
59            }})
60        })
61        .collect();
62    // The API rejects an empty text block; a picture shown without words gets a nudge.
63    let text = if t.text.trim().is_empty() { "(The child is showing you this.)" } else { t.text.as_str() };
64    blocks.push(json!({ "type": "text", "text": text }));
65    json!(blocks)
66}

68impl Model for Gateway {
69    fn complete(&self, system: &str, history: &[Turn]) -> Result<String, ModelError> {
70        // The API wants the first message from the user; a window cut can
71        // start on one of Whiskers' own lines.
72        debug!("gateway request: system {} chars, {} turns", system.len(), history.len());
73        let started = Instant::now();
74        let messages: Vec<_> = history
75            .iter()
76            .skip_while(|t| t.speaker == Speaker::Whiskers)
77            .map(|t| {
78                let role = if t.speaker == Speaker::Child { "user" } else { "assistant" };
79                json!({ "role": role, "content": content(t) })
80            })
81            .collect();
82        let body = json!({
83            "model": self.model,
84            "max_tokens": self.max_tokens,
85            "system": system,
86            "messages": messages,
87        });
88        let mut resp = self
89            .agent
90            .post(format!("{}/v1/messages", self.base_url))
91            .header("content-type", "application/json")
92            .header("anthropic-version", "2023-06-01")
93            // The gateway ignores the key; the header keeps SDK-shaped
94            // servers happy.
95            .header("x-api-key", "gateway")
96            .send(body.to_string())
97            .map_err(|e| {
98                error!("gateway unreachable after {} ms: {e}", started.elapsed().as_millis());
99                ModelError::new(e.to_string())
100            })?;
101        let status = resp.status();
102        let text = resp.body_mut().read_to_string().map_err(|e| {
103            error!("gateway body unreadable (status {status}) after {} ms: {e}", started.elapsed().as_millis());
104            ModelError::new(e.to_string())
105        })?;
106        info!("gateway answered {status} in {} ms, {} bytes", started.elapsed().as_millis(), text.len());
107        if status.as_u16() == 429 && text.contains("\"budget\"") {
108            // The service's allowance for thinking is spent; not a fault, and said differently.
109            warn!("gateway says the thinking allowance is spent (429 budget)");
110            return Err(ModelError::resting("the thinking allowance is spent"));
111        }
112        if !status.is_success() {
113            error!("gateway refused with {status} ({} byte body)", text.len());
114            return Err(ModelError::new(format!("gateway answered {status}: {text}")));
115        }
116        let parsed: Response = serde_json::from_str(&text).map_err(|e| {
117            error!("gateway reply does not parse ({} bytes): {e}", text.len());
118            ModelError::new(e.to_string())
119        })?;
120        // Thinking and tool blocks are not for her ears; only text is kept.
121        let answer: String = parsed
122            .content
123            .into_iter()
124            .filter_map(|b| match b {
125                Block::Text { text } => Some(text),
126                Block::Other => None,
127            })
128            .collect();
129        if answer.trim().is_empty() {
130            warn!("gateway answered with no text blocks");
131            return Err(ModelError::new("gateway answered with no text"));
132        }
133        debug!("gateway answer: {} chars", answer.len());
134        Ok(answer)
135    }
136}