use super::GatewayState;
use axum::Json;
use axum::body::Body;
use axum::extract::{ConnectInfo, Query, State};
use axum::http::{HeaderMap, Request, StatusCode, header};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use base64::Engine;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;

const CODE_TTL: Duration = Duration::from_secs(5 * 60);
const LOCKOUT_DURATION: Duration = Duration::from_secs(5 * 60);
const MAX_FAILED_ATTEMPTS: u32 = 5;
const MAX_TRACKED_CLIENTS: usize = 4096;
const MAX_PAIRED_TOKENS: usize = 128;
const AUTH_COOKIE: &str = "picobot_auth";
pub const ADMIN_TOKEN_HEADER: &str = "X-Picobot-Admin-Token";

#[derive(Debug, Clone, Serialize, Deserialize)]
struct AuthStore {
    version: u32,
    token_hashes: Vec<String>,
}

impl Default for AuthStore {
    fn default() -> Self {
        Self {
            version: 1,
            token_hashes: Vec::new(),
        }
    }
}

#[derive(Debug, Clone)]
struct PendingCode {
    value: String,
    expires_at: Instant,
}

#[derive(Debug, Clone)]
struct FailedAttempts {
    count: u32,
    locked_until: Option<Instant>,
    last_attempt: Instant,
}

#[derive(Debug)]
struct AuthState {
    token_hashes: HashSet<String>,
    pending_code: Option<PendingCode>,
    failures: HashMap<String, FailedAttempts>,
}

#[derive(Debug, Clone)]
pub struct AuthManager {
    required: bool,
    path: PathBuf,
    admin_token_hash: String,
    state: Arc<Mutex<AuthState>>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AuthIdentity {
    PairingDisabled,
    Paired { token_hash: String },
    LocalAdmin,
}

#[derive(Debug)]
pub enum PairError {
    Invalid,
    Locked(u64),
    Capacity,
    Persistence(std::io::Error),
}

impl AuthManager {
    pub async fn load(required: bool, path: PathBuf) -> Result<Self, Box<dyn std::error::Error>> {
        let store = match tokio::fs::read(&path).await {
            Ok(bytes) => serde_json::from_slice::<AuthStore>(&bytes)
                .map_err(|error| format!("invalid WebUI auth store {}: {error}", path.display()))?,
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => AuthStore::default(),
            Err(error) => return Err(error.into()),
        };
        if store.version != 1 {
            return Err(format!("unsupported WebUI auth store version: {}", store.version).into());
        }
        let admin_token_path = path
            .parent()
            .unwrap_or_else(|| Path::new("."))
            .join("web_admin_token");
        let admin_token = load_or_create_admin_token(&admin_token_path).await?;
        Ok(Self {
            required,
            path,
            admin_token_hash: hash_token(&admin_token),
            state: Arc::new(Mutex::new(AuthState {
                token_hashes: store.token_hashes.into_iter().collect(),
                pending_code: None,
                failures: HashMap::new(),
            })),
        })
    }

    pub fn required(&self) -> bool {
        self.required
    }

    pub async fn authenticate(&self, token: Option<&str>) -> Option<AuthIdentity> {
        if !self.required {
            return Some(AuthIdentity::PairingDisabled);
        }
        let hash = hash_token(token?);
        self.state
            .lock()
            .await
            .token_hashes
            .contains(&hash)
            .then_some(AuthIdentity::Paired { token_hash: hash })
    }

    pub async fn identity_is_active(&self, identity: &AuthIdentity) -> bool {
        match identity {
            AuthIdentity::PairingDisabled => !self.required,
            AuthIdentity::LocalAdmin => true,
            AuthIdentity::Paired { token_hash } => {
                self.required && self.state.lock().await.token_hashes.contains(token_hash)
            }
        }
    }

    pub fn authenticate_admin(&self, token: Option<&str>) -> bool {
        token.is_some_and(|token| constant_time_eq(&hash_token(token), &self.admin_token_hash))
    }

    pub async fn issue_code(&self, revoke_all: bool) -> Result<(String, i64), std::io::Error> {
        let mut state = self.state.lock().await;
        if revoke_all && !state.token_hashes.is_empty() {
            persist_hashes(&self.path, &HashSet::new()).await?;
            state.token_hashes.clear();
        }
        let code = generate_code();
        let expires_at_unix = chrono::Utc::now().timestamp() + CODE_TTL.as_secs() as i64;
        state.pending_code = Some(PendingCode {
            value: code.clone(),
            expires_at: Instant::now() + CODE_TTL,
        });
        Ok((code, expires_at_unix))
    }

    pub async fn try_pair(&self, code: &str, client: &str) -> Result<String, PairError> {
        let now = Instant::now();
        let mut state = self.state.lock().await;
        prune_failures(&mut state.failures, now);
        if let Some(attempts) = state.failures.get(client)
            && let Some(until) = attempts.locked_until
            && now < until
        {
            return Err(PairError::Locked((until - now).as_secs().max(1)));
        }

        let valid = state.pending_code.as_ref().is_some_and(|pending| {
            now < pending.expires_at && constant_time_eq(code.trim(), pending.value.as_str())
        });
        if !valid {
            record_failure(&mut state.failures, client.to_string(), now);
            return Err(PairError::Invalid);
        }
        if state.token_hashes.len() >= MAX_PAIRED_TOKENS {
            return Err(PairError::Capacity);
        }

        let token = generate_token();
        let mut next_hashes = state.token_hashes.clone();
        next_hashes.insert(hash_token(&token));
        persist_hashes(&self.path, &next_hashes)
            .await
            .map_err(PairError::Persistence)?;
        state.token_hashes = next_hashes;
        state.pending_code = None;
        state.failures.remove(client);
        Ok(token)
    }
}

fn token_from_headers(headers: &HeaderMap) -> Option<&str> {
    if let Some(token) = headers
        .get(header::AUTHORIZATION)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.strip_prefix("Bearer "))
        .filter(|value| !value.is_empty())
    {
        return Some(token);
    }
    headers
        .get(header::COOKIE)
        .and_then(|value| value.to_str().ok())
        .and_then(|cookies| {
            cookies.split(';').find_map(|cookie| {
                let (name, value) = cookie.trim().split_once('=')?;
                (name == AUTH_COOKIE && !value.is_empty()).then_some(value)
            })
        })
}

pub async fn require_auth(
    State(auth): State<AuthManager>,
    mut request: Request<Body>,
    next: Next,
) -> Response {
    let mut identity = auth
        .authenticate(token_from_headers(request.headers()))
        .await;
    if identity.is_none()
        && request.uri().path() == "/ws"
        && request
            .extensions()
            .get::<ConnectInfo<SocketAddr>>()
            .is_some_and(|ConnectInfo(peer)| peer.ip().is_loopback())
        && auth.authenticate_admin(
            request
                .headers()
                .get(ADMIN_TOKEN_HEADER)
                .and_then(|value| value.to_str().ok()),
        )
    {
        identity = Some(AuthIdentity::LocalAdmin);
    }
    let Some(identity) = identity else {
        return (
            StatusCode::UNAUTHORIZED,
            Json(serde_json::json!({ "error": "browser is not paired" })),
        )
            .into_response();
    };
    request.extensions_mut().insert(identity);
    next.run(request).await
}

#[derive(Serialize)]
pub struct AuthStatus {
    require_pairing: bool,
    authenticated: bool,
}

pub async fn status(
    State(state): State<Arc<GatewayState>>,
    headers: HeaderMap,
) -> Json<AuthStatus> {
    let authenticated = state
        .auth
        .authenticate(token_from_headers(&headers))
        .await
        .is_some();
    Json(AuthStatus {
        require_pairing: state.auth.required(),
        authenticated,
    })
}

#[derive(Debug, Deserialize)]
pub struct PairRequest {
    code: String,
}

pub async fn pair(
    ConnectInfo(peer): ConnectInfo<SocketAddr>,
    State(state): State<Arc<GatewayState>>,
    Json(request): Json<PairRequest>,
) -> Response {
    if !state.auth.required() {
        return Json(serde_json::json!({ "paired": true })).into_response();
    }
    match state
        .auth
        .try_pair(&request.code, &peer.ip().to_string())
        .await
    {
        Ok(token) => {
            tracing::info!(client_ip = %peer.ip(), "New browser device paired");
            let cookie = format!(
                "{AUTH_COOKIE}={token}; Path=/; HttpOnly; SameSite=Strict; Max-Age=31536000"
            );
            (
                StatusCode::OK,
                [(header::SET_COOKIE, cookie)],
                Json(serde_json::json!({ "paired": true, "token": token })),
            )
                .into_response()
        }
        Err(PairError::Locked(seconds)) => (
            StatusCode::TOO_MANY_REQUESTS,
            Json(serde_json::json!({ "error": "too many attempts", "retry_after": seconds })),
        )
            .into_response(),
        Err(PairError::Capacity) => (
            StatusCode::CONFLICT,
            Json(serde_json::json!({ "error": "paired device limit reached" })),
        )
            .into_response(),
        Err(PairError::Persistence(error)) => {
            tracing::error!(error = %error, "Failed to persist browser pairing");
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({ "error": "pairing could not be persisted" })),
            )
                .into_response()
        }
        Err(PairError::Invalid) => (
            StatusCode::UNAUTHORIZED,
            Json(serde_json::json!({ "error": "invalid or expired pairing code" })),
        )
            .into_response(),
    }
}

#[derive(Debug, Default, Deserialize)]
pub struct CodeQuery {
    #[serde(default)]
    revoke_all: bool,
}

pub async fn issue_code(
    ConnectInfo(peer): ConnectInfo<SocketAddr>,
    State(state): State<Arc<GatewayState>>,
    Query(query): Query<CodeQuery>,
    headers: HeaderMap,
) -> Response {
    let admin_token = headers
        .get(ADMIN_TOKEN_HEADER)
        .and_then(|value| value.to_str().ok());
    if !peer.ip().is_loopback() || !state.auth.authenticate_admin(admin_token) {
        return (
            StatusCode::FORBIDDEN,
            Json(serde_json::json!({ "error": "pairing codes can only be issued locally" })),
        )
            .into_response();
    }
    if !state.auth.required() {
        return (
            StatusCode::CONFLICT,
            Json(serde_json::json!({ "error": "pairing is disabled" })),
        )
            .into_response();
    }
    match state.auth.issue_code(query.revoke_all).await {
        Ok((code, expires_at)) => Json(serde_json::json!({
            "pairing_code": code,
            "expires_at": expires_at,
            "revoke_all": query.revoke_all,
        }))
        .into_response(),
        Err(error) => {
            tracing::error!(error = %error, "Failed to persist pairing token revocation");
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({ "error": "could not update pairing state" })),
            )
                .into_response()
        }
    }
}

fn generate_code() -> String {
    let bytes = *uuid::Uuid::new_v4().as_bytes();
    let value = u64::from_le_bytes(bytes[..8].try_into().expect("eight UUID bytes"));
    format!("{:08}", value % 100_000_000)
}

fn generate_token() -> String {
    let mut bytes = Vec::with_capacity(32);
    bytes.extend_from_slice(uuid::Uuid::new_v4().as_bytes());
    bytes.extend_from_slice(uuid::Uuid::new_v4().as_bytes());
    format!(
        "pb_{}",
        base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
    )
}

fn hash_token(token: &str) -> String {
    let digest = Sha256::digest(token.as_bytes());
    digest.iter().map(|byte| format!("{byte:02x}")).collect()
}

fn constant_time_eq(left: &str, right: &str) -> bool {
    if left.len() != right.len() {
        return false;
    }
    left.bytes()
        .zip(right.bytes())
        .fold(0_u8, |diff, (a, b)| diff | (a ^ b))
        == 0
}

fn record_failure(failures: &mut HashMap<String, FailedAttempts>, client: String, now: Instant) {
    if failures.len() >= MAX_TRACKED_CLIENTS
        && !failures.contains_key(&client)
        && let Some(oldest) = failures
            .iter()
            .min_by_key(|(_, value)| value.last_attempt)
            .map(|(key, _)| key.clone())
    {
        failures.remove(&oldest);
    }
    let attempts = failures.entry(client).or_insert(FailedAttempts {
        count: 0,
        locked_until: None,
        last_attempt: now,
    });
    attempts.count += 1;
    attempts.last_attempt = now;
    if attempts.count >= MAX_FAILED_ATTEMPTS {
        attempts.locked_until = Some(now + LOCKOUT_DURATION);
    }
}

fn prune_failures(failures: &mut HashMap<String, FailedAttempts>, now: Instant) {
    failures.retain(|_, attempts| now.duration_since(attempts.last_attempt) < LOCKOUT_DURATION * 2);
}

async fn persist_hashes(path: &Path, hashes: &HashSet<String>) -> Result<(), std::io::Error> {
    let store = AuthStore {
        version: 1,
        token_hashes: hashes.iter().cloned().collect(),
    };
    let bytes = serde_json::to_vec_pretty(&store)
        .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))?;
    let parent = path.parent().unwrap_or_else(|| Path::new("."));
    tokio::fs::create_dir_all(parent).await?;
    let temporary = parent.join(format!(".web-auth-{}.tmp", uuid::Uuid::new_v4().simple()));
    let result = async {
        tokio::fs::write(&temporary, &bytes).await?;
        #[cfg(unix)]
        tokio::fs::set_permissions(
            &temporary,
            std::os::unix::fs::PermissionsExt::from_mode(0o600),
        )
        .await?;
        tokio::fs::rename(&temporary, path).await
    }
    .await;
    if result.is_err() {
        let _ = tokio::fs::remove_file(&temporary).await;
    }
    result
}

async fn load_or_create_admin_token(path: &Path) -> Result<String, std::io::Error> {
    match tokio::fs::read_to_string(path).await {
        Ok(token) if !token.trim().is_empty() => {
            #[cfg(unix)]
            tokio::fs::set_permissions(path, std::os::unix::fs::PermissionsExt::from_mode(0o600))
                .await?;
            return Ok(token.trim().to_string());
        }
        Ok(_) => {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "WebUI admin token file is empty",
            ));
        }
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
        Err(error) => return Err(error),
    }
    let parent = path.parent().unwrap_or_else(|| Path::new("."));
    tokio::fs::create_dir_all(parent).await?;
    let token = generate_token();
    let temporary = parent.join(format!(".web-admin-{}.tmp", uuid::Uuid::new_v4().simple()));
    let result = async {
        #[cfg(unix)]
        {
            use std::io::Write;
            use std::os::unix::fs::OpenOptionsExt;
            let mut file = std::fs::OpenOptions::new()
                .create_new(true)
                .write(true)
                .mode(0o600)
                .open(&temporary)?;
            file.write_all(token.as_bytes())?;
            file.sync_all()?;
        }
        #[cfg(not(unix))]
        tokio::fs::write(&temporary, &token).await?;
        match tokio::fs::rename(&temporary, path).await {
            Ok(()) => Ok(()),
            Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => Ok(()),
            Err(error) => Err(error),
        }
    }
    .await;
    if let Err(error) = result {
        let _ = tokio::fs::remove_file(&temporary).await;
        return Err(error);
    }
    // Another process may have won the creation race. Always use what is on disk.
    tokio::fs::read_to_string(path)
        .await
        .map(|value| value.trim().to_string())
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{Extension, Router, middleware, routing};
    use tower::ServiceExt;

    #[tokio::test]
    async fn pairing_is_one_time_and_persists_only_hashes() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("auth.json");
        let manager = AuthManager::load(true, path.clone()).await.unwrap();
        let (code, _) = manager.issue_code(false).await.unwrap();

        let token = manager.try_pair(&code, "127.0.0.1").await.unwrap();
        assert!(manager.authenticate(Some(&token)).await.is_some());
        assert!(matches!(
            manager.try_pair(&code, "127.0.0.2").await,
            Err(PairError::Invalid)
        ));
        let stored = tokio::fs::read_to_string(path).await.unwrap();
        assert!(!stored.contains(&token));

        let restored = AuthManager::load(true, dir.path().join("auth.json"))
            .await
            .unwrap();
        assert!(restored.authenticate(Some(&token)).await.is_some());
    }

    #[tokio::test]
    async fn failed_pairing_attempts_are_locked_out() {
        let dir = tempfile::tempdir().unwrap();
        let manager = AuthManager::load(true, dir.path().join("auth.json"))
            .await
            .unwrap();
        manager.issue_code(false).await.unwrap();
        for _ in 0..MAX_FAILED_ATTEMPTS {
            assert!(matches!(
                manager.try_pair("wrong", "client").await,
                Err(PairError::Invalid)
            ));
        }
        assert!(matches!(
            manager.try_pair("wrong", "client").await,
            Err(PairError::Locked(_))
        ));
    }

    #[tokio::test]
    async fn revoke_all_is_durable_and_invalidates_active_identity() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("auth.json");
        let manager = AuthManager::load(true, path.clone()).await.unwrap();
        let (code, _) = manager.issue_code(false).await.unwrap();
        let token = manager.try_pair(&code, "client").await.unwrap();
        let identity = manager.authenticate(Some(&token)).await.unwrap();

        manager.issue_code(true).await.unwrap();
        assert!(!manager.identity_is_active(&identity).await);
        assert!(manager.authenticate(Some(&token)).await.is_none());

        let restored = AuthManager::load(true, path).await.unwrap();
        assert!(restored.authenticate(Some(&token)).await.is_none());
    }

    #[test]
    fn bearer_token_takes_precedence_over_cookie() {
        let mut headers = HeaderMap::new();
        headers.insert(header::COOKIE, "picobot_auth=cookie-token".parse().unwrap());
        headers.insert(
            header::AUTHORIZATION,
            "Bearer header-token".parse().unwrap(),
        );
        assert_eq!(token_from_headers(&headers), Some("header-token"));
    }

    #[tokio::test]
    async fn auth_middleware_rejects_unknown_devices_and_accepts_paired_tokens() {
        let dir = tempfile::tempdir().unwrap();
        let manager = AuthManager::load(true, dir.path().join("auth.json"))
            .await
            .unwrap();
        let app = Router::new()
            .route("/protected", routing::get(|| async { StatusCode::OK }))
            .route_layer(middleware::from_fn_with_state(
                manager.clone(),
                require_auth,
            ));

        let unauthorized = app
            .clone()
            .oneshot(Request::get("/protected").body(Body::empty()).unwrap())
            .await
            .unwrap();
        assert_eq!(unauthorized.status(), StatusCode::UNAUTHORIZED);

        let (code, _) = manager.issue_code(false).await.unwrap();
        let token = manager.try_pair(&code, "client").await.unwrap();
        let authorized = app
            .oneshot(
                Request::get("/protected")
                    .header(header::AUTHORIZATION, format!("Bearer {token}"))
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(authorized.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn local_admin_auth_is_limited_to_loopback_websockets() {
        let dir = tempfile::tempdir().unwrap();
        let manager = AuthManager::load(true, dir.path().join("auth.json"))
            .await
            .unwrap();
        let admin_token = tokio::fs::read_to_string(dir.path().join("web_admin_token"))
            .await
            .unwrap();
        let app = Router::new()
            .route(
                "/ws",
                routing::get(|Extension(identity): Extension<AuthIdentity>| async move {
                    assert_eq!(identity, AuthIdentity::LocalAdmin);
                    StatusCode::OK
                }),
            )
            .route("/protected", routing::get(|| async { StatusCode::OK }))
            .route_layer(middleware::from_fn_with_state(manager, require_auth));

        let mut local_ws = Request::get("/ws")
            .header(ADMIN_TOKEN_HEADER, admin_token.trim())
            .body(Body::empty())
            .unwrap();
        local_ws
            .extensions_mut()
            .insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 42000))));
        assert_eq!(
            app.clone().oneshot(local_ws).await.unwrap().status(),
            StatusCode::OK
        );

        let mut remote_ws = Request::get("/ws")
            .header(ADMIN_TOKEN_HEADER, admin_token.trim())
            .body(Body::empty())
            .unwrap();
        remote_ws
            .extensions_mut()
            .insert(ConnectInfo(SocketAddr::from(([192, 0, 2, 10], 42000))));
        assert_eq!(
            app.clone().oneshot(remote_ws).await.unwrap().status(),
            StatusCode::UNAUTHORIZED
        );

        let mut local_api = Request::get("/protected")
            .header(ADMIN_TOKEN_HEADER, admin_token.trim())
            .body(Body::empty())
            .unwrap();
        local_api
            .extensions_mut()
            .insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 42000))));
        assert_eq!(
            app.oneshot(local_api).await.unwrap().status(),
            StatusCode::UNAUTHORIZED
        );
    }
}
