mod audit_emit;
mod auth;
mod config;
mod error;
mod handlers;
mod orchestrator;
mod persistence;
mod recovery;
mod types;

use std::{path::PathBuf, sync::Arc};

use async_nats::{
    jetstream::{self, consumer::pull, AckKind},
    ConnectOptions,
};
use clap::Parser;
use futures::StreamExt;
use surrealdb::{engine::any, Surreal};
use tracing::{error, info, warn};

use audit_emit::{emit_ack, emit_audit};
use auth::JwtAuth;
use config::ControllerConfig;
use error::ControllerError;
use orchestrator::OrchestratorClient;
use persistence::Persistence;
use recovery::{check_idempotency, reconcile_pending};
use types::AckResult;

#[derive(Parser)]
#[command(name = "ops-controller", about = "WorkQueue subscriber for ops.cmd.* (ADR-037)")]
struct Cli {
    #[arg(short, long, default_value = "ops-controller.toml")]
    config: PathBuf,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .init();

    let cli = Cli::parse();
    let cfg = ControllerConfig::from_toml(&cli.config)?;

    // Load public key for JWT verification
    let public_pem = std::fs::read(&cfg.keeper_public_key_path)?;
    let jwt_auth = Arc::new(JwtAuth::from_public_pem(&public_pem, cfg.workspace.clone())?);

    // Connect SurrealDB
    let db: Surreal<any::Any> = any::connect(&cfg.surrealdb_url).await
        .map_err(|e| format!("SurrealDB connect to {}: {e}", cfg.surrealdb_url))?;
    let persistence = Arc::new(Persistence::new(db, cfg.workspace.clone()).await?);

    // Connect orchestrator HTTP client
    let orc = Arc::new(OrchestratorClient::new(
        cfg.orchestrator_url.clone(),
        cfg.orchestrator_auth_token.clone(),
    )?);

    // Connect NATS JetStream
    let nats_opts = match &cfg.nats_auth_token {
        Some(token) => ConnectOptions::new().token(token.clone()),
        None => ConnectOptions::new(),
    };
    let nats = nats_opts.connect(cfg.nats_url.as_str()).await
        .map_err(|e| format!("NATS connect: {e}"))?;
    let js = Arc::new(jetstream::new(nats));

    // Recovery pass on startup
    reconcile_pending(
        Arc::clone(&persistence),
        Arc::clone(&orc),
        Arc::clone(&js),
        &cfg.workspace,
    ).await;

    // Subscribe to OPS_CMD_<workspace> as a single durable WorkQueue consumer
    let stream_name = format!("OPS_CMD_{}", cfg.workspace.to_uppercase());
    let consumer_name = format!("{}-ops-controller", cfg.workspace);
    let stream = js.get_stream(&stream_name).await
        .map_err(|e| format!("get stream {stream_name}: {e}"))?;
    let consumer = stream
        .get_or_create_consumer(
            &consumer_name,
            pull::Config {
                durable_name: Some(consumer_name.clone()),
                filter_subject: format!("ops.cmd.{}.>", cfg.workspace),
                ack_policy: async_nats::jetstream::consumer::AckPolicy::Explicit,
                ..Default::default()
            },
        )
        .await
        .map_err(|e| format!("create consumer {consumer_name}: {e}"))?;

    info!(stream = stream_name, consumer = consumer_name, "ops-controller ready");

    let mut messages = consumer.messages().await
        .map_err(|e| format!("messages: {e}"))?;

    while let Some(msg_result) = messages.next().await {
        let msg = match msg_result {
            Ok(m) => m,
            Err(e) => { error!("NATS error: {e}"); continue; }
        };

        let token = match std::str::from_utf8(&msg.payload) {
            Ok(t) => t.trim().to_string(),
            Err(e) => {
                error!(error = %e, "payload is not UTF-8; acking and skipping");
                let _ = msg.ack().await;
                continue;
            }
        };

        let claims = match jwt_auth.validate(&token) {
            Ok(c) => c,
            Err(e) => {
                error!(error = %e, "JWT validation failed; acking and skipping (bad message)");
                let _ = msg.ack().await;
                continue;
            }
        };

        let jti = &claims.jti;
        let op = &claims.op;

        // Idempotency check
        match check_idempotency(&persistence, jti).await {
            Ok(true) => {
                info!(jti, "duplicate jti — idempotent skip");
                let _ = msg.ack().await;
                continue;
            }
            Ok(false) => {}
            Err(e) => {
                warn!(jti, error = %e, "idempotency check inconclusive; NACKing for retry");
                let _ = msg.ack_with(AckKind::Nak(None)).await;
                continue;
            }
        }

        // Version check
        match persistence.current_version().await {
            Ok(current) if current != claims.expected_state_version => {
                let reason = format!(
                    "state version mismatch: expected '{}', current '{current}'",
                    claims.expected_state_version
                );
                error!(jti, reason, "version mismatch — rejecting op");
                if let Err(e) = emit_ack(&js, &cfg.workspace, op, jti, AckResult::Rejected, Some(reason.clone())).await {
                    warn!(jti, error = %e, "ack emit failed");
                }
                let _ = msg.ack().await;
                continue;
            }
            Err(e) => {
                warn!(jti, error = %e, "version lookup failed; NACKing for retry");
                let _ = msg.ack_with(AckKind::Nak(None)).await;
                continue;
            }
            Ok(_) => {}
        }

        // Persist as Pending before applying (crash-safe)
        if let Err(e) = persistence.insert_pending(jti, &op.op_type, &op.target).await {
            match e {
                ControllerError::DuplicateJti(_) => {
                    info!(jti, "duplicate jti on insert — idempotent skip");
                    let _ = msg.ack().await;
                    continue;
                }
                other => {
                    error!(jti, error = %other, "DB insert failed; NACKing for retry");
                    let _ = msg.ack_with(AckKind::Nak(None)).await;
                    continue;
                }
            }
        }

        // Apply via orchestrator
        match handlers::dispatch(op, &orc).await {
            Ok(()) => {
                info!(jti, op_type = %op.op_type, target = %op.target, "op applied");
                let version = persistence.increment_version().await.unwrap_or_else(|_| "unknown".into());
                if let Err(e) = persistence.mark_applied(jti).await {
                    warn!(jti, error = %e, "mark_applied failed");
                }
                if let Err(e) = emit_audit(&js, &cfg.workspace, op, jti, "applied").await {
                    warn!(jti, error = %e, "audit emit failed (non-fatal)");
                }
                if let Err(e) = emit_ack(&js, &cfg.workspace, op, jti, AckResult::Applied, None).await {
                    warn!(jti, error = %e, "ack emit failed (non-fatal)");
                }
                info!(jti, new_version = version, "state version incremented");
                let _ = msg.ack().await;
            }
            Err(e) => {
                let reason = e.to_string();
                error!(jti, error = reason, "orchestrator apply failed");
                if let Err(db_e) = persistence.mark_rejected(jti, &reason).await {
                    warn!(jti, error = %db_e, "mark_rejected failed");
                }
                if let Err(e) = emit_audit(&js, &cfg.workspace, op, jti, "rejected").await {
                    warn!(jti, error = %e, "audit emit failed");
                }
                if let Err(e) = emit_ack(&js, &cfg.workspace, op, jti, AckResult::Rejected, Some(reason)).await {
                    warn!(jti, error = %e, "ack emit failed");
                }
                let _ = msg.ack().await;
            }
        }
    }

    warn!("NATS message stream ended; ops-controller exiting");
    Ok(())
}
