//! 3D models with fal.ai (generate_3d): from an image (plus optional back/left/right
//! views), from a text prompt, or from an image steered by text – the 3D models take no
//! text with an image, so the image is first refined with the prompt (image edit model)
//! and then converted. The result is a GLB file (plus a preview image) in the project.

use crate::app::App;
use crate::config::expand_env;
use crate::store::UsageRecord;
use crate::tools::{Tool, ToolCtx, ToolOutput};
use crate::types::now_ms;
use anyhow::{Context, Result, anyhow, bail};
use async_trait::async_trait;
use serde_json::{Value, json};
use std::path::{Path, PathBuf};
use std::time::Duration;

const STYLES: &[&str] = &["textured", "geometry", "lowpoly"];

/// Request body for a fal.ai 3D endpoint (the families name their fields differently).
fn request_body(model: &str, prompt: Option<&str>, images: &Views, style: &str, pbr: bool, faces: Option<u64>) -> Value {
    let mut b = json!({});
    if model.contains("tripo") {
        // Tripo: image_url / prompt, texture + pbr flags, face_limit.
        if let Some(front) = &images.front {
            b["image_url"] = json!(front);
        }
        if let Some(p) = prompt {
            b["prompt"] = json!(p);
        }
        b["texture"] = json!(style != "geometry");
        b["pbr"] = json!(style != "geometry" && pbr);
        if let Some(f) = faces {
            b["face_limit"] = json!(f);
        }
        if style == "lowpoly" {
            b["face_limit"] = json!(faces.unwrap_or(20_000));
        }
    } else if model.contains("trellis") {
        if let Some(front) = &images.front {
            b["image_url"] = json!(front);
        }
    } else {
        // Hunyuan3D: front + optional views, generate_type, enable_pbr, face_count.
        if let Some(front) = &images.front {
            b["input_image_url"] = json!(front);
        }
        for (k, v) in [("back_image_url", &images.back), ("left_image_url", &images.left), ("right_image_url", &images.right)] {
            if let Some(v) = v {
                b[k] = json!(v);
            }
        }
        if let Some(p) = prompt {
            b["prompt"] = json!(p);
        }
        b["generate_type"] = json!(match style {
            "geometry" => "Geometry",
            "lowpoly" => "LowPoly",
            _ => "Normal",
        });
        b["enable_pbr"] = json!(style != "geometry" && pbr);
        if let Some(f) = faces {
            b["face_count"] = json!(f.clamp(40_000, 1_500_000));
        }
    }
    b
}

/// Data URIs of the views (front required for image-to-3D).
#[derive(Default)]
struct Views {
    front: Option<String>,
    back: Option<String>,
    left: Option<String>,
    right: Option<String>,
}

impl Views {
    fn extra(&self) -> usize {
        [&self.back, &self.left, &self.right].iter().filter(|v| v.is_some()).count()
    }
}

/// Price of one generation (USD) by the published fal.ai prices; a price configured for
/// the model under [images] models wins.
fn price(app: &App, model: &str, style: &str, pbr: bool, multiview: bool, faces: bool) -> f64 {
    if let Some(p) = app.config().images.models.iter().find(|m| m.id == model).and_then(|m| m.price) {
        return p;
    }
    if model.contains("tripo") {
        return match style {
            "geometry" => 0.20,
            _ => 0.30,
        };
    }
    if model.contains("hunyuan") {
        let base = match style {
            "geometry" => 0.225,
            "lowpoly" => 0.45,
            _ => 0.375,
        };
        let extras = [pbr && style != "geometry", multiview, faces].iter().filter(|x| **x).count() as f64;
        return base + 0.15 * extras;
    }
    0.0
}

/// URL of the queue API for fal.ai (requests that take minutes).
fn queue_base(app: &App) -> String {
    let base = expand_env(&app.config().images.base_url);
    if base.contains("://fal.run") { "https://queue.fal.run".into() } else { base.trim_end_matches('/').to_string() }
}

/// Submit to fal.ai's queue and wait for the result, reporting progress.
async fn run_queued(app: &App, ctx: &ToolCtx, model: &str, body: &Value) -> Result<Value> {
    let key = expand_env(&app.config().images.api_key);
    let client = crate::provider::http_client();
    let submit = client
        .post(format!("{}/{}", queue_base(app), model.trim_matches('/')))
        .header("Authorization", format!("Key {key}"))
        .timeout(Duration::from_secs(120))
        .json(body)
        .send()
        .await
        .context("fal.ai request failed")?;
    if !submit.status().is_success() {
        return Err(crate::provider::http_error(submit).await);
    }
    let sub: Value = submit.json().await?;
    let status_url = sub["status_url"].as_str().context("fal.ai: no status_url")?.to_string();
    let response_url = sub["response_url"].as_str().context("fal.ai: no response_url")?.to_string();
    let started = std::time::Instant::now();
    loop {
        if ctx.cancel.is_cancelled() {
            bail!("cancelled");
        }
        if started.elapsed() > Duration::from_secs(20 * 60) {
            bail!("fal.ai did not finish within 20 minutes");
        }
        tokio::time::sleep(Duration::from_secs(4)).await;
        let st: Value = client.get(&status_url).header("Authorization", format!("Key {key}")).timeout(Duration::from_secs(30)).send().await?.json().await?;
        match st["status"].as_str().unwrap_or("") {
            "COMPLETED" => break,
            "IN_QUEUE" => ctx.progress(format!("waiting in fal.ai's queue (position {}) – {} s", st["queue_position"].as_u64().unwrap_or(0), started.elapsed().as_secs())),
            "IN_PROGRESS" => ctx.progress(format!("generating the 3D model – {} s", started.elapsed().as_secs())),
            other => bail!("fal.ai: unexpected status {other}: {st}"),
        }
    }
    let resp = client.get(&response_url).header("Authorization", format!("Key {key}")).timeout(Duration::from_secs(60)).send().await?;
    if !resp.status().is_success() {
        return Err(crate::provider::http_error(resp).await);
    }
    Ok(resp.json().await?)
}

/// URL of the mesh and of the preview image in a fal.ai 3D result.
fn result_urls(out: &Value) -> (Option<String>, Option<String>) {
    let s = |v: &Value| v.as_str().map(String::from);
    let mesh = s(&out["model_glb"]["url"])
        .or_else(|| s(&out["model_mesh"]["url"]))
        .or_else(|| s(&out["model_urls"]["glb"]["url"]))
        .or_else(|| s(&out["model_file"]["url"]))
        .or_else(|| s(&out["mesh"]["url"]));
    let thumb = s(&out["thumbnail"]["url"]).or_else(|| s(&out["rendered_image"]["url"]));
    (mesh, thumb)
}

/// File extension of a downloaded mesh: from the URL, else GLB (what the models return).
fn ext_of(url: &str) -> &'static str {
    let path = url.split(['?', '#']).next().unwrap_or(url).to_ascii_lowercase();
    ["glb", "gltf", "fbx", "obj", "usdz", "stl"].into_iter().find(|e| path.ends_with(&format!(".{e}"))).unwrap_or("glb")
}

/// An image the user attached to the chat ("attached" = the first image of the latest
/// message with images, "attached:2" = its second), written to `dir` so it can be used
/// like a project file.
async fn attached_image(app: &App, session_id: &str, spec: &str, dir: &Path) -> Result<PathBuf> {
    let n: usize = spec.strip_prefix("attached").map(|r| r.trim_start_matches(':').trim()).filter(|r| !r.is_empty()).map(|r| r.parse()).transpose().map_err(|_| anyhow!("use \"attached\" or \"attached:2\""))?.unwrap_or(1);
    let root = app.root_session_id(session_id);
    let messages = app.store.messages(&root)?;
    let images: Vec<crate::types::ImageData> = messages
        .iter()
        .rev()
        .filter(|m| m.role == crate::types::Role::User)
        .map(|m| m.parts.iter().filter_map(|p| if let crate::types::Part::Image { image } = p { Some(image.clone()) } else { None }).collect::<Vec<_>>())
        .find(|imgs| !imgs.is_empty())
        .ok_or_else(|| anyhow!("no image attached to this chat"))?;
    let img = images.get(n.saturating_sub(1)).ok_or_else(|| anyhow!("the latest message has only {} image(s)", images.len()))?;
    let (mime, bytes) = match &img.blob {
        Some(h) => app.store.blob(h)?.ok_or_else(|| anyhow!("attached image not found"))?,
        None => (img.media_type.clone(), base64::Engine::decode(&base64::engine::general_purpose::STANDARD, img.data.as_bytes())?),
    };
    let ext = match mime.as_str() {
        "image/jpeg" => "jpg",
        "image/webp" => "webp",
        _ => "png",
    };
    tokio::fs::create_dir_all(dir).await?;
    let path = dir.join(format!("{}-attached-{n}.{ext}", chrono::Local::now().format("%Y%m%d-%H%M%S")));
    tokio::fs::write(&path, &bytes).await?;
    Ok(path)
}

/// An image argument: a project file, or an image attached to the chat.
async fn image_arg(ctx: &ToolCtx, spec: &str, dir: &Path) -> Result<PathBuf> {
    if spec.starts_with("attached") {
        attached_image(&ctx.app, &ctx.session_id, spec, dir).await
    } else {
        ctx.resolve(spec)
    }
}

pub struct Generate3d;

#[async_trait]
impl Tool for Generate3d {
    fn name(&self) -> &str {
        "generate_3d"
    }
    fn description(&self) -> String {
        "Generate a 3D model (GLB) with fal.ai and save it in the project – from an image of the \
         object (a file, or image: \"attached\" for an image the user pasted into the chat; best: a \
         clean front view; optional back/left/right views improve the shape), from a \
         text prompt alone, or from an image steered by a prompt (the image is refined with the \
         prompt first, then converted). style: \"textured\" (default), \"geometry\" (untextured – for \
         3D printing), \"lowpoly\" (games). Takes one to several minutes and costs money (about \
         $0.25–0.75): generate only what the user asked for. The user can open the model in the 3D \
         view of the file manager."
            .into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "prompt": {"type": "string", "description": "What the model should be (English). With an image: how to change/clean it up before the conversion."},
                "image": {"type": "string", "description": "Image of the object (front view): a file path, or \"attached\" for the image the user attached to the chat (\"attached:2\" = the second of that message)"},
                "back": {"type": "string", "description": "Optional back view image (path or \"attached:N\")"},
                "left": {"type": "string", "description": "Optional left view image (path or \"attached:N\")"},
                "right": {"type": "string", "description": "Optional right view image (path or \"attached:N\")"},
                "style": {"type": "string", "enum": STYLES},
                "pbr": {"type": "boolean", "description": "PBR materials (metal/roughness), textured models only; costs extra"},
                "face_count": {"type": "integer", "description": "Target number of faces (optional)"},
                "directory": {"type": "string", "description": "Target directory relative to the working directory (default: <images.output_dir>/3d)"},
                "model": {"type": "string", "description": "fal.ai model id – only if the user asks for a specific one"}
            }
        })
    }
    fn title(&self, input: &Value) -> String {
        let what = input["prompt"].as_str().map(|p| p.chars().take(70).collect::<String>()).or_else(|| input["image"].as_str().map(String::from)).unwrap_or_default();
        format!("3D model: {what}")
    }
    fn read_only(&self) -> bool {
        false
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let app = ctx.app.clone();
        let cfg = app.config();
        if !cfg.images.enabled {
            bail!("fal.ai is disabled – enable image generation in the settings");
        }
        if expand_env(&cfg.images.api_key).is_empty() {
            bail!("no fal.ai API key configured");
        }
        let owner = app.session(&ctx.session_id).ok().and_then(|s| s.owner_id);
        crate::budget::check(&app, "fal", owner.as_deref())?;
        let prompt = input["prompt"].as_str().map(str::trim).filter(|p| !p.is_empty()).map(String::from);
        let dir_arg = input["directory"].as_str().map(String::from).unwrap_or_else(|| format!("{}/3d", cfg.images.output_dir.trim_end_matches('/')));
        let dir = ctx.resolve(&dir_arg)?;
        let image = match input["image"].as_str() {
            Some(spec) => Some(image_arg(ctx, spec, &dir).await?),
            None => None,
        };
        if prompt.is_none() && image.is_none() {
            bail!("give an image of the object, a prompt, or both");
        }
        let style = input["style"].as_str().filter(|s| STYLES.contains(s)).unwrap_or("textured");
        let pbr = input["pbr"].as_bool().unwrap_or(false);
        let faces = input["face_count"].as_u64();
        let mut notes = Vec::new();

        // Image + prompt: refine the image with the prompt first (3D models take no text).
        let mut front = image.clone();
        if let (Some(img), Some(p)) = (&image, &prompt) {
            ctx.progress("refining the image with the prompt");
            let edit = crate::images::choose_model(&app, "final", Some("illustration"), true);
            if edit.reference_param.is_some() {
                let records = crate::images::generate(
                    &app,
                    crate::images::GenerateRequest {
                        prompt: format!(
                            "{p}. Show exactly one object, complete and centered, as a clean product shot: plain light neutral background, soft even lighting, no text, no other objects."
                        ),
                        model: Some(edit.id.clone()),
                        size: None,
                        count: 1,
                        dir: dir.clone(),
                        user_id: owner.clone(),
                        session_id: Some(ctx.session_id.clone()),
                        references: vec![img.clone()],
                    },
                )
                .await
                .map_err(|e| anyhow!("refining the image failed: {e:#}"))?;
                if let Some(r) = records.first() {
                    front = Some(PathBuf::from(&r.path));
                    notes.push(format!("Refined image: {}", rel(&r.path, &ctx.cwd)));
                }
            } else {
                notes.push("No image edit model configured – the prompt was not applied to the image.".into());
            }
        }

        let mut views = Views::default();
        if let Some(f) = &front {
            views.front = crate::images::reference_uris(std::slice::from_ref(f)).await?.pop();
            for (field, slot) in [("back", &mut views.back), ("left", &mut views.left), ("right", &mut views.right)] {
                if let Some(p) = input[field].as_str() {
                    *slot = crate::images::reference_uris(&[image_arg(ctx, p, &dir).await?]).await?.pop();
                }
            }
        }
        let model = input["model"].as_str().filter(|m| !m.is_empty()).map(String::from).unwrap_or_else(|| if front.is_some() { cfg.images.model_3d.clone() } else { cfg.images.model_3d_text.clone() });
        if views.extra() > 0 && !model.contains("hunyuan") {
            notes.push("This model takes only the front view; the other views were not used.".into());
        }
        let text = if front.is_none() { prompt.as_deref() } else { None };
        let body = request_body(&model, text, &views, style, pbr, faces);
        ctx.progress(format!("sending to {model}"));
        let out = run_queued(&app, ctx, &model, &body).await?;
        let (mesh_url, thumb_url) = result_urls(&out);
        let mesh_url = mesh_url.ok_or_else(|| anyhow!("fal.ai returned no 3D model: {}", out.to_string().chars().take(300).collect::<String>()))?;

        tokio::fs::create_dir_all(&dir).await?;
        let stamp = chrono::Local::now().format("%Y%m%d-%H%M%S");
        // Named after the prompt, else the image (without a time stamp it already carries).
        let stem = image.as_ref().and_then(|p| p.file_stem()).and_then(|s| s.to_str()).unwrap_or("model");
        let stem = if stem.len() > 16 && stem.as_bytes()[8] == b'-' && stem[..8].bytes().all(|b| b.is_ascii_digit()) { &stem[16..] } else { stem };
        let base = crate::images::slug(prompt.as_deref().unwrap_or(stem).trim_start_matches('-'));
        let (bytes, _) = crate::images::download(&mesh_url).await?;
        let mesh_path = dir.join(format!("{stamp}-{base}.{}", ext_of(&mesh_url)));
        tokio::fs::write(&mesh_path, &bytes).await?;
        let mut thumb_path = None;
        if let Some(t) = thumb_url {
            if let Ok((tb, _)) = crate::images::download(&t).await {
                let p = dir.join(format!("{stamp}-{base}-preview.png"));
                if tokio::fs::write(&p, &tb).await.is_ok() {
                    thumb_path = Some(p);
                }
            }
        }

        // Cost into the statistics (provider "fal", like images).
        let cost = price(&app, &model, style, pbr, views.extra() > 0 && model.contains("hunyuan"), faces.is_some());
        let session = ctx.session_id.clone();
        let _ = app.store.add_usage(&UsageRecord {
            ts: now_ms(),
            session_id: session.clone(),
            root_session_id: app.root_session_id(&session),
            provider: "fal".into(),
            model: model.clone(),
            agent: "3d".into(),
            input_tokens: 0,
            output_tokens: 0,
            cache_read_tokens: 0,
            cache_write_tokens: 0,
            ttft_ms: 0,
            duration_ms: 0,
            cost,
            error: false,
            user_id: owner,
            allowance: None,
        });

        let mesh_rel = rel(&mesh_path.display().to_string(), &ctx.cwd);
        let mut text = format!("3D model ({model}, {style}{}): {mesh_rel} ({} KB)", if pbr && style != "geometry" { ", PBR" } else { "" }, bytes.len() / 1024);
        if let Some(t) = &thumb_path {
            text.push_str(&format!("\nPreview: {}", rel(&t.display().to_string(), &ctx.cwd)));
        }
        if cost > 0.0 {
            text.push_str(&format!("\nCost: about ${cost:.2}"));
        }
        for n in &notes {
            text.push_str(&format!("\n{n}"));
        }
        text.push_str("\nThe user can open the model in the 3D view of the file manager.");
        let images = thumb_path.as_deref().and_then(crate::images::preview_for_model).into_iter().collect();
        Ok(ToolOutput {
            content: text,
            images,
            meta: Some(json!({
                "model_file": mesh_path.display().to_string(),
                "preview": thumb_path.map(|p| p.display().to_string()),
            })),
            ..Default::default()
        })
    }
}

fn rel(path: &str, cwd: &Path) -> String {
    Path::new(path).strip_prefix(cwd).map(|p| p.display().to_string()).unwrap_or_else(|_| path.to_string())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn bodies_per_model_family() {
        let views = Views { front: Some("data:a".into()), back: Some("data:b".into()), ..Default::default() };
        let h = request_body("fal-ai/hunyuan3d-v3/image-to-3d", None, &views, "geometry", true, Some(10));
        assert_eq!(h["input_image_url"], "data:a");
        assert_eq!(h["back_image_url"], "data:b");
        assert_eq!(h["generate_type"], "Geometry");
        assert_eq!(h["enable_pbr"], false);
        assert_eq!(h["face_count"], 40_000);
        let t = request_body("tripo3d/h3.1/image-to-3d", None, &views, "textured", true, None);
        assert_eq!((t["image_url"].as_str(), t["texture"].as_bool(), t["pbr"].as_bool()), (Some("data:a"), Some(true), Some(true)));
        let txt = request_body("fal-ai/hunyuan3d-v3/text-to-3d", Some("a chair"), &Views::default(), "textured", false, None);
        assert_eq!((txt["prompt"].as_str(), txt["generate_type"].as_str()), (Some("a chair"), Some("Normal")));
    }

    #[test]
    fn result_urls_of_the_families() {
        let h = json!({"model_glb": {"url": "https://x/m.glb"}, "thumbnail": {"url": "https://x/t.png"}});
        assert_eq!(result_urls(&h), (Some("https://x/m.glb".into()), Some("https://x/t.png".into())));
        let t = json!({"model_mesh": {"url": "https://x/m.glb"}, "rendered_image": {"url": "https://x/r.webp"}});
        assert_eq!(result_urls(&t).1.as_deref(), Some("https://x/r.webp"));
        assert_eq!(ext_of("https://x/a.FBX?sig=1"), "fbx");
        assert_eq!(ext_of("https://x/blob"), "glb");
    }
}