use std::time::Duration;

use atproto_identity::key::{KeyData, KeyType, to_public};
use atproto_identity::model::Document;
use atproto_identity::traits::IdentityResolver;
use atproto_xrpcs::{ServiceAuthPolicy, mint_service_auth, verify_service_auth};
use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
use tower::ServiceExt;
use zkwid_h::{Host, run_stage, seed, stage_bytes};

const DID: &str = "did:web:h.zkwid.test";

#[test]
fn stages_build_deterministically() {
    for name in ["shim", "stage2"] {
        let a = stage_bytes(name).unwrap();
        let b = stage_bytes(name).unwrap();
        assert!(a.starts_with(b"\0asm"));
        assert_eq!(a, b);
    }
}

#[test]
fn shim_derives_public_key() {
    let shim = stage_bytes("shim").unwrap();
    let seed = seed();
    let key = run_stage(&shim, &seed).unwrap();
    let expected = to_public(&KeyData::new(KeyType::K256Private, seed.to_vec())).unwrap();
    assert_eq!(key.len(), 33);
    assert_eq!(key, expected.bytes());
}

#[test]
fn shim_traps_on_bad_input() {
    let shim = stage_bytes("shim").unwrap();
    assert!(run_stage(&shim, &[1; 31]).is_err());
    assert!(run_stage(&shim, &[0; 32]).is_err());
}

#[test]
fn stage2_emits_demo_log() {
    let stage2 = stage_bytes("stage2").unwrap();
    let out = run_stage(
        &stage2,
        br#"{"instance":"run-1","createdAt":"2026-11-05T00:00:00.000Z"}"#,
    )
    .unwrap();
    assert_eq!(
        String::from_utf8(out).unwrap(),
        r#"{"$type":"at.zkwid.demo.log","text":"hello from run-1","createdAt":"2026-11-05T00:00:00.000Z"}"#
    );
    let out = run_stage(&stage2, br#"{ "createdAt" : "t", "instance" : "i" }"#).unwrap();
    assert_eq!(
        out,
        br#"{"$type":"at.zkwid.demo.log","text":"hello from i","createdAt":"t"}"#
    );
    assert!(run_stage(&stage2, br#"{"instance":"i"}"#).is_err());
    assert!(run_stage(&stage2, &vec![b' '; 4097]).is_err());
}

struct Fixed(Document);

#[async_trait::async_trait]
impl IdentityResolver for Fixed {
    async fn resolve(&self, subject: &str) -> anyhow::Result<Document> {
        anyhow::ensure!(subject == self.0.id, "unknown {subject}");
        Ok(self.0.clone())
    }
}

#[tokio::test]
async fn did_document_serves_host_key() {
    let host = Host::new(
        DID,
        "https://h.zkwid.test",
        "https://zkwid.test",
        "did:web:zkwid.test",
        reqwest::Client::new(),
    )
    .unwrap();
    let response = host
        .router()
        .oneshot(
            Request::get("/.well-known/did.json")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    let body = response.into_body().collect().await.unwrap().to_bytes();
    let doc: Document = serde_json::from_slice(&body).unwrap();
    assert_eq!(doc.id, DID);
    let public = to_public(&host.signing_key).unwrap().to_string();
    assert_eq!(
        doc.verification_method_multibase("atproto"),
        public.strip_prefix("did:key:")
    );

    let aud = "did:web:zkwid.test#zkwid";
    let token = mint_service_auth(
        &host.signing_key,
        DID,
        aud,
        "at.zkwid.createKey",
        Duration::from_secs(60),
    )
    .unwrap();
    let claims = verify_service_auth(
        &token,
        &Fixed(doc),
        &ServiceAuthPolicy::new("at.zkwid.createKey", aud),
        None,
    )
    .await
    .unwrap();
    assert_eq!(claims.iss, DID);
}