use crate::{
  auth::{bearer_token::BearerTokenClaims, session::Session},
  error::{Error, Result},
};
use axum::{
  extract::FromRequestParts,
  http::{HeaderMap, StatusCode, request::Parts},
};
use axum_extra::extract::CookieJar;
use chrono::Utc;
use jsonwebtoken::{DecodingKey, Validation, decode};
use sha2::{Digest, Sha256};
use siren_core::data;

pub const COOKIE_NAME: &str = "siren_session";

/// Wraps an optional authenticated session.
///
/// Handlers using this extractor work for both authenticated and
/// unauthenticated callers. A valid `siren_session` cookie grants a
/// `Some(session)`.
pub struct SessionAuthorization(pub Option<Session>);

impl<S> FromRequestParts<S> for SessionAuthorization
where
  S: Send + Sync,
{
  type Rejection = std::convert::Infallible;

  async fn from_request_parts(
    parts: &mut Parts,
    state: &S,
  ) -> std::result::Result<Self, Self::Rejection> {
    let jar = CookieJar::from_request_parts(parts, state).await.unwrap();

    if let Some(cookie) = jar.get(COOKIE_NAME) {
      let ip = extract_ip(&parts.headers);
      let user_agent = parts
        .headers
        .get("user-agent")
        .and_then(|h| h.to_str().ok())
        .unwrap_or("unknown")
        .to_string();

      if let Ok(session) = check_cookie_auth(cookie.value(), &ip, &user_agent).await {
        parts.extensions.insert(session.clone());
        return Ok(Self(Some(session)));
      }
    }

    Ok(Self(None))
  }
}

/// Extract the client IP from common proxy headers, falling back to "unknown".
pub fn extract_ip(headers: &HeaderMap) -> String {
  if let Some(forwarded) = headers.get("x-forwarded-for") {
    if let Ok(val) = forwarded.to_str() {
      if let Some(ip) = val.split(',').next() {
        return ip.trim().to_string();
      }
    }
  }
  if let Some(real_ip) = headers.get("x-real-ip") {
    if let Ok(val) = real_ip.to_str() {
      return val.trim().to_string();
    }
  }
  "unknown".to_string()
}

/// Compute a fingerprint from client IP and User-Agent
///
/// Stored in the Redis session at login time and re-checked on every
/// authenticated request so that a stolen cookie is detected when used
/// from a different device or IP.
pub fn compute_fingerprint(ip: &str, user_agent: &str) -> String {
  let mut hasher = Sha256::new();
  hasher.update(format!("{ip}:{user_agent}"));
  format!("{:x}", hasher.finalize())
}

/// Validate a JWT cookie value, look up the Redis session, and verify the
/// fingerprint against the current request's IP / User-Agent.
pub async fn check_cookie_auth(token: &str, ip: &str, user_agent: &str) -> Result<Session> {
  let jwt_secret = std::env::var("JWT_SECRET").expect("JWT_SECRET must be set");
  let decoding_key = DecodingKey::from_secret(jwt_secret.as_bytes());

  let token_data = decode::<BearerTokenClaims>(token, &decoding_key, &Validation::default())
    .map_err(|_| StatusCode::UNAUTHORIZED)?;

  let claims = token_data.claims;

  if claims.exp < Utc::now().timestamp() {
    return Err(StatusCode::UNAUTHORIZED.into());
  }

  let session = match Session::find(&claims.jti).await? {
    Some(s) => s,
    None => return Err(StatusCode::UNAUTHORIZED.into()),
  };

  // Reject if the request comes from a different device / network
  let expected = compute_fingerprint(ip, user_agent);
  if session.fingerprint != expected {
    log::warn!(
      "Fingerprint mismatch for session {}: stored={} request={}",
      claims.jti,
      session.fingerprint,
      expected
    );
    return Err(StatusCode::UNAUTHORIZED.into());
  }

  Ok(session)
}

/// Parse the raw `Cookie:` header string and validate the siren_session
/// value.  Used by the WebSocket upgrade handler where we cannot use the
/// normal `FromRequestParts` machinery.
pub async fn check_cookie_from_header_str(
  cookie_header: &str,
  ip: &str,
  user_agent: &str,
) -> Option<Session> {
  for pair in cookie_header.split(';') {
    let pair = pair.trim();
    if let Some(value) = pair.strip_prefix("siren_session=") {
      return check_cookie_auth(value, ip, user_agent).await.ok();
    }
  }
  None
}

/// Extractor that requires the caller to be an authenticated site admin.
///
/// Returns `401 Unauthorized` if there is no valid session, or
/// `403 Forbidden` if the user's role is not `"admin"`.
/// On success, the inner `Session` is available to the handler.
pub struct AdminAuthorization(pub Session);

impl<S> FromRequestParts<S> for AdminAuthorization
where
  S: Send + Sync,
{
  type Rejection = Error;

  async fn from_request_parts(
    parts: &mut Parts,
    state: &S,
  ) -> std::result::Result<Self, Self::Rejection> {
    let SessionAuthorization(maybe_session) =
      SessionAuthorization::from_request_parts(parts, state)
        .await
        .unwrap();

    let session = maybe_session.ok_or_else(|| Error::from(StatusCode::UNAUTHORIZED))?;

    // Verify admin role in the database
    let pool = data::pool();
    let role: Option<String> = sqlx::query_scalar("SELECT role FROM users WHERE id = $1")
      .bind(session.user_id)
      .fetch_optional(pool)
      .await
      .map_err(|e| {
        log::error!("DB error checking admin role: {e}");
        Error::from(StatusCode::INTERNAL_SERVER_ERROR)
      })?;

    match role.as_deref() {
      Some("admin") => Ok(AdminAuthorization(session)),
      Some(_) => Err(Error::from(StatusCode::FORBIDDEN)),
      None => Err(Error::from(StatusCode::UNAUTHORIZED)),
    }
  }
}
