use agentos_core::{CommitId, WorldlineId};
use agentos_kernel::ProcessCredentials;
use agentos_runtime::{BackendSelection, Runtime, RuntimeConfig, doctor};
use agentos_worldline::WorldlineStore;
use anyhow::{Result, bail};
use clap::{Parser, Subcommand};
use serde_json::json;
use std::path::PathBuf;

#[derive(Debug, Parser)]
#[command(name = "agentos", version, about = "Bare-Linux operating-system agent")]
struct Cli {
    #[arg(long, global = true)]
    state_dir: Option<PathBuf>,
    #[arg(long, global = true)]
    backend: Option<String>,
    #[arg(long, global = true)]
    fake_model: bool,
    #[arg(long, global = true)]
    json: bool,
    #[arg(long)]
    agent_once: Option<String>,
    #[arg(long)]
    doctor: bool,
    #[command(subcommand)]
    command: Option<Command>,
}

#[derive(Debug, Subcommand)]
enum Command {
    Doctor,
    Agent {
        #[command(subcommand)]
        command: AgentCommand,
    },
    Worldline {
        #[command(subcommand)]
        command: WorldlineCommand,
    },
}

#[derive(Debug, Subcommand)]
enum AgentCommand {
    Once { prompt: String },
}

#[derive(Debug, Subcommand)]
enum WorldlineCommand {
    Init,
    Status,
    Branch {
        id: Option<String>,
    },
    Diff {
        id: String,
    },
    Commit {
        id: String,
        #[arg(short, long)]
        message: String,
    },
    Rollback {
        commit: String,
        #[arg(short, long, default_value = "rollback")]
        message: String,
    },
    Discard {
        id: String,
    },
    Log {
        #[arg(short, long, default_value_t = 20)]
        limit: usize,
    },
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    if (cli.agent_once.is_some() || cli.doctor) && cli.command.is_some() {
        bail!("legacy flags --doctor/--agent-once cannot be combined with a subcommand");
    }
    let mut config = RuntimeConfig::from_env()?;
    if let Some(state_dir) = &cli.state_dir {
        config.state_dir.clone_from(state_dir);
    }
    if let Some(backend) = &cli.backend {
        config.backend = backend.parse()?;
    }
    if cli.fake_model {
        config.backend = BackendSelection::Fake;
    }

    if cli.doctor || matches!(cli.command, Some(Command::Doctor)) {
        let report = doctor(&config)?;
        print_value(&serde_json::to_value(report)?, cli.json);
        return Ok(());
    }
    if let Some(prompt) = cli.agent_once {
        return run_agent(config, &prompt, cli.json).await;
    }
    match cli.command {
        Some(Command::Agent {
            command: AgentCommand::Once { prompt },
        }) => run_agent(config, &prompt, cli.json).await,
        Some(Command::Worldline { command }) => run_worldline(&config, command, cli.json),
        Some(Command::Doctor) => unreachable!(),
        None => bail!("select --doctor, --agent-once, or a subcommand"),
    }
}

async fn run_agent(config: RuntimeConfig, prompt: &str, json_output: bool) -> Result<()> {
    let runtime = Runtime::initialize(config)?;
    let result = runtime.run_once(prompt).await?;
    if json_output {
        println!("{}", serde_json::to_string_pretty(&result)?);
    } else {
        println!("{}", result.content);
    }
    Ok(())
}

fn run_worldline(
    config: &RuntimeConfig,
    command: WorldlineCommand,
    json_output: bool,
) -> Result<()> {
    let root = config.worldline_root();
    let store = match command {
        WorldlineCommand::Init => {
            let store = WorldlineStore::initialize(&root)?;
            print_value(
                &json!({
                    "root": store.root(),
                    "backend": store.backend(),
                    "current": store.current_path(),
                }),
                json_output,
            );
            return Ok(());
        }
        _ if root.join("backend.json").is_file() => WorldlineStore::open(&root)?,
        _ => WorldlineStore::initialize(&root)?,
    };
    let uid = ProcessCredentials::current().effective_uid;
    match command {
        WorldlineCommand::Init => unreachable!(),
        WorldlineCommand::Status => print_value(
            &json!({
                "root": store.root(),
                "backend": store.backend(),
                "head": store.head()?,
                "current": store.current_path(),
            }),
            json_output,
        ),
        WorldlineCommand::Branch { id } => {
            let id = id.map_or_else(|| Ok(WorldlineId::new()), WorldlineId::parse)?;
            let path = store.create_branch(&id)?;
            print_value(&json!({"worldline_id": id, "path": path}), json_output);
        }
        WorldlineCommand::Diff { id } => {
            let id = WorldlineId::parse(id)?;
            print_value(&serde_json::to_value(store.diff_branch(&id)?)?, json_output);
        }
        WorldlineCommand::Commit { id, message } => {
            let commit = store.commit_branch(&WorldlineId::parse(id)?, message, uid)?;
            print_value(&serde_json::to_value(commit)?, json_output);
        }
        WorldlineCommand::Rollback { commit, message } => {
            let commit = store.rollback_to(&CommitId::parse(commit)?, message, uid)?;
            print_value(&serde_json::to_value(commit)?, json_output);
        }
        WorldlineCommand::Discard { id } => {
            let id = WorldlineId::parse(id)?;
            print_value(
                &json!({"worldline_id": id, "discarded": store.discard_branch(&id)?}),
                json_output,
            );
        }
        WorldlineCommand::Log { limit } => {
            print_value(&serde_json::to_value(store.log(limit)?)?, json_output);
        }
    }
    Ok(())
}

fn print_value(value: &serde_json::Value, pretty: bool) {
    if pretty {
        println!(
            "{}",
            serde_json::to_string_pretty(value).expect("JSON value serializes")
        );
    } else if let Some(content) = value.as_str() {
        println!("{content}");
    } else {
        println!(
            "{}",
            serde_json::to_string_pretty(value).expect("JSON value serializes")
        );
    }
}