use rig::client::{AgentClientExt, Nothing};
use rig::completion::Prompt;
use rig::providers::ollama;
use crate::deep_research::review::{self, Review};
use crate::deep_research::tools::{FetchPage, SearchWeb};

/// The tool-calling research loop needs to reliably decide what to search
/// for, when a page is worth fetching, and when it has enough evidence —
/// that's a reasoning-heavy job best given to the largest local Gemma
/// variant. Turning the gathered notes into prose afterwards is comparatively
/// mechanical, so the smaller/faster variant handles that pass instead.
const RESEARCHER_MODEL: &str = "gemma4:26b";
const WRITER_MODEL: &str = "gemma4-e4b:latest";
const MAX_RESEARCH_TURNS: usize = 12;
/// Research/review rounds before giving up and writing the report from
/// whatever the last pass produced, rather than looping forever on a topic
/// the reviewer can never be satisfied with.
const MAX_RESEARCH_ROUNDS: usize = 3;

const DEFAULT_TOPIC: &str =
    "What are the latest advances in running large language models locally, on consumer hardware?";

fn initialize_observability() {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE)
        .with_writer(std::io::stderr)
        .init();
}

/// Performs deep research via a two-stage agentic flow: a tool-calling agent
/// gathers and cross-checks evidence from the web, then a second agent turns
/// those raw notes into a structured report.
pub(crate) async fn start() -> anyhow::Result<()> {
    initialize_observability();

    let topic = std::env::args().nth(1).unwrap_or_else(|| DEFAULT_TOPIC.to_string());

    let report = research(&topic).await?;

    println!("{report}");

    Ok(())
}

/// The least-agentic shape that fits: a plain Rust loop putting *this code*,
/// not a model, in charge of when to stop — re-running research with the
/// reviewer's feedback folded in until it approves or the round budget runs
/// out, then writing the report from whatever the last pass produced.
async fn research(topic: &str) -> anyhow::Result<String> {
    let client = ollama::Client::new(Nothing)?;

    let mut findings = String::new();
    let mut feedback: Option<Review> = None;

    for round in 1..=MAX_RESEARCH_ROUNDS {
        findings = gather_findings(&client, topic, feedback.as_ref(), round).await?;

        let review = review::review_findings(&client, topic, &findings).await?;
        let approved = review.approved;

        tracing::info!(round, approved, "review verdict");

        if approved || round == MAX_RESEARCH_ROUNDS {
            break;
        }

        feedback = Some(review);
    }

    write_report(&client, topic, &findings).await
}

/// Wraps the tool-calling research loop in its own span so it's visible as a
/// single unit in traces, distinct from the writing and review phases and
/// nesting rig's own per-turn `chat`/`execute_tool` spans underneath it.
#[tracing::instrument(skip(client, feedback), fields(gen_ai.agent.name = "researcher"))]
async fn gather_findings(
    client: &ollama::Client,
    topic: &str,
    feedback: Option<&Review>,
    round: usize,
) -> anyhow::Result<String> {
    let researcher = client
        .agent(RESEARCHER_MODEL)
        .name("researcher")
        .preamble(
            "You are a meticulous research assistant. Use the search_web and fetch_page tools to \
             investigate the user's topic: run several searches with varied phrasing, fetch the \
             most promising pages, and cross-check claims across at least two sources before \
             trusting them. Once you are confident you have enough evidence, stop calling tools \
             and reply with a plain-text dump of every fact you gathered, the source URL it came \
             from, and any open questions or contradictions between sources. This is raw research \
             material for a writer, not a final report, so favor completeness over polish.",
        )
        .tool(SearchWeb)
        .tool(FetchPage)
        .build();

    let task = match feedback {
        None => topic.to_string(),
        Some(review) => format!(
            "Topic: {topic}\n\n\
             You already ran a research pass on this topic. A reviewer checked it against its \
             cited sources and found it insufficient. Do more research to address the reviewer's \
             feedback, then produce an updated findings dump: carry forward what's solid, and \
             add, correct, or better-source whatever the gaps call for.\n\n\
             Solid findings from the last pass — keep and build on these:\n{}\n\n\
             Gaps the reviewer found — conclusions not actually backed by their source, sources \
             that don't line up with the conclusion drawn from them, or parts of the topic still \
             uncovered:\n{}",
            review.solid_findings, review.gaps
        ),
    };

    let findings = researcher
        .runner(task)
        .max_turns(MAX_RESEARCH_TURNS)
        .run()
        .await?
        .output;

    tracing::info!(round, findings = %findings, "research phase complete");

    Ok(findings)
}

#[tracing::instrument(skip(client, findings), fields(gen_ai.agent.name = "writer"))]
async fn write_report(client: &ollama::Client, topic: &str, findings: &str) -> anyhow::Result<String> {
    let writer = client
        .agent(WRITER_MODEL)
        .name("writer")
        .preamble(
            "You turn raw research notes into a clear, well-organized report for the reader. \
             Structure the report with headings, cite source URLs inline next to the claims they \
             support, and call out any open questions or contradictions the research turned up. Do \
             not invent facts beyond what the notes provide.",
        )
        .build();

    let report = writer
        .prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}"))
        .await?;

    Ok(report)
}