// SPDX-FileCopyrightText: 2025 Jonathan D.A. Jewell <j.d.a.jewell@open.ac.uk>
// SPDX-License-Identifier: MPL-2.0
// ECHIDNA REST API Server

use axum::{
    routing::{delete, get, post},
    Router,
};
use echidna::core::ProofState;
use echidna::provers::ProverBackend;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{Mutex, RwLock};
use tower_http::cors::CorsLayer;
use utoipa::OpenApi;
use utoipa_redoc::{Redoc, Servable};

mod ffi_wrapper;
mod handlers;
mod models;

use models::*;

/// Application state shared across handlers
#[derive(Clone)]
pub struct AppState {
    /// Active proof sessions (proof_id -> ProofSession)
    pub sessions: Arc<RwLock<HashMap<String, Arc<Mutex<ProofSession>>>>>,
    /// HTTP client for Julia ML API
    pub ml_client: reqwest::Client,
    /// Julia ML API base URL
    pub ml_api_url: String,
    /// Whether FFI layer is initialized
    pub ffi_initialized: bool,
    /// LLM advisor — drives /api/v1/consult, talks to BoJ's
    /// model-router-mcp cartridge for cost-aware model selection.
    /// Mutex-wrapped because check_health takes &mut self.
    pub llm_advisor: Arc<Mutex<echidna::llm::LlmAdvisor>>,
}

/// A proof session
pub struct ProofSession {
    pub id: String,
    pub prover_kind: echidna::ProverKind,
    pub prover: Box<dyn ProverBackend>,
    pub state: Option<ProofState>,
    pub goal: String,
    pub status: ProofStatus,
    pub history: Vec<String>,
    pub start_time: std::time::Instant,
}

/// Runs the REST API with permissive CORS, ReDoc at `/redoc` and the OpenAPI
/// JSON specification at `/api-docs/openapi.json`.
///
/// Binds to `ECHIDNA_REST_ADDR`, defaulting to `127.0.0.1:8000` when the variable
/// is absent or is not valid Unicode. Probes Julia ML at `http://127.0.0.1:8090`
/// and initialises the FFI layer; failures of either check do not prevent startup.
///
/// # Panics
///
/// Panics if tracing initialisation, HTTP client creation, listener binding or
/// serving fails.
#[tokio::main]
async fn main() {
    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .init();

    let ml_client = reqwest::Client::builder()
        .timeout(Duration::from_secs(10))
        .build()
        .expect("Failed to create HTTP client");

    let ml_api_url = "http://127.0.0.1:8090".to_string();

    // Test Julia ML connection
    match ml_client.get(format!("{}/health", ml_api_url)).send().await {
        Ok(resp) if resp.status().is_success() => {
            tracing::info!("Connected to Julia ML API at {}", ml_api_url);
        },
        _ => {
            tracing::info!(
                "Julia ML API not available at {} (will use fallback)",
                ml_api_url
            );
        },
    }

    // Initialize FFI layer
    let ffi_initialized = ffi_wrapper::init_ffi().is_ok();
    if ffi_initialized {
        tracing::info!("FFI layer initialized successfully");
    } else {
        tracing::warn!("FFI layer initialization failed, falling back to direct Rust calls");
    }

    let state = AppState {
        sessions: Arc::new(RwLock::new(HashMap::new())),
        ml_client,
        ml_api_url,
        ffi_initialized,
        llm_advisor: Arc::new(Mutex::new(echidna::llm::LlmAdvisor::new())),
    };

    let app = Router::new()
        // Health check
        .route("/health", get(health_check))
        // Prover endpoints
        .route("/api/v1/provers", get(handlers::list_provers))
        .route("/api/v1/provers/{kind}", get(handlers::get_prover))
        // Proof endpoints
        .route("/api/v1/proofs", post(handlers::submit_proof))
        .route("/api/v1/proofs", get(handlers::list_proofs))
        .route("/api/v1/proofs/{id}", get(handlers::get_proof))
        .route("/api/v1/proofs/{id}", delete(handlers::cancel_proof))
        // Tactic endpoints
        .route("/api/v1/proofs/{id}/tactics", post(handlers::apply_tactic))
        .route(
            "/api/v1/proofs/{id}/tactics/suggest",
            get(handlers::suggest_tactics),
        )
        // Cross-prover exchange endpoints (OpenTheory / Dedukti).
        // Export is session-scoped; import is stateless.
        .route("/api/v1/proofs/{id}/export", get(handlers::export_proof))
        .route("/api/v1/exchange/import", post(handlers::import_proof))
        // Consultant-mode Q&A — free-form question + optional context →
        // LLM-shaped markdown answer via BoJ's cartridge router.
        .route("/api/v1/consult", post(handlers::consult))
        // OpenAPI documentation: the raw spec, and ReDoc rendering it.
        // ReDoc replaced Swagger UI because utoipa-swagger-ui builds
        // against zip ^8.6, which has no fix for RUSTSEC-2026-0367..0369.
        .route(
            "/api-docs/openapi.json",
            get(|| async { axum::Json(ApiDoc::openapi()) }),
        )
        .merge(Redoc::with_url("/redoc", ApiDoc::openapi()))
        .layer(CorsLayer::permissive())
        .with_state(state);

    let addr = std::env::var("ECHIDNA_REST_ADDR").unwrap_or_else(|_| "127.0.0.1:8000".to_string());
    tracing::info!("REST API listening on http://{}", addr);
    tracing::info!("API docs (ReDoc) available at http://{}/redoc", addr);

    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

async fn health_check() -> &'static str {
    "OK"
}

#[derive(OpenApi)]
#[openapi(
    paths(
        handlers::list_provers,
        handlers::get_prover,
        handlers::submit_proof,
        handlers::list_proofs,
        handlers::get_proof,
        handlers::cancel_proof,
        handlers::apply_tactic,
        handlers::suggest_tactics,
        handlers::export_proof,
        handlers::import_proof,
    ),
    components(
        schemas(ProverKind, ProverInfo, ProofStatus, ProofRequest, ProofResponse, TacticRequest, TacticResponse, ExchangeFormat, ExportResponse, ImportRequest, ImportResponse)
    ),
    tags(
        (name = "provers", description = "Theorem prover management"),
        (name = "proofs", description = "Proof submission and management"),
        (name = "tactics", description = "Tactic application and suggestions"),
        (name = "exchange", description = "Cross-prover proof exchange (OpenTheory, Dedukti)")
    )
)]
struct ApiDoc;