mizuki.guru/src/main.rs

use std::{
    fs,
    io::{BufRead, BufReader, ErrorKind, Read, Write},
    net::{TcpListener, TcpStream},
    path::Path,
    process::Command,
    sync::Arc,
    thread,
    time::Duration,
};

use serde::Deserialize;
use serde_json::{json, Value};

const PORT: u16 = 1108;
const API: &str = "https://akiyama.mizuki.guru/api/post/random/url";
const CONFIG_PATH: &str = "config.toml";
const MODEL_ID: &str = "gemma-4-26b-a4b-it";
const GENERATE_CONTENT_API: &str = "streamGenerateContent";

#[derive(Clone, Deserialize)]
struct Config {
    google: GoogleConfig,
}

#[derive(Clone, Deserialize)]
struct GoogleConfig {
    #[serde(rename = "api-key")]
    api_key: String,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let config = Arc::new(load_config()?);
    let listener = TcpListener::bind(("0.0.0.0", PORT))?;
    println!("listening on :{PORT}");

    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                let config = Arc::clone(&config);
                thread::spawn(move || {
                    if let Err(e) = handle_client(stream, &config) {
                        eprintln!("client error: {e}");
                    }
                });
            }
            Err(e) => eprintln!("accept error: {e}"),
        }
    }

    Ok(())
}

fn load_config() -> Result<Config, Box<dyn std::error::Error>> {
    let config = fs::read_to_string(CONFIG_PATH)?;
    let config = toml::from_str(&config)?;
    Ok(config)
}

fn handle_client(
    mut stream: TcpStream,
    config: &Config,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut reader = BufReader::new(stream.try_clone()?);
    let mut history = Vec::new();
    let url = ureq::get(API).call()?.into_string()?;
    let url = url.trim();

    writeln!(
        stream,
        r#"███╗   ███╗██╗███████╗██╗   ██╗██╗  ██╗██╗    ██████╗ ██╗   ██╗██████╗ ██╗   ██╗
████╗ ████║██║╚══███╔╝██║   ██║██║ ██╔╝██║   ██╔════╝ ██║   ██║██╔══██╗██║   ██║
██╔████╔██║██║  ███╔╝ ██║   ██║█████╔╝ ██║   ██║  ███╗██║   ██║██████╔╝██║   ██║
██║╚██╔╝██║██║ ███╔╝  ██║   ██║██╔═██╗ ██║   ██║   ██║██║   ██║██╔══██╗██║   ██║
██║ ╚═╝ ██║██║███████╗╚██████╔╝██║  ██╗██║██╗╚██████╔╝╚██████╔╝██║  ██║╚██████╔╝
╚═╝     ╚═╝╚═╝╚══════╝ ╚═════╝ ╚═╝  ╚═╝╚═╝╚═╝ ╚═════╝  ╚═════╝ ╚═╝  ╚═╝ ╚═════╝
"#,
    )?;

    writeln!(stream, "URL: {url}")?;
    writeln!(stream)?;

    let mut response = ureq::get(url).call()?.into_reader();

    let mut file = tempfile::Builder::new().suffix(".img").tempfile()?;

    let mut buf = Vec::new();
    response.read_to_end(&mut buf)?;
    file.write_all(&buf)?;

    write_image(&mut stream, file.path())?;

    loop {
        writeln!(stream)?;
        write!(stream, "> ")?;
        stream.flush()?;

        let Some(prompt) = read_prompt(&mut reader, Duration::from_secs(300))? else {
            break;
        };

        if matches!(prompt.as_str(), "exit" | "quit" | "종료") {
            writeln!(stream, "bye~")?;
            break;
        }

        writeln!(stream)?;
        write!(stream, ": ")?;
        stream.flush()?;

        match generate_gemma_content(config, &mut history, &prompt) {
            Ok(text) => writeln!(stream, "{text}")?,
            Err(e) => writeln!(stream, "Gemma request failed: {e}")?,
        }
    }

    Ok(())
}

fn write_image(stream: &mut TcpStream, image_path: &Path) -> Result<(), Box<dyn std::error::Error>> {
    let output = Command::new("chafa")
        .args([
            "-f",
            "symbols",
            "--colors=256",
            "--size=100x60",
            image_path.to_str().unwrap(),
        ])
        .output();

    match output {
        Ok(output) => {
            stream.write_all(&output.stdout)?;

            if !output.status.success() {
                stream.write_all(&output.stderr)?;
            }
        }
        Err(e) if e.kind() == ErrorKind::NotFound => {
            writeln!(stream, "chafa not found; cannot render image in terminal.")?;
        }
        Err(e) => return Err(e.into()),
    }

    Ok(())
}

fn read_prompt(
    reader: &mut BufReader<TcpStream>,
    timeout: Duration,
) -> Result<Option<String>, Box<dyn std::error::Error>> {
    reader.get_ref().set_read_timeout(Some(timeout))?;

    let mut line = String::new();
    let read_result = reader.read_line(&mut line);

    reader.get_ref().set_read_timeout(None)?;

    match read_result {
        Ok(0) => Ok(None),
        Ok(_) => {
            let prompt = line.trim().to_string();
            Ok((!prompt.is_empty()).then_some(prompt))
        }
        Err(e) if matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => Ok(None),
        Err(e) => Err(e.into()),
    }
}

fn generate_gemma_content(
    config: &Config,
    history: &mut Vec<Value>,
    prompt: &str,
) -> Result<String, Box<dyn std::error::Error>> {
    let mut contents = history.clone();
    let user_content = content("user", prompt);
    contents.push(user_content.clone());

    let request = json!({
        "contents": contents,
        "generationConfig": {
            "thinkingConfig": {
                "thinkingLevel": "MINIMAL"
            }
        },
        "safetySettings": [
            {
                "category": "HARM_CATEGORY_HARASSMENT",
                "threshold": "BLOCK_NONE"
            },
            {
                "category": "HARM_CATEGORY_HATE_SPEECH",
                "threshold": "BLOCK_NONE"
            },
            {
                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                "threshold": "BLOCK_NONE"
            },
            {
                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                "threshold": "BLOCK_NONE"
            }
        ],
        "tools": [
            {
                "googleSearch": {}
            }
        ],
        "systemInstruction": {
            "parts": [
                {
                    "text": "당신은 프로젝트 세카이의 아키야마 미즈키입니다."
                }
            ]
        }
    });

    let url = format!(
        "https://generativelanguage.googleapis.com/v1beta/models/{MODEL_ID}:{GENERATE_CONTENT_API}?key={}",
        config.google.api_key
    );

    let response = match ureq::post(&url)
        .set("Content-Type", "application/json")
        .send_string(&request.to_string())
    {
        Ok(response) => response,
        Err(ureq::Error::Status(code, response)) => {
            let body = response
                .into_string()
                .unwrap_or_else(|_| "failed to read error body".to_string());
            return Err(std::io::Error::new(
                ErrorKind::Other,
                format!("Gemma API returned HTTP {code}: {body}"),
            )
            .into());
        }
        Err(e) => return Err(e.into()),
    };

    let body = response.into_string()?;
    let text = extract_gemma_text(&body).unwrap_or(body);

    history.push(user_content);
    history.push(content("model", &text));

    Ok(text)
}

fn content(role: &str, text: &str) -> Value {
    json!({
        "role": role,
        "parts": [
            {
                "text": text
            }
        ]
    })
}

fn extract_gemma_text(body: &str) -> Option<String> {
    let value: Value = serde_json::from_str(body).ok()?;
    let mut texts = Vec::new();

    match value {
        Value::Array(chunks) => {
            for chunk in &chunks {
                collect_candidate_text(chunk, &mut texts);
            }
        }
        value => collect_candidate_text(&value, &mut texts),
    }

    (!texts.is_empty()).then(|| texts.join(""))
}

fn collect_candidate_text(value: &Value, texts: &mut Vec<String>) {
    let Some(candidates) = value.get("candidates").and_then(Value::as_array) else {
        return;
    };

    for candidate in candidates {
        let Some(parts) = candidate
            .get("content")
            .and_then(|content| content.get("parts"))
            .and_then(Value::as_array)
        else {
            continue;
        };

        for part in parts {
            if let Some(text) = part.get("text").and_then(Value::as_str) {
                texts.push(text.to_string());
            }
        }
    }
}