ninox-core 0.2.1

use crate::{
    config::AppConfig,
    events::{Engine, Event},
    github::{split_repo, CheckRun},
    hooks,
    lifecycle::{
        enrichment::EnrichmentCache,
        probe::is_pid_alive,
        usage,
    },
    types::{
        CIStatus, Comment, Notification, NotificationKind, PrId, SessionStatus, PR,
    },
};
use std::{collections::HashMap, sync::Arc, time::Duration};
use tokio_util::sync::CancellationToken;

/// Unix epoch milliseconds "now" — used to stamp `Notification::created_at`.
fn now_millis() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_millis() as i64)
        .unwrap_or(0)
}

pub struct Poller {
    engine:           Arc<Engine>,
    enrichment_cache: Arc<std::sync::Mutex<EnrichmentCache>>,
}

impl Poller {
    pub fn new(engine: Arc<Engine>) -> Self {
        Self {
            engine,
            enrichment_cache: Arc::new(std::sync::Mutex::new(HashMap::new())),
        }
    }

    pub async fn start(self, token: CancellationToken) {
        let mut pid_interval    = tokio::time::interval(Duration::from_secs(5));
        let mut usage_interval  = tokio::time::interval(Duration::from_secs(10));
        let mut github_interval = tokio::time::interval(Duration::from_secs(30));
        // Prevent a missed tick from causing back-to-back polls.
        github_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
        usage_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
        loop {
            tokio::select! {
                _ = token.cancelled()      => break,
                _ = pid_interval.tick()    => self.poll_pids().await,
                _ = usage_interval.tick()  => self.poll_usage().await,
                _ = github_interval.tick() => self.poll_github().await,
            }
        }
    }

    // ── PID liveness ────────────────────────────────────────────────────────

    async fn poll_pids(&self) {
        let Ok(sessions) = self.engine.store.list_sessions() else { return };
        for mut session in sessions {
            if matches!(session.status, SessionStatus::Done | SessionStatus::Terminated) {
                continue;
            }
            if let Some(pid) = session.pid {
                if !is_pid_alive(pid) {
                    session.status = SessionStatus::Terminated;
                    let _ = self.engine.store.upsert_session(&session);
                    self.engine.emit(Event::SessionUpdated(session));
                    continue;
                }
            }

            // Poll metadata files for PR number on working sessions that have none yet.
            if matches!(session.status, SessionStatus::Working | SessionStatus::Spawning)
                && session.pr_number.is_none()
            {
                let sessions_dir = AppConfig::sessions_dir();
                if let Ok(meta) = hooks::read_session_metadata(&sessions_dir, &session.id) {
                    if let Some(pr_num) = meta.pr_number {
                        session.pr_number = Some(pr_num);
                        session.status    = SessionStatus::PrOpen;
                        let _ = self.engine.store.upsert_session(&session);
                        self.engine.emit(Event::SessionUpdated(session.clone()));
                        tracing::info!(
                            "session {} PR #{pr_num} detected via metadata hook",
                            session.id
                        );
                    }
                }
            }
        }
    }

    // ── Cost / context-window usage ─────────────────────────────────────────

    /// Ingest cost/token usage for every active session by reading `claude`'s
    /// own transcript for the session's workspace directory (see
    /// `lifecycle::usage`). Sessions without a workspace, or whose transcript
    /// has no usage yet (agent hasn't taken a turn), are left untouched.
    /// Only writes + emits when something actually changed, so this doesn't
    /// spam the store/UI every tick for idle sessions.
    async fn poll_usage(&self) {
        let Ok(sessions) = self.engine.store.list_sessions() else { return };
        for mut session in sessions {
            if matches!(session.status, SessionStatus::Done | SessionStatus::Terminated) {
                continue;
            }
            let Some(workspace) = session.workspace_path.clone() else { continue };
            let Some(snapshot) = usage::ingest_usage_for_workspace(&workspace) else { continue };

            let cost_changed = (session.cost_usd - snapshot.cost_usd).abs() > 1e-9;
            let context_changed = session.context_tokens != Some(snapshot.context_tokens);
            if !cost_changed && !context_changed {
                continue;
            }

            session.cost_usd = snapshot.cost_usd;
            session.context_tokens = Some(snapshot.context_tokens);
            if session.model.is_none() {
                session.model = snapshot.model;
            }
            let _ = self.engine.store.upsert_session(&session);
            self.engine.emit(Event::SessionUpdated(session));
        }
    }

    // ── GitHub enrichment ────────────────────────────────────────────────────

    async fn poll_github(&self) {
        let Some(gh) = &self.engine.github else { return };
        let Ok(sessions) = self.engine.store.list_sessions() else { return };

        for session in sessions {
            if matches!(session.status, SessionStatus::Done | SessionStatus::Terminated) {
                continue;
            }
            let Some(pr_number) = session.pr_number else { continue };
            let Some((owner, repo)) = split_repo(&session.repo) else { continue };

            // -- PR state --
            let pr_status = match gh.get_pr_status(&owner, &repo, pr_number).await {
                Ok(s)  => s,
                Err(e) => { tracing::warn!("github pr status: {e}"); continue }
            };

            let pr_id: PrId = pr_number as i64;

            // -- Merge detection — handle before CI (no point polling CI on merged PR) --
            if pr_status.merged && !matches!(session.status, SessionStatus::Done) {
                self.engine.emit(Event::Notification(Notification {
                    id:         format!("merged-{}", session.id),
                    kind:       NotificationKind::WorkerDone,
                    title:      format!("PR merged — {}", session.name),
                    body:       format!("#{} merged successfully", pr_number),
                    session_id: Some(session.id.clone()),
                    created_at: now_millis(),
                }));
                if let Err(e) = self.engine.cleanup_session(&session.id).await {
                    tracing::warn!("cleanup_session {}: {e}", session.id);
                }
                // Remove enrichment state for this session — it's done
                {
                    let mut cache = self.enrichment_cache.lock().unwrap();
                    cache.remove(&session.id);
                }
                continue; // skip further enrichment for this session
            }

            // Upsert PR record — only when not merged (merged sessions stay Done after cleanup)
            {
                let pr = PR {
                    id:         pr_id,
                    number:     pr_number,
                    title:      pr_status.title.clone(),
                    url:        format!("https://github.com/{owner}/{repo}/pull/{pr_number}"),
                    body:       String::new(),
                    session_id: session.id.clone(),
                };
                let _ = self.engine.store.upsert_pr(&pr);
                self.engine.emit(Event::PrOpened { session_id: session.id.clone(), pr });
            }

            // -- CI checks --
            let checks = match gh.get_ci_checks(&owner, &repo, &pr_status.head_sha).await {
                Ok(c)  => c,
                Err(e) => { tracing::warn!("github ci checks: {e}"); vec![] }
            };
            let ci = summarize_checks(pr_id, &checks);
            let _ = self.engine.store.upsert_ci_status(&ci);
            self.engine.emit(Event::CiUpdated { pr_id, status: ci.clone() });

            // -- Detect CI transition and update session status --
            let (newly_failing, ci_reaction_already_sent) = {
                let mut cache = self.enrichment_cache.lock().unwrap();
                let state = cache.entry(session.id.clone()).or_default();

                let newly_failing = state.prev_failing.is_none_or(|p| p == 0)
                    && ci.failing > 0;
                state.prev_failing = Some(ci.failing);

                let already_sent = state.ci_reaction_sent;
                if newly_failing && !already_sent {
                    state.ci_reaction_sent = true;
                }
                if ci.failing == 0 {
                    state.ci_reaction_sent = false;
                }
                (newly_failing, already_sent)
            };

            if newly_failing && !ci_reaction_already_sent {
                self.engine.emit(Event::Notification(Notification {
                    id:         format!("ci-{}", session.id),
                    kind:       NotificationKind::CiFailure,
                    title:      format!("CI failing — {}", session.name),
                    body:       format!("{}/{} checks failing", ci.failing, ci.total),
                    session_id: Some(session.id.clone()),
                    created_at: now_millis(),
                }));
                // Send reaction to the agent in the tmux session
                let failing_names: Vec<String> = checks.iter()
                    .filter(|c| c.conclusion.as_deref() == Some("failure")
                             || c.conclusion.as_deref() == Some("timed_out"))
                    .map(|c| c.name.clone())
                    .collect();
                let msg = crate::lifecycle::reactions::format_ci_reaction(
                    &session, &ci, &failing_names
                );
                if let Err(e) = self.engine.send_to_session(&session.id, &msg).await {
                    tracing::warn!("send ci reaction to {}: {e}", session.id);
                }
            }

            // -- Review threads (throttled via seen_comment_ids) --
            let threads = match gh.get_review_threads(&owner, &repo, pr_number).await {
                Ok(t)  => t,
                Err(e) => { tracing::warn!("github review threads: {e}"); vec![] }
            };

            let has_changes_requested = threads.iter().any(|t| t.state == "CHANGES_REQUESTED");

            let (has_new, review_reaction_already_sent, new_comments) = {
                let mut cache = self.enrichment_cache.lock().unwrap();
                let state = cache.entry(session.id.clone()).or_default();
                let mut has_new = false;
                let mut new_comments: Vec<Comment> = Vec::new();

                for thread in &threads {
                    if thread.state == "CHANGES_REQUESTED"
                        && !state.seen_comment_ids.contains(&thread.id)
                    {
                        state.seen_comment_ids.insert(thread.id);
                        has_new = true;
                        let comment = Comment {
                            id:         thread.id,
                            pr_id,
                            author:     thread.author.clone(),
                            body:       thread.body.clone(),
                            path:       thread.path.clone(),
                            line:       thread.line,
                            created_at: 0,
                        };
                        let _ = self.engine.store.upsert_comment(&comment);
                        self.engine.emit(Event::ReviewComment { pr_id, comment: comment.clone() });
                        new_comments.push(comment);
                    }
                }

                let already_sent = state.review_reaction_sent;
                if has_new && !already_sent {
                    state.review_reaction_sent = true;
                }
                // Reset when all CHANGES_REQUESTED are resolved
                if !has_changes_requested {
                    state.review_reaction_sent = false;
                }
                (has_new, already_sent, new_comments)
            };

            // Update session status in DB (after review threads so has_changes_requested is known)
            let new_status = derive_session_status(&session.status, &pr_status, &ci, has_changes_requested);
            let mut updated = session.clone();
            updated.status = new_status;
            if updated.status != session.status {
                let _ = self.engine.store.upsert_session(&updated);
                self.engine.emit(Event::SessionUpdated(updated.clone()));
            }

            if has_new && !review_reaction_already_sent {
                self.engine.emit(Event::Notification(Notification {
                    id:         format!("review-{}", session.id),
                    kind:       NotificationKind::PrNeedsAttention,
                    title:      format!("Review comments — {}", session.name),
                    body:       "Changes requested on your PR".to_string(),
                    session_id: Some(session.id.clone()),
                    created_at: now_millis(),
                }));
                if !new_comments.is_empty() {
                    let msg = crate::lifecycle::reactions::format_review_reaction(
                        &session, &new_comments
                    );
                    if let Err(e) = self.engine.send_to_session(&session.id, &msg).await {
                        tracing::warn!("send review reaction to {}: {e}", session.id);
                    }
                }
            }
        }
    }
}

// ── Helpers ──────────────────────────────────────────────────────────────────

fn summarize_checks(pr_id: PrId, checks: &[CheckRun]) -> CIStatus {
    let total   = checks.len() as u32;
    let failing = checks.iter().filter(|c| {
        c.conclusion.as_deref() == Some("failure")
            || c.conclusion.as_deref() == Some("timed_out")
    }).count() as u32;
    let passing = checks.iter().filter(|c| {
        c.conclusion.as_deref() == Some("success")
    }).count() as u32;
    let pending = total - failing - passing;
    CIStatus { pr_id, total, failing, passing, pending }
}

fn derive_session_status(
    current:               &SessionStatus,
    pr_status:             &crate::github::PrStatus,
    ci:                    &CIStatus,
    has_changes_requested: bool,
) -> SessionStatus {
    // Terminal states are never overwritten.
    if matches!(current, SessionStatus::Done | SessionStatus::Terminated) {
        return current.clone();
    }
    if pr_status.merged {
        return SessionStatus::Done;
    }
    if ci.failing > 0 {
        return SessionStatus::CiFailed;
    }
    if has_changes_requested {
        return SessionStatus::ReviewPending;
    }
    if pr_status.mergeable == Some(true) && ci.failing == 0 && ci.pending == 0 {
        return SessionStatus::Mergeable;
    }
    SessionStatus::PrOpen
}

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

    #[test]
    fn summarize_checks_counts_failures() {
        let checks = vec![
            CheckRun { name: "lint".into(), status: "completed".into(), conclusion: Some("success".into()) },
            CheckRun { name: "test".into(), status: "completed".into(), conclusion: Some("failure".into()) },
            CheckRun { name: "build".into(), status: "in_progress".into(), conclusion: None },
        ];
        let ci = summarize_checks(1, &checks);
        assert_eq!(ci.total,   3);
        assert_eq!(ci.passing, 1);
        assert_eq!(ci.failing, 1);
        assert_eq!(ci.pending, 1);
    }

    #[test]
    fn derive_status_merged_becomes_done() {
        let pr = crate::github::PrStatus {
            merged: true, state: "closed".into(), mergeable: None,
            title: "t".into(), number: 1, head_sha: String::new(),
        };
        let ci = CIStatus { pr_id: 1, total: 0, failing: 0, passing: 0, pending: 0 };
        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, false);
        assert!(matches!(s, SessionStatus::Done));
    }

    #[test]
    fn derive_status_ci_failure_overrides_open() {
        let pr = crate::github::PrStatus {
            merged: false, state: "open".into(), mergeable: Some(true),
            title: "t".into(), number: 1, head_sha: String::new(),
        };
        let ci = CIStatus { pr_id: 1, total: 3, failing: 1, passing: 2, pending: 0 };
        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, false);
        assert!(matches!(s, SessionStatus::CiFailed));
    }

    #[test]
    fn derive_status_all_green_becomes_mergeable() {
        let pr = crate::github::PrStatus {
            merged: false, state: "open".into(), mergeable: Some(true),
            title: "t".into(), number: 1, head_sha: String::new(),
        };
        let ci = CIStatus { pr_id: 1, total: 3, failing: 0, passing: 3, pending: 0 };
        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, false);
        assert!(matches!(s, SessionStatus::Mergeable));
    }

    #[test]
    fn derive_status_preserves_done() {
        let pr = crate::github::PrStatus {
            merged: false, state: "open".into(), mergeable: Some(true),
            title: "t".into(), number: 1, head_sha: String::new(),
        };
        let ci = CIStatus { pr_id: 1, total: 0, failing: 0, passing: 0, pending: 0 };
        let s  = derive_session_status(&SessionStatus::Done, &pr, &ci, false);
        assert!(matches!(s, SessionStatus::Done));
    }

    #[test]
    fn derive_status_preserves_terminated() {
        let pr = crate::github::PrStatus {
            merged: true, state: "closed".into(), mergeable: None,   // merged=true!
            title: "t".into(), number: 1, head_sha: String::new(),
        };
        let ci = CIStatus { pr_id: 1, total: 0, failing: 0, passing: 0, pending: 0 };
        let s  = derive_session_status(&SessionStatus::Terminated, &pr, &ci, false);
        assert!(matches!(s, SessionStatus::Terminated));  // must not become Done
    }

    #[test]
    fn derive_status_changes_requested_becomes_review_pending() {
        let pr = crate::github::PrStatus {
            merged: false, state: "open".into(), mergeable: Some(true),
            title: "t".into(), number: 1, head_sha: String::new(),
        };
        let ci = CIStatus { pr_id: 1, total: 3, failing: 0, passing: 3, pending: 0 };
        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, true);
        assert!(matches!(s, SessionStatus::ReviewPending));
    }

    fn test_session(id: &str, workspace: &str) -> crate::types::Session {
        crate::types::Session {
            id: id.into(), orchestrator_id: None, name: id.into(),
            repo: String::new(), status: SessionStatus::Working,
            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
            pr_number: None, pr_id: None,
            workspace_path: Some(workspace.into()), pid: None,
            model: None, context_tokens: None, catalogue_path: None,
        }
    }

    /// End-to-end (within-process) proof that the poller closes the gap
    /// documented in `lifecycle::usage`: given a workspace whose `claude`
    /// transcript directory has usage recorded, `poll_usage` writes the
    /// derived cost/context/model back into the store and emits
    /// `SessionUpdated` — the exact path the UI's $0.0000 / missing-tokens
    /// symptom traces back to when this ingestion doesn't happen.
    // The `ENV_TEST_GUARD` mutex is intentionally held across the `.await`
    // points below — it serializes access to the process-global
    // `NINOX_CLAUDE_PROJECTS_DIR` env var against other tests (in this file
    // and in `lifecycle::usage`) for this single-threaded `#[tokio::test]`,
    // and must stay held for the env var's entire lifetime, not just around
    // the sync portions.
    #[allow(clippy::await_holding_lock)]
    #[tokio::test]
    async fn poll_usage_ingests_transcript_into_store_and_emits_update() {
        use crate::{lifecycle::usage::{claude_project_slug, ENV_TEST_GUARD}, store::Store};
        use std::io::Write;

        let _guard = ENV_TEST_GUARD.lock().unwrap_or_else(|e| e.into_inner());
        let projects_dir = tempfile::tempdir().unwrap();
        let workspace = "/tmp/poller-usage-probe-workspace";
        let project_dir = projects_dir.path().join(claude_project_slug(workspace));
        std::fs::create_dir_all(&project_dir).unwrap();
        let mut f = std::fs::File::create(project_dir.join("s.jsonl")).unwrap();
        writeln!(
            f,
            r#"{{"type":"assistant","timestamp":"2026-07-05T13:00:00.000Z","message":{{"model":"claude-fable-5","usage":{{"input_tokens":2,"output_tokens":300,"cache_creation_input_tokens":500,"cache_read_input_tokens":45000}}}}}}"#
        ).unwrap();
        drop(f);

        let prior = std::env::var("NINOX_CLAUDE_PROJECTS_DIR").ok();
        std::env::set_var("NINOX_CLAUDE_PROJECTS_DIR", projects_dir.path());

        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
        store.upsert_session(&test_session("s1", workspace)).unwrap();
        let engine = Engine::new(store.clone());
        let mut rx = engine.subscribe();
        let poller = Poller::new(engine);

        poller.poll_usage().await;

        match prior {
            Some(v) => std::env::set_var("NINOX_CLAUDE_PROJECTS_DIR", v),
            None    => std::env::remove_var("NINOX_CLAUDE_PROJECTS_DIR"),
        }

        let updated = store.get_session("s1").unwrap().unwrap();
        assert!(updated.cost_usd > 0.0, "cost_usd should be ingested, not 0.0000");
        assert_eq!(updated.context_tokens, Some(2 + 500 + 45000));
        assert_eq!(updated.model.as_deref(), Some("claude-fable-5"));

        let evt = tokio::time::timeout(std::time::Duration::from_secs(1), rx.recv())
            .await
            .expect("SessionUpdated should be emitted")
            .unwrap();
        assert!(matches!(evt, Event::SessionUpdated(s) if s.id == "s1" && s.cost_usd > 0.0));
    }

    #[tokio::test]
    async fn poll_usage_leaves_sessions_without_workspace_or_usage_untouched() {
        use crate::store::Store;

        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
        let mut no_ws = test_session("no-ws", "/does/not/matter");
        no_ws.workspace_path = None;
        store.upsert_session(&no_ws).unwrap();
        let engine = Engine::new(store.clone());
        let poller = Poller::new(engine);

        poller.poll_usage().await;

        let unchanged = store.get_session("no-ws").unwrap().unwrap();
        assert_eq!(unchanged.cost_usd, 0.0);
        assert_eq!(unchanged.context_tokens, None);
    }
}