apalis

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TripRequest {
    pub origin: String,
    pub destination: String,
    pub departure_date: String,
    pub return_date: String,
}

use apalis::prelude::BoxDynError;
use rig::providers::openai;
use rig::completion::Prompt;

pub async fn build_itinerary(trip: TripRequest) -> Result<String, BoxDynError> {
    let client = openai::Client::from_env();
    let agent = client
        .agent("gpt-4")
        .preamble("You are a travel expert. Given an origin, destination, and dates, \
                   suggest a day-by-day itinerary with activities, landmarks, \
                   and local experiences. Be specific and practical.")
        .build();

    let prompt = format!(
        "Create a 7-day itinerary for a trip from {} to {}, \
         departing {} and returning {}.",
        trip.origin, trip.destination, trip.departure_date, trip.return_date
    );

    Ok(agent.prompt(&prompt).await?)
}

pub async fn find_requirements(trip: TripRequest) -> Result<String, BoxDynError> {
    let client = openai::Client::from_env();
    let agent = client
        .agent("gpt-4")
        .preamble("You are a travel documentation specialist. Provide accurate, \
                   current information about visa requirements, recommended and \
                   mandatory vaccinations, travel advisories, and any health \
                   documentation needed for entry.")
        .build();

    let prompt = format!(
        "What are the visa requirements, vaccination recommendations, \
         and entry documentation needed for a traveller from {} visiting {}?",
        trip.origin, trip.destination
    );

    Ok(agent.prompt(&prompt).await?)
}

pub async fn find_hotels(trip: TripRequest) -> Result<String, BoxDynError> {
    let client = openai::Client::from_env();
    let agent = client
        .agent("gpt-4")
        .preamble("You are a hotel and accommodation expert. Recommend a range of \
                   accommodation options from budget to luxury, including neighbourhoods \
                   to stay in, what to look for, and approximate nightly rates.")
        .build();

    let prompt = format!(
        "Recommend accommodation options in {} for a visitor from {}. \
         Include budget, mid-range, and luxury options with approximate prices \
         and the best neighbourhoods to stay in.",
        trip.destination, trip.origin
    );

    Ok(agent.prompt(&prompt).await?)
}

pub async fn find_flights(trip: TripRequest) -> Result<String, BoxDynError> {
    let client = openai::Client::from_env();
    let agent = client
        .agent("gpt-4")
        .preamble("You are a flight search assistant. Provide practical guidance \
                   on finding the best flights, including typical airlines on the route, \
                   approximate prices, layover options, and booking advice.")
        .build();

    let prompt = format!(
        "What are the best flight options from {} to {} departing around {}? \
         Include typical airlines, approximate costs, direct vs connecting options, \
         and the best time to book.",
        trip.origin, trip.destination, trip.departure_date
    );

    Ok(agent.prompt(&prompt).await?)
}

pub async fn generate_report(
    (itinerary, requirements, hotels, flights): (String, String, String, String),
    wrk: WorkerContext,
) -> Result<String, BoxDynError> {
    let client = openai::Client::from_env();
    let agent = client
        .agent("gpt-4")
        .preamble("You are a professional travel writer. Given structured research \
                   on a destination, produce a beautifully formatted markdown travel \
                   report that a traveller can save and reference throughout their trip. \
                   Include all sections, use clear headings, and write in an engaging tone.")
        .build();

    let prompt = format!(
        "Produce a complete markdown travel report using the following research:\n\n\
         ## Itinerary\n{itinerary}\n\n\
         ## Entry Requirements\n{requirements}\n\n\
         ## Accommodation\n{hotels}\n\n\
         ## Flights\n{flights}"
    );

    let report = agent.prompt(&prompt).await?;

    // Write the report to disk so the caller can retrieve it.
    tokio::fs::write("travel_report.md", &report).await?;
    tracing::info!("Travel report written to travel_report.md");

    wrk.stop().unwrap();
    Ok(report)
}

use apalis::prelude::*;
use apalis_file_storage::JsonStorage;
use apalis_workflow::{DagFlow, WorkflowSink};
use tracing::info;

#[tokio::main]
async fn main() -> Result<(), BoxDynError> {
    tracing_subscriber::fmt::init();

    // 1. Set up the backend.
    let mut backend = JsonStorage::new_temp().unwrap();

    // 2. Push one or more trip requests as starting inputs.
    //    Each request will run through the full DAG independently.
    backend
        .push_start(vec![
            TripRequest {
                origin: "Nairobi".into(),
                destination: "Tokyo".into(),
                departure_date: "2026-07-01".into(),
                return_date: "2026-07-14".into(),
            }
        ])
        .await?;

    // 3. Declare the DAG.
    let dag = DagFlow::new("travel-planner");

    let itinerary    = dag.node(build_itinerary);
    let requirements = dag.node(find_requirements);
    let hotels       = dag.node(find_hotels);
    let flights      = dag.node(find_flights);

    // The collector depends on all four — order here must match
    // the tuple in generate_report's signature.
    dag.node(generate_report)
        .depends_on((&itinerary, &requirements, &hotels, &flights));

    // 4. Validate before running — catches cycles and missing dependencies.
    dag.validate()?;

    // 5. Print the graph structure for debugging.
    //    Paste the output at https://dreampuf.github.io/GraphvizOnline/
    info!("DAG structure:\n{dag}");

    // 6. Run — the worker drives the DAG like any other backend.
    WorkerBuilder::new("travel-planner-worker")
        .backend(backend)
        .enable_tracing()
        .on_event(|_ctx, event| info!("{event}"))
        .build(dag)
        .run()
        .await?;

    Ok(())
}