//! Inference relay (#298): when this runner carries the `inference` label,
//! it polls Anvil for inference requests, runs them against its OWN local
//! model (Ollama), and posts the result back — exactly like the CI job loop,
//! so the Anvil server never makes an outbound model call.
//!
//! The runner does the Ollama wire translation here (mirroring
//! `Anvil.Agents.Providers.Ollama`): Anvil's internal messages/tools →
//! Ollama's `/api/chat` shape, and Ollama's response → the Anvil provider
//! response shape (`content` / `tool_uses` / `stop_reason` / `usage`).

use crate::runner::RunnerConfig;
use reqwest::header::{HeaderValue, AUTHORIZATION, CONTENT_TYPE};
use serde_json::{json, Map, Value};
use std::sync::Arc;
use tokio::sync::Notify;

const MIN_BACKOFF_MS: u64 = 2_000;
const MAX_BACKOFF_MS: u64 = 15_000;

/// Inference polling loop. Mirrors `loop_runner::poller_loop` for the
/// inference-claim endpoint. Runs until `shutdown` is signalled.
pub async fn poller_loop(config: RunnerConfig, shutdown: Arc<Notify>) {
    let client = reqwest::Client::new();
    let mut backoff_ms = MIN_BACKOFF_MS;
    eprintln!(
        "[inference] polling for inference requests (ollama: {})",
        config.ollama_url
    );

    loop {
        let poll = tokio::select! {
            _ = shutdown.notified() => return,
            result = claim(&client, &config) => result,
        };

        match poll {
            Ok(Some(req)) => {
                backoff_ms = MIN_BACKOFF_MS;
                let id = req
                    .get("id")
                    .and_then(|v| v.as_str())
                    .unwrap_or("unknown")
                    .to_string();
                eprintln!("[inference] claimed {}", &id[..8.min(id.len())]);

                let outcome = run_one(&config, &req).await;
                report(&client, &config, &id, outcome).await;
            }
            Ok(None) => {
                tokio::select! {
                    _ = shutdown.notified() => return,
                    _ = tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)) => {},
                }
                backoff_ms = (backoff_ms * 2).min(MAX_BACKOFF_MS);
            }
            Err(e) => {
                eprintln!("[inference] poll error: {e}");
                drop(e);
                tokio::select! {
                    _ = shutdown.notified() => return,
                    _ = tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)) => {},
                }
                backoff_ms = (backoff_ms * 2).min(MAX_BACKOFF_MS);
            }
        }
    }
}

async fn claim(
    client: &reqwest::Client,
    config: &RunnerConfig,
) -> Result<Option<Value>, Box<dyn std::error::Error + Send + Sync>> {
    let url = config.api_url(&format!("/runners/{}/inference/claim", config.runner_id));
    let resp = client
        .post(&url)
        .header(AUTHORIZATION, HeaderValue::from_str(&config.auth_header())?)
        .send()
        .await?;

    match resp.status().as_u16() {
        200 => {
            let body: Value = resp.json().await?;
            Ok(body.get("request").cloned())
        }
        204 => Ok(None),
        status => {
            let text = resp.text().await.unwrap_or_default();
            Err(format!("inference claim failed ({status}): {text}").into())
        }
    }
}

/// Run one request against the local model. Returns the Anvil-shaped
/// response on success, or an error string to report as a failure.
async fn run_one(config: &RunnerConfig, req: &Value) -> Result<Value, String> {
    match req.get("provider").and_then(|v| v.as_str()) {
        Some("ollama") => call_ollama(config, req).await,
        Some(other) => Err(format!("runner cannot serve provider '{other}'")),
        None => Err("inference request missing 'provider'".to_string()),
    }
}

async fn call_ollama(config: &RunnerConfig, req: &Value) -> Result<Value, String> {
    let model = req.get("model").and_then(|v| v.as_str()).unwrap_or("");
    let options = req.get("options").cloned().unwrap_or(json!({}));

    let mut body = Map::new();
    body.insert("model".into(), json!(model));
    body.insert("messages".into(), json!(ollama_messages(req, &options)));
    body.insert("stream".into(), json!(false));

    let tools = ollama_tools(req);
    if !tools.is_empty() {
        body.insert("tools".into(), json!(tools));
    }
    if let Some(opts) = ollama_options(&options) {
        body.insert("options".into(), opts);
    }

    let url = format!("{}/api/chat", config.ollama_url.trim_end_matches('/'));
    let client = reqwest::Client::new();
    let resp = client
        .post(&url)
        .header(CONTENT_TYPE, HeaderValue::from_static("application/json"))
        .json(&Value::Object(body))
        .send()
        .await
        .map_err(|e| format!("ollama request failed: {e}"))?;

    let status = resp.status();
    let raw: Value = resp
        .json()
        .await
        .map_err(|e| format!("ollama returned non-JSON: {e}"))?;

    if !status.is_success() {
        let msg = raw
            .get("error")
            .and_then(|v| v.as_str())
            .unwrap_or("unknown error");
        return Err(format!("ollama error ({status}): {msg}"));
    }

    Ok(parse_ollama_response(&raw))
}

// ── Anvil internal → Ollama wire ───────────────────────────────────────

// Build Ollama's `messages`, prefixing a system message from options.system.
fn ollama_messages(req: &Value, options: &Value) -> Vec<Value> {
    let mut out = Vec::new();

    if let Some(system) = options.get("system").and_then(|v| v.as_str()) {
        out.push(json!({"role": "system", "content": system}));
    }

    if let Some(messages) = req.get("messages").and_then(|v| v.as_array()) {
        for msg in messages {
            out.extend(translate_message(msg));
        }
    }

    out
}

// A message's content is either a plain string (initial turn) or a list of
// typed blocks (after a tool round-trip). Mirrors the Elixir provider's
// format_message/1.
fn translate_message(msg: &Value) -> Vec<Value> {
    let role = msg.get("role").and_then(|v| v.as_str()).unwrap_or("user");

    match msg.get("content") {
        Some(Value::String(text)) => vec![json!({"role": role, "content": text})],
        Some(Value::Array(blocks)) => translate_blocks(role, blocks),
        _ => vec![json!({"role": role, "content": ""})],
    }
}

fn translate_blocks(role: &str, blocks: &[Value]) -> Vec<Value> {
    if role == "user" {
        // tool_result blocks become separate {role: "tool"} messages; text
        // blocks become a user message.
        blocks
            .iter()
            .map(|b| match block_type(b) {
                "tool_result" => json!({
                    "role": "tool",
                    "content": stringify(b.get("content")),
                    "tool_call_id": b.get("tool_use_id").and_then(|v| v.as_str()).unwrap_or("")
                }),
                _ => json!({"role": "user", "content": b.get("text").and_then(|v| v.as_str()).unwrap_or("")}),
            })
            .collect()
    } else {
        // assistant: collect text + tool_use blocks into one message.
        let text: String = blocks
            .iter()
            .filter(|b| block_type(b) == "text")
            .filter_map(|b| b.get("text").and_then(|v| v.as_str()))
            .collect();

        let tool_calls: Vec<Value> = blocks
            .iter()
            .filter(|b| block_type(b) == "tool_use")
            .map(|b| {
                json!({
                    "id": b.get("id").and_then(|v| v.as_str()).unwrap_or(""),
                    "type": "function",
                    "function": {
                        "name": b.get("name").and_then(|v| v.as_str()).unwrap_or(""),
                        "arguments": b.get("input").cloned().unwrap_or(json!({}))
                    }
                })
            })
            .collect();

        let mut msg = Map::new();
        msg.insert("role".into(), json!(role));
        msg.insert("content".into(), json!(text));
        if !tool_calls.is_empty() {
            msg.insert("tool_calls".into(), json!(tool_calls));
        }
        vec![Value::Object(msg)]
    }
}

fn block_type(b: &Value) -> &str {
    b.get("type").and_then(|v| v.as_str()).unwrap_or("")
}

fn stringify(v: Option<&Value>) -> String {
    match v {
        Some(Value::String(s)) => s.clone(),
        Some(other) => other.to_string(),
        None => String::new(),
    }
}

// Anvil tool defs ({name, description, input_schema}) → Ollama function tools.
fn ollama_tools(req: &Value) -> Vec<Value> {
    req.get("tools")
        .and_then(|v| v.as_array())
        .map(|tools| {
            tools
                .iter()
                .map(|t| {
                    json!({
                        "type": "function",
                        "function": {
                            "name": t.get("name").and_then(|v| v.as_str()).unwrap_or(""),
                            "description": t.get("description").and_then(|v| v.as_str()).unwrap_or(""),
                            "parameters": t.get("input_schema").cloned().unwrap_or(json!({}))
                        }
                    })
                })
                .collect()
        })
        .unwrap_or_default()
}

fn ollama_options(options: &Value) -> Option<Value> {
    let mut opts = Map::new();
    if let Some(t) = options.get("temperature") {
        if !t.is_null() {
            opts.insert("temperature".into(), t.clone());
        }
    }
    if let Some(mt) = options.get("max_tokens") {
        if !mt.is_null() {
            opts.insert("num_predict".into(), mt.clone());
        }
    }
    if opts.is_empty() {
        None
    } else {
        Some(Value::Object(opts))
    }
}

// ── Ollama response → Anvil provider response ──────────────────────────

fn parse_ollama_response(raw: &Value) -> Value {
    let message = raw.get("message").cloned().unwrap_or(json!({}));
    let tool_uses = parse_tool_calls(message.get("tool_calls"));
    let has_tools = !tool_uses.is_empty();

    json!({
        "content": message.get("content").and_then(|v| v.as_str()).unwrap_or(""),
        "tool_uses": tool_uses,
        "stop_reason": stop_reason(raw.get("done_reason").and_then(|v| v.as_str()), has_tools),
        "usage": {
            "input_tokens": raw.get("prompt_eval_count").and_then(|v| v.as_u64()).unwrap_or(0),
            "output_tokens": raw.get("eval_count").and_then(|v| v.as_u64()).unwrap_or(0)
        }
    })
}

fn parse_tool_calls(calls: Option<&Value>) -> Vec<Value> {
    calls
        .and_then(|v| v.as_array())
        .map(|arr| {
            arr.iter()
                .enumerate()
                .map(|(i, call)| {
                    let f = call.get("function").cloned().unwrap_or(json!({}));
                    json!({
                        "id": call.get("id").and_then(|v| v.as_str()).map(String::from)
                            .unwrap_or_else(|| format!("ollama-{i}")),
                        "name": f.get("name").and_then(|v| v.as_str()).unwrap_or(""),
                        "input": normalize_arguments(f.get("arguments"))
                    })
                })
                .collect()
        })
        .unwrap_or_default()
}

// Some models return arguments as a JSON string rather than an object.
fn normalize_arguments(args: Option<&Value>) -> Value {
    match args {
        Some(Value::Object(_)) => args.cloned().unwrap(),
        Some(Value::String(s)) => serde_json::from_str(s).unwrap_or(json!({"_raw": s})),
        _ => json!({}),
    }
}

fn stop_reason(done_reason: Option<&str>, has_tools: bool) -> &'static str {
    if has_tools {
        "tool_use"
    } else {
        match done_reason {
            Some("length") => "max_tokens",
            _ => "end_turn",
        }
    }
}

// ── Report result back to Anvil ────────────────────────────────────────

async fn report(
    client: &reqwest::Client,
    config: &RunnerConfig,
    request_id: &str,
    outcome: Result<Value, String>,
) {
    let url = config.api_url(&format!(
        "/runners/{}/inference/{request_id}/result",
        config.runner_id
    ));

    let body = match &outcome {
        Ok(response) => json!({"response": response}),
        Err(error) => json!({"error": error}),
    };

    match &outcome {
        Ok(_) => eprintln!(
            "[inference] {} completed",
            &request_id[..8.min(request_id.len())]
        ),
        Err(e) => eprintln!(
            "[inference] {} failed: {e}",
            &request_id[..8.min(request_id.len())]
        ),
    }

    let result = client
        .post(&url)
        .header(
            AUTHORIZATION,
            HeaderValue::from_str(&config.auth_header()).unwrap(),
        )
        .header(CONTENT_TYPE, HeaderValue::from_static("application/json"))
        .json(&body)
        .send()
        .await;

    if let Err(e) = result {
        eprintln!("[inference] failed to report result: {e}");
    }
}