lmocr 0.1.1

use anyhow::{Context, Result};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use reqwest::header::{AUTHORIZATION, HeaderMap, HeaderValue};
use serde::{Deserialize, Serialize};

pub const DEFAULT_MODEL: &str = "google/gemini-3-flash-preview";

#[derive(Debug, Clone)]
pub struct Usage {
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub total_tokens: u64,
}

pub struct OcrResult {
    pub content: String,
    pub usage: Usage,
}

pub struct OpenRouterClient {
    client: reqwest::Client,
    api_key: String,
}

impl OpenRouterClient {
    pub fn new(api_key: String) -> Result<Self> {
        let client = reqwest::Client::builder()
            .user_agent("lmocr/0.1")
            .timeout(std::time::Duration::from_secs(600))
            .build()?;
        Ok(Self { client, api_key })
    }

    pub async fn ocr_page(
        &self,
        model: &str,
        system_prompt: &str,
        user_text: &str,
        image_bytes: Option<&[u8]>,
        mime_type: &str,
    ) -> Result<OcrResult> {
        let image_data = image_bytes
            .map(|bytes| format!("data:{};base64,{}", mime_type, STANDARD.encode(bytes)));
        let mut user_content = vec![Content::Text {
            text: user_text.to_string(),
        }];
        if let Some(data) = image_data {
            user_content.push(Content::ImageUrl {
                image_url: ImageUrl { url: data },
            });
        }
        let request = ChatRequest {
            model: model.to_string(),
            temperature: 0.7,
            messages: vec![
                ChatMessage {
                    role: "system".to_string(),
                    content: vec![Content::Text {
                        text: system_prompt.to_string(),
                    }],
                },
                ChatMessage {
                    role: "user".to_string(),
                    content: user_content,
                },
            ],
        };

        let mut headers = HeaderMap::new();
        headers.insert(
            AUTHORIZATION,
            HeaderValue::from_str(&format!("Bearer {}", self.api_key))?,
        );
        headers.insert(
            "HTTP-Referer",
            HeaderValue::from_static("https://lmocr.local"),
        );
        headers.insert("X-Title", HeaderValue::from_static("lmocr"));

        let response = self
            .client
            .post("https://openrouter.ai/api/v1/chat/completions")
            .headers(headers)
            .json(&request)
            .send()
            .await?;

        let status = response.status();
        let body = response.bytes().await?;
        if !status.is_success() {
            let text = String::from_utf8_lossy(&body);
            anyhow::bail!("OpenRouter error {}: {}", status, text);
        }

        let parsed: ChatResponse = serde_json::from_slice(&body).with_context(|| {
            format!(
                "Failed to parse OpenRouter response: {}",
                String::from_utf8_lossy(&body)
            )
        })?;
        let choice = parsed
            .choices
            .into_iter()
            .next()
            .context("OpenRouter response missing choices")?;
        let usage = parsed.usage;
        Ok(OcrResult {
            content: choice.message.content,
            usage: Usage {
                prompt_tokens: usage.prompt_tokens,
                completion_tokens: usage.completion_tokens,
                total_tokens: usage.total_tokens,
            },
        })
    }
}

#[derive(Serialize)]
struct ChatRequest {
    model: String,
    temperature: f32,
    messages: Vec<ChatMessage>,
}

#[derive(Serialize)]
struct ChatMessage {
    role: String,
    content: Vec<Content>,
}

#[derive(Serialize)]
#[serde(tag = "type")]
enum Content {
    #[serde(rename = "text")]
    Text { text: String },
    #[serde(rename = "image_url")]
    ImageUrl { image_url: ImageUrl },
}

#[derive(Serialize)]
struct ImageUrl {
    url: String,
}

#[derive(Deserialize)]
struct ChatResponse {
    choices: Vec<Choice>,
    #[serde(default)]
    usage: UsageResponse,
}

#[derive(Deserialize)]
struct Choice {
    message: AssistantMessage,
}

#[derive(Deserialize)]
struct AssistantMessage {
    content: String,
}

#[derive(Deserialize, Default)]
struct UsageResponse {
    #[serde(default)]
    prompt_tokens: u64,
    #[serde(default)]
    completion_tokens: u64,
    #[serde(default)]
    total_tokens: u64,
}