//! LLM brain: async worker pool talking to Ollama, wake scheduling, action
//! application. The render loop never blocks on a model call.

use bevy::prelude::*;
use serde::Deserialize;
use serde_json::json;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use crate::agent::{cell_of, world_of, Brain, Memory, Names, Needs, Pathing, Speech, Villager};
use crate::persona::MODEL;
use crate::world::WorldMap;

const WORKERS: usize = 6; // in-flight cap
const WAKE_MIN: f32 = 45.0; // seconds between a villager's scheduled decisions
const WAKE_JITTER: f32 = 75.0;
const THINK_BUDGET: u32 = 64; // max new tokens per decision

pub struct DecisionJob {
    pub entity: Entity,
    pub name: String,
    pub system: String,
    pub user: String,
}

pub struct DecisionResult {
    pub entity: Entity,
    pub name: String,
    pub action: Result<Action, String>,
    pub elapsed: f32,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(tag = "action", rename_all = "snake_case")]
pub enum Action {
    GoTo { place: String, reason: String },
    Talk { to: String, say: String, reason: String },
    Work { reason: String },
    Rest { reason: String },
    Idle { reason: String },
}

/// `Receiver` is `Send` but not `Sync`, so it rides behind a mutex to be a
/// valid Bevy `Resource`. Workers share the job receiver the same way.
#[derive(Resource)]
pub struct LlmPool {
    pub jobs_tx: Sender<DecisionJob>,
    pub results_rx: Mutex<Receiver<DecisionResult>>,
    pub in_flight: usize,
    pub completed: u64,
    pub failed: u64,
    pub peak_in_flight: usize,
}

pub fn start_pool(mut commands: Commands) {
    let (jobs_tx, jobs_rx) = channel::<DecisionJob>();
    let (results_tx, results_rx) = channel::<DecisionResult>();
    let url = std::env::var("OLLAMA_URL").unwrap_or_else(|_| "http://localhost:11434".into());
    let jobs_rx = Arc::new(Mutex::new(jobs_rx));

    for _ in 0..WORKERS {
        let jobs_rx = Arc::clone(&jobs_rx);
        let results_tx = results_tx.clone();
        let url = url.clone();
        std::thread::spawn(move || worker_loop(jobs_rx, results_tx, &url));
    }
    commands.insert_resource(LlmPool {
        jobs_tx,
        results_rx: Mutex::new(results_rx),
        in_flight: 0,
        completed: 0,
        failed: 0,
        peak_in_flight: 0,
    });
}

fn worker_loop(
    jobs_rx: Arc<Mutex<Receiver<DecisionJob>>>,
    results_tx: Sender<DecisionResult>,
    url: &str,
) {
    let client = reqwest::blocking::Client::builder()
        .timeout(Duration::from_secs(120))
        .build()
        .expect("http client");
    loop {
        // Blocking recv while holding the lock is fine: workers take turns
        // waiting, and the handoff is instant when the queue is non-empty.
        let job = {
            let rx = jobs_rx.lock().unwrap_or_else(|e| e.into_inner());
            rx.recv()
        };
        let Ok(job) = job else { break };
        let started = Instant::now();
        let action = request_action(&client, url, &job.system, &job.user);
        let _ = results_tx.send(DecisionResult {
            entity: job.entity,
            name: job.name,
            action,
            elapsed: started.elapsed().as_secs_f32(),
        });
    }
}

fn request_action(
    client: &reqwest::blocking::Client,
    url: &str,
    system: &str,
    user: &str,
) -> Result<Action, String> {
    let schema = json!({
        "type": "object",
        "properties": {
            "action": {"type": "string",
                       "enum": ["go_to", "talk", "work", "rest", "idle"]},
            "place": {"type": "string"},
            "to": {"type": "string"},
            "say": {"type": "string"},
            "reason": {"type": "string"}
        },
        "required": ["action", "reason"]
    });

    let body = json!({
        "model": MODEL,
        "stream": false,
        "think": false,
        "format": schema,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "options": {
            "temperature": 0.9,
            "num_predict": THINK_BUDGET,
            "num_ctx": 4096,
            "keep_alive": "30m",
        }
    });

    let resp = client
        .post(format!("{url}/api/chat"))
        .json(&body)
        .send()
        .map_err(|e| format!("http: {e}"))?;
    if !resp.status().is_success() {
        return Err(format!("status {}", resp.status()));
    }
    let value: serde_json::Value = resp.json().map_err(|e| format!("json: {e}"))?;
    let content = value["message"]["content"].as_str().unwrap_or_default();
    let content = content
        .trim()
        .trim_start_matches("```json")
        .trim_start_matches("```")
        .trim_end_matches("```")
        .trim();
    serde_json::from_str::<Action>(content).map_err(|e| format!("{e}: {content:.120}"))
}

// ---------------------------------------------------------------------------
// Perception helpers
// ---------------------------------------------------------------------------

/// Rough named place for a position (nearest named POI within 6 units).
fn place_at(map: &WorldMap, pos: Vec3) -> String {
    let cell = cell_of(pos);
    let mut best = format!("({}, {})", cell.x, cell.y);
    let mut best_d = 6.0f32;
    for (name, poi) in &map.pois {
        if name.starts_with("home_") {
            continue;
        }
        let d = (world_of(*poi) - pos).length();
        if d < best_d {
            best_d = d;
            best = name.clone();
        }
    }
    best
}

fn build_nearby(
    map: &WorldMap,
    positions: &[(Entity, Vec3, String)],
    me: Entity,
) -> Vec<(String, String)> {
    let my_pos = match positions.iter().find(|(e, _, _)| *e == me) {
        Some((_, p, _)) => *p,
        None => return vec![],
    };
    let mut out = Vec::new();
    for (e, pos, name) in positions {
        if *e == me || pos.distance(my_pos) > 12.0 || out.len() >= 5 {
            continue;
        }
        out.push((name.clone(), place_at(map, *pos)));
    }
    out
}

// ---------------------------------------------------------------------------
// Scheduling
// ---------------------------------------------------------------------------

pub fn schedule_wakes(
    time: Res<Time>,
    mut pool: ResMut<LlmPool>,
    map: Res<WorldMap>,
    mut query: Query<(Entity, &Villager, &Needs, &Memory, &Transform, &mut Brain, &Pathing)>,
    all: Query<(Entity, &Transform, &Villager)>,
) {
    if pool.in_flight >= WORKERS {
        return;
    }
    let now = time.elapsed_secs();
    let positions: Vec<(Entity, Vec3, String)> = all
        .iter()
        .map(|(e, t, v)| (e, t.translation, v.persona.name.clone()))
        .collect();

    for (entity, villager, needs, memory, transform, mut brain, pathing) in query.iter_mut() {
        if brain.in_flight || now < brain.next_wake || pool.in_flight >= WORKERS {
            continue;
        }

        let p = &villager.persona;
        let home = map
            .place(&format!("home_{}", p.home_index))
            .map(|c| (c.x, c.y))
            .unwrap_or((32, 32));

        let at_place = if pathing.is_moving() {
            format!("walking to {}", pathing.destination)
        } else {
            place_at(&map, transform.translation)
        };
        let nearby = build_nearby(&map, &positions, entity);
        let nearby_str = if nearby.is_empty() {
            "nobody is near you".to_string()
        } else {
            nearby
                .iter()
                .map(|(n, pl)| format!("{n} (at {pl})"))
                .collect::<Vec<_>>()
                .join(", ")
        };
        let memories: Vec<String> = memory.recent(6).cloned().collect();
        let memories_str = if memories.is_empty() {
            "  (you remember nothing yet)".to_string()
        } else {
            memories
                .iter()
                .enumerate()
                .map(|(i, m)| format!("  {}. {m}", i + 1))
                .collect::<Vec<_>>()
                .join("\n")
        };
        let reply_hint = if brain.expect_reply {
            "\nSomeone recently spoke to you. You may respond to them with a talk action, \
             or do whatever you actually want to do."
        } else {
            ""
        };

        let user = format!(
            "[minute {:.0}] You are at: {at_place}.\n\
             Needs: hunger {hunger:.0}/100 (higher = hungrier), \
             energy {energy:.0}/100, social {social:.0}/100 (lower = lonelier).\n\
             Nearby: {nearby_str}.{reply_hint}\n\
             Your memories:\n{memories_str}\n\
             Choose your one action as JSON.",
            now / 60.0,
            hunger = needs.hunger,
            energy = needs.energy,
            social = needs.social,
        );

        let job = DecisionJob {
            entity,
            name: p.name.clone(),
            system: p.system_prompt(home),
            user,
        };
        if pool.jobs_tx.send(job).is_ok() {
            brain.in_flight = true;
            brain.expect_reply = false;
            pool.in_flight += 1;
            pool.peak_in_flight = pool.peak_in_flight.max(pool.in_flight);
        }
    }
}

// ---------------------------------------------------------------------------
// Results + action application
// ---------------------------------------------------------------------------

pub fn collect_results(
    time: Res<Time>,
    mut pool: ResMut<LlmPool>,
    map: Res<WorldMap>,
    mut events: MessageWriter<VillagerEvent>,
    mut query: Query<(
        Entity,
        &Villager,
        &Transform,
        &mut Brain,
        &mut Pathing,
        &mut Speech,
        &mut Memory,
        &mut Needs,
    )>,
) {
    loop {
        // Guard must drop before we mutate `pool` (edition-2021 while-let
        // would keep it alive through the body).
        let next = {
            let rx = pool.results_rx.lock().unwrap_or_else(|e| e.into_inner());
            rx.try_recv()
        };
        let Ok(result) = next else { break };
        pool.in_flight = pool.in_flight.saturating_sub(1);
        let now = time.elapsed_secs();

        let Ok((entity, villager, transform, mut brain, mut pathing, mut speech, mut memory, mut needs)) =
            query.get_mut(result.entity)
        else {
            continue;
        };

        brain.in_flight = false;
        brain.next_wake = now + WAKE_MIN + jitter(&villager.persona.name, now) * WAKE_JITTER;

        match result.action {
            Ok(action) => {
                pool.completed += 1;
                apply_action(
                    action,
                    &map,
                    &villager.persona.name,
                    transform.translation,
                    &mut pathing,
                    &mut speech,
                    &mut memory,
                    &mut needs,
                    now,
                    &mut events,
                    entity,
                );
            }
            Err(err) => {
                pool.failed += 1;
                brain.next_wake = now + 10.0; // don't stall the person
                memory.push(format!("(could not decide: {err})"));
            }
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn apply_action(
    action: Action,
    map: &WorldMap,
    my_name: &str,
    my_pos: Vec3,
    pathing: &mut Pathing,
    speech: &mut Speech,
    memory: &mut Memory,
    needs: &mut Needs,
    now: f32,
    events: &mut MessageWriter<VillagerEvent>,
    me: Entity,
) {
    match action {
        Action::GoTo { place, reason } => {
            let target = map
                .place(&place)
                .or_else(|| map.place("plaza"))
                .unwrap_or(IVec2::new(32, 32));
            match map.path(cell_of(my_pos), target) {
                Some(cells) => {
                    pathing.cells = cells;
                    pathing.step_progress = 0.0;
                    pathing.destination = place.clone();
                    memory.push(format!("Decided to go to the {place}: {reason}"));
                    events.write(VillagerEvent::Action { entity: me, text: format!("→ {place}") });
                }
                None => {
                    memory.push(format!("Wanted to go to the {place} but could not find a way."));
                }
            }
            needs.social = (needs.social - 0.2).max(0.0);
        }
        Action::Talk { to, say, reason } => {
            speech.text = say.clone();
            speech.expires = now + 7.0;
            memory.push(format!("Said to {to}: \"{say}\" ({reason})"));
            events.write(VillagerEvent::Talk { from: my_name.to_string(), to, say });
        }
        Action::Work { reason } => {
            memory.push(format!("Worked ({reason})"));
            needs.energy = (needs.energy - 4.0).max(0.0);
            needs.hunger = (needs.hunger + 2.0).min(100.0);
            events.write(VillagerEvent::Action { entity: me, text: "work".into() });
        }
        Action::Rest { reason } => {
            memory.push(format!("Rested ({reason})"));
            needs.energy = (needs.energy + 12.0).min(100.0);
            events.write(VillagerEvent::Action { entity: me, text: "rest".into() });
        }
        Action::Idle { reason } => {
            memory.push(format!("Loitered ({reason})"));
            events.write(VillagerEvent::Action { entity: me, text: "idle".into() });
        }
    }
}

/// Deterministic pseudo-jitter per (name, cycle) — no rand crate.
fn jitter(name: &str, t: f32) -> f32 {
    let mut h: u32 = 2166136261;
    for b in name.bytes() {
        h ^= b as u32;
        h = h.wrapping_mul(16777619);
    }
    h ^= (t / 1000.0) as u32;
    (h % 1000) as f32 / 1000.0
}

#[derive(Message, Debug)]
pub enum VillagerEvent {
    Talk { from: String, to: String, say: String },
    Action { entity: Entity, text: String },
}

/// Record what was said into the listener's memory + wake them to respond.
pub fn deliver_talks(
    time: Res<Time>,
    mut events: MessageReader<VillagerEvent>,
    names: Res<Names>,
    mut query: Query<(&mut Memory, &mut Brain)>,
) {
    let now = time.elapsed_secs();
    for ev in events.read() {
        if let VillagerEvent::Talk { from, to, say } = ev {
            if let Some(&target) = names.0.get(&to.to_lowercase()) {
                if let Ok((mut memory, mut brain)) = query.get_mut(target) {
                    memory.push(format!("{from} said: \"{say}\""));
                    if !brain.in_flight {
                        brain.next_wake = now + 1.0 + jitter(to, now) * 2.0;
                        brain.expect_reply = true;
                    }
                }
            }
        }
    }
}