import at_record_server/catalog_index
import at_record_server/context.{type Context, Context}
import at_record_server/handlers/oauth as oauth_handler
import at_record_server/oauth/assertion
import at_record_server/oauth/authed
import at_record_server/oauth/config
import at_record_server/oauth/dpop
import at_record_server/oauth/keys
import at_record_server/oauth/pkce
import at_record_server/oauth/session_store
import at_record_server/oauth/sessions
import at_record_server/oauth/sessions_memory
import at_record_server/oauth/store
import at_record_server/oauth/tokens
import atproto/xrpc
import gleam/bit_array
import gleam/crypto
import gleam/http
import gleam/http/request
import gleam/http/response
import gleam/list
import gleam/option.{None, Some}
import gleam/result
import gleam/string
import gose
import kryptos/ec
import support
import wisp
import wisp/simulate

fn segment(jwt: String, index: Int) -> String {
  case list.drop(string.split(jwt, "."), index) {
    [seg, ..] ->
      case bit_array.base64_url_decode(seg) {
        Ok(bits) ->
          case bit_array.to_string(bits) {
            Ok(s) -> s
            Error(_) -> ""
          }
        Error(_) -> ""
      }
    _ -> ""
  }
}

pub fn pkce_challenge_is_s256_of_verifier_test() {
  let p = pkce.generate()
  let recomputed =
    bit_array.base64_url_encode(
      crypto.hash(crypto.Sha256, bit_array.from_string(p.verifier)),
      False,
    )
  assert p.challenge == recomputed
}

pub fn dpop_proof_has_dpop_header_and_bound_claims_test() {
  let key = keys.load().key
  let assert Ok(proof) =
    dpop.proof(
      key,
      method: "POST",
      url: "https://as.example/par",
      nonce: None,
      ath: None,
    )

  assert list.length(string.split(proof, ".")) == 3

  let header = segment(proof, 0)
  assert string.contains(header, "\"typ\":\"dpop+jwt\"")
  assert string.contains(header, "\"alg\":\"ES256\"")
  assert string.contains(header, "\"jwk\"")

  let payload = segment(proof, 1)
  assert string.contains(payload, "\"htm\":\"POST\"")
  assert string.contains(payload, "\"htu\":\"https://as.example/par\"")
}

pub fn dpop_proof_includes_nonce_when_given_test() {
  let key = keys.load().key
  let assert Ok(proof) =
    dpop.proof(
      key,
      method: "POST",
      url: "https://as.example/par",
      nonce: Some("nonce-123"),
      ath: None,
    )
  assert string.contains(segment(proof, 1), "\"nonce\":\"nonce-123\"")
}

pub fn client_assertion_is_signed_with_kid_and_audience_test() {
  let ck = keys.load()
  let assert Ok(jwt) =
    assertion.client_assertion(
      ck,
      client_id: "https://app.example/client-metadata.json",
      audience: "https://as.example",
    )
  assert string.contains(segment(jwt, 0), "\"kid\":\"" <> ck.jkt)
  let payload = segment(jwt, 1)
  assert string.contains(
    payload,
    "\"iss\":\"https://app.example/client-metadata.json\"",
  )
  assert string.contains(payload, "\"aud\":\"https://as.example\"")
}

pub fn config_localhost_is_public_client_test() {
  let cfg =
    support.stub_config_with(
      support.unreachable_client(),
      "https://resolver",
      "http://localhost:8080",
    )
  assert cfg.confidential == False
  assert string.starts_with(cfg.client_id, "http://localhost?")
  assert string.contains(cfg.redirect_uri, "127.0.0.1")
}

pub fn config_https_is_confidential_client_test() {
  let cfg =
    support.stub_config_with(
      support.unreachable_client(),
      "https://resolver",
      "https://app.example",
    )
  assert cfg.confidential == True
  assert cfg.client_id == "https://app.example/client-metadata.json"
  assert cfg.redirect_uri == "https://app.example/api/oauth/callback"
}

const token_response = "{\"access_token\":\"at\",\"refresh_token\":\"rt\",\"token_type\":\"DPoP\",\"expires_in\":3600,\"sub\":\"did:plc:abc\",\"scope\":\"atproto\"}"

pub fn token_exchange_parses_dpop_bound_tokens_test() {
  let stub =
    xrpc.Client(send: fn(_req) {
      Ok(response.Response(200, [], bit_array.from_string(token_response)))
    })
  let cfg = support.stub_config_with(stub, "r", "http://localhost:8080")
  let flow =
    store.PendingFlow(
      handle: "h.test",
      pds: "https://pds",
      issuer: "https://as",
      token_endpoint: "https://as/token",
      dpop_key: gose.generate_ec(ec.P256),
      pkce_verifier: "verifier",
      client_id: cfg.client_id,
      confidential: False,
    )
  let assert Ok(tok) = tokens.exchange_code(cfg, flow, "code-123")
  assert tok.access_token == "at"
  assert tok.refresh_token == "rt"
  assert tok.sub == "did:plc:abc"
  assert tok.expires_in == 3600
}

pub fn sessions_create_get_delete_roundtrip_test() {
  let assert Ok(ss) = sessions_memory.start()
  let session =
    support.stub_session_with(
      access_token: "at",
      refresh_token: "rt",
      expires_at: 0,
    )
  let assert Ok(id) = session_store.create(ss, session)
  let assert Some(got) = session_store.get_by_id(ss, id)
  assert got.did == "did:plc:x"
  session_store.delete(ss, id)
  assert session_store.get_by_id(ss, id) == None
}

const far_future = 9_999_999_999

fn test_cfg(stub: xrpc.Client) -> config.Config {
  support.stub_config_with(stub, "r", "http://localhost:8080")
}

fn test_session(
  access_token: String,
  expires_at: Int,
) -> sessions.OauthSession {
  support.stub_session_with(access_token:, refresh_token: "rt", expires_at:)
}

pub fn authed_attaches_dpop_authorization_and_ath_test() {
  let stub =
    xrpc.Client(send: fn(req) {
      let auth =
        request.get_header(req, "authorization") |> result.unwrap("none")
      let dpop_hdr = request.get_header(req, "dpop") |> result.unwrap("none")
      Ok(response.Response(
        200,
        [],
        bit_array.from_string(auth <> "||" <> dpop_hdr),
      ))
    })
  let cfg = test_cfg(stub)
  let assert Ok(#(client, _)) =
    authed.prepare(cfg, "id", test_session("AT", far_future))
  let assert Ok(req) =
    request.to(
      "https://pds.example/xrpc/com.atproto.repo.listRecords?repo=did&collection=x",
    )
  let assert Ok(resp) = xrpc.send_text(client, req)

  let assert [auth, dpop_jwt] = string.split(resp.body, "||")
  assert auth == "DPoP AT"
  assert string.contains(segment(dpop_jwt, 0), "\"typ\":\"dpop+jwt\"")
  let payload = segment(dpop_jwt, 1)
  assert string.contains(payload, "\"htm\":\"GET\"")
  assert string.contains(
    payload,
    "\"htu\":\"https://pds.example/xrpc/com.atproto.repo.listRecords\"",
  )
  assert string.contains(payload, "\"ath\"")
}

pub fn authed_retries_with_server_nonce_test() {
  let stub =
    xrpc.Client(send: fn(req) {
      let dpop_hdr = request.get_header(req, "dpop") |> result.unwrap("")
      case string.contains(segment(dpop_hdr, 1), "\"nonce\"") {
        True -> Ok(response.Response(200, [], bit_array.from_string("ok")))
        False ->
          Ok(response.Response(
            401,
            [#("dpop-nonce", "srv-1"), #("content-type", "application/json")],
            bit_array.from_string("{\"error\":\"use_dpop_nonce\"}"),
          ))
      }
    })
  let cfg = test_cfg(stub)
  let assert Ok(#(client, _)) =
    authed.prepare(cfg, "id", test_session("AT", far_future))
  let assert Ok(req) = request.to("https://pds.example/xrpc/x")
  let assert Ok(resp) = xrpc.send_text(client, req)
  assert resp.status == 200
  assert resp.body == "ok"
}

pub fn prepare_refreshes_expired_session_test() {
  let stub =
    xrpc.Client(send: fn(_req) {
      Ok(response.Response(200, [], bit_array.from_string(token_response)))
    })
  let cfg = test_cfg(stub)
  let assert Ok(#(_, session)) =
    authed.prepare(cfg, "id", test_session("OLD", 0))
  assert session.access_token == "at"
  assert session.refresh_token == "rt"
}

pub fn prepare_keeps_valid_session_test() {
  let stub =
    xrpc.Client(send: fn(_req) {
      Ok(response.Response(
        500,
        [],
        bit_array.from_string("should not be called"),
      ))
    })
  let cfg = test_cfg(stub)
  let assert Ok(#(_, session)) =
    authed.prepare(cfg, "id", test_session("KEEP", far_future))
  assert session.access_token == "KEEP"
}

pub fn authed_ignores_nonce_signal_in_binary_bodies_test() {
  let stub =
    xrpc.Client(send: fn(_req) {
      Ok(response.Response(
        401,
        [#("dpop-nonce", "srv-1"), #("content-type", "image/jpeg")],
        bit_array.from_string("use_dpop_nonce"),
      ))
    })
  let cfg = test_cfg(stub)
  let assert Ok(#(client, _)) =
    authed.prepare(cfg, "id", test_session("AT", far_future))
  let assert Ok(req) = request.to("https://pds.example/xrpc/x")
  // No retry: the signal only counts in JSON bodies or WWW-Authenticate.
  let assert Ok(resp) = xrpc.send_text(client, req)
  assert resp.status == 401
}

// -- browser-facing callback/login error redirects ------------------------
//
// The login and callback routes are full-page browser navigations: on
// failure they must redirect to `/?error=<code>` rather than return a JSON
// body, so the app shell (not raw JSON) is what the user sees.

const state_cookie = "ar_oauth_state"

fn test_context(cfg: config.Config) -> Context {
  support.stub_context(cfg)
}

/// A request carrying the `ar_oauth_state` cookie signed for this process,
/// mirroring what the browser sends back after `login` set it.
fn request_with_state_cookie(path: String, state: String) -> wisp.Request {
  let req = simulate.request(http.Get, path)
  let carrier =
    wisp.ok() |> wisp.set_cookie(req, state_cookie, state, wisp.Signed, 600)
  simulate.session(req, req, carrier)
}

pub fn callback_as_error_passes_code_through_test() {
  let ctx = test_context(test_cfg(support.unreachable_client()))
  let req =
    simulate.request(http.Get, "/api/oauth/callback?error=access_denied")
  let resp = oauth_handler.callback(req, ctx)
  assert resp.status == 303
  assert response.get_header(resp, "location") == Ok("/?error=access_denied")
}

pub fn callback_missing_code_or_state_redirects_bad_callback_test() {
  let ctx = test_context(test_cfg(support.unreachable_client()))
  let req = simulate.request(http.Get, "/api/oauth/callback")
  let resp = oauth_handler.callback(req, ctx)
  assert resp.status == 303
  assert response.get_header(resp, "location") == Ok("/?error=bad_callback")
}

pub fn callback_without_matching_cookie_redirects_state_mismatch_test() {
  let ctx = test_context(test_cfg(support.unreachable_client()))
  let req = simulate.request(http.Get, "/api/oauth/callback?code=abc&state=xyz")
  let resp = oauth_handler.callback(req, ctx)
  assert resp.status == 303
  assert response.get_header(resp, "location") == Ok("/?error=state_mismatch")
}

pub fn callback_token_exchange_failure_redirects_test() {
  let cfg = test_cfg(support.unreachable_client())
  let ctx = test_context(cfg)
  let state = "test-state"
  let issuer = "https://as.example"
  store.put(
    cfg.store,
    state,
    store.PendingFlow(
      handle: "h.test",
      pds: "https://pds.example",
      issuer:,
      token_endpoint: issuer <> "/token",
      dpop_key: gose.generate_ec(ec.P256),
      pkce_verifier: "verifier",
      client_id: cfg.client_id,
      confidential: False,
    ),
  )
  let req =
    request_with_state_cookie(
      "/api/oauth/callback?code=abc&state=" <> state <> "&iss=" <> issuer,
      state,
    )
  let resp = oauth_handler.callback(req, ctx)
  assert resp.status == 303
  assert response.get_header(resp, "location")
    == Ok("/?error=token_exchange_failed")
}

const list_records_response = "{\"records\":[{\"uri\":\"at://did:plc:abc/dev.mokkenstorm.crate.catalog.release/r1\",\"cid\":\"bafycid\",\"value\":{\"createdAt\":\"2024-01-01T00:00:00Z\",\"title\":\"Loveless\"}}]}"

/// A successful login must backfill the just-logged-in user's own existing
/// releases into `catalog_index`: that repo's contents predate this login and
/// Jetstream only tails commits from here forward, so this is the only chance
/// to seed them.
pub fn callback_success_backfills_new_users_releases_into_catalog_index_test() {
  let stub =
    xrpc.Client(send: fn(req) {
      case string.contains(req.path, "listRecords") {
        True ->
          Ok(response.Response(
            200,
            [],
            bit_array.from_string(list_records_response),
          ))
        False ->
          Ok(response.Response(200, [], bit_array.from_string(token_response)))
      }
    })
  let cfg = test_cfg(stub)
  let assert Ok(index) = catalog_index.start()
  let ctx = Context(..test_context(cfg), catalog_index: index)
  let state = "test-state-success"
  let issuer = "https://as.example"
  store.put(
    cfg.store,
    state,
    store.PendingFlow(
      handle: "h.test",
      pds: "https://pds.example",
      issuer:,
      token_endpoint: issuer <> "/token",
      dpop_key: gose.generate_ec(ec.P256),
      pkce_verifier: "verifier",
      client_id: cfg.client_id,
      confidential: False,
    ),
  )
  let req =
    request_with_state_cookie(
      "/api/oauth/callback?code=abc&state=" <> state <> "&iss=" <> issuer,
      state,
    )
  let resp = oauth_handler.callback(req, ctx)
  assert resp.status == 303

  let assert [row] = index.releases.list()
  assert row.title == "Loveless"
  assert row.publisher_did == "did:plc:abc"
}