use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{Context, Result};
use axum::{
    Router,
    extract::{Path as AxumPath, State},
    http::{HeaderValue, StatusCode, header::CONTENT_TYPE},
    response::{IntoResponse, Response},
    routing::get,
};
use clap::Parser;
use serde::{Deserialize, Serialize};

mod db;

// Vendored from plox/src/opcodec.rs; keep the storage format implementation in sync.
#[allow(dead_code)]
mod opcodec;

#[derive(Parser)]
struct Args {
    #[arg(long, env = "PLC_MIRROR_DB", default_value = "../plox/data/plox.db")]
    db: PathBuf,

    #[arg(long, env = "PLC_MIRROR_LISTEN", default_value = "127.0.0.1:2486")]
    listen: SocketAddr,

    #[arg(
        long,
        env = "PLC_MIRROR_DB_POOL_SIZE",
        default_value_t = 16,
        value_parser = clap::value_parser!(u32).range(1..)
    )]
    db_pool_size: u32,
}

#[derive(Deserialize)]
#[serde(tag = "type")]
enum Operation {
    #[serde(rename = "plc_operation")]
    Regular(RegularOperation),
    #[serde(rename = "create")]
    Legacy(LegacyOperation),
    #[serde(rename = "plc_tombstone")]
    Tombstone,
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct RegularOperation {
    verification_methods: indexmap::IndexMap<String, String>,
    also_known_as: Vec<String>,
    services: indexmap::IndexMap<String, OperationService>,
}

#[derive(Deserialize)]
struct OperationService {
    #[serde(rename = "type")]
    kind: String,
    endpoint: String,
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct LegacyOperation {
    signing_key: String,
    handle: String,
    service: String,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct DidDocument {
    #[serde(rename = "@context")]
    context: Vec<String>,
    id: String,
    also_known_as: Vec<String>,
    verification_method: Vec<VerificationMethod>,
    service: Vec<Service>,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct VerificationMethod {
    id: String,
    #[serde(rename = "type")]
    kind: &'static str,
    controller: String,
    public_key_multibase: String,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct Service {
    id: String,
    #[serde(rename = "type")]
    kind: String,
    service_endpoint: String,
}

enum Resolution {
    Document(DidDocument),
    Tombstone,
}

fn resolution_from_operation(did: &str, operation: serde_json::Value) -> Result<Resolution> {
    let operation: Operation =
        serde_json::from_value(operation).context("parsing PLC operation")?;
    // legacy create ops resolve through the same path as regular ops once
    // normalized, mirroring plc.directory's normalizeOp
    let (verification_methods, also_known_as, services) = match operation {
        Operation::Regular(op) => (op.verification_methods, op.also_known_as, op.services),
        Operation::Legacy(op) => (
            [("atproto".to_string(), op.signing_key)]
                .into_iter()
                .collect(),
            vec![ensure_atproto_prefix(&op.handle)],
            [(
                "atproto_pds".to_string(),
                OperationService {
                    kind: "AtprotoPersonalDataServer".into(),
                    endpoint: ensure_http_prefix(&op.service),
                },
            )]
            .into_iter()
            .collect(),
        ),
        Operation::Tombstone => return Ok(Resolution::Tombstone),
    };

    let mut context = vec![
        "https://www.w3.org/ns/did/v1".to_string(),
        "https://w3id.org/security/multikey/v1".to_string(),
    ];
    let verification_method = verification_methods
        .into_iter()
        .map(|(name, key)| {
            let multibase = key.strip_prefix("did:key:").unwrap_or(&key);
            if let Some(suite) = suite_context(multibase)
                && !context.iter().any(|existing| existing == suite)
            {
                context.push(suite.into());
            }
            VerificationMethod {
                id: format!("{did}#{name}"),
                kind: "Multikey",
                controller: did.into(),
                public_key_multibase: multibase.into(),
            }
        })
        .collect();
    let service = services
        .into_iter()
        .map(|(name, service)| Service {
            id: format!("#{name}"),
            kind: service.kind,
            service_endpoint: service.endpoint,
        })
        .collect();

    Ok(Resolution::Document(DidDocument {
        context,
        id: did.into(),
        also_known_as,
        verification_method,
        service,
    }))
}

/// Suite context for a verification key, mirroring plc.directory's
/// formatKeyAndContext: the key type comes from the did:key multibase
/// prefix (zDnae = P-256, zQ3sh = secp256k1); other types get none.
fn suite_context(public_key_multibase: &str) -> Option<&'static str> {
    if public_key_multibase.starts_with("zDnae") {
        Some("https://w3id.org/security/suites/ecdsa-2019/v1")
    } else if public_key_multibase.starts_with("zQ3sh") {
        Some("https://w3id.org/security/suites/secp256k1-2019/v1")
    } else {
        None
    }
}

fn ensure_http_prefix(s: &str) -> String {
    if s.starts_with("http://") || s.starts_with("https://") {
        s.to_string()
    } else {
        format!("https://{s}")
    }
}

fn ensure_atproto_prefix(s: &str) -> String {
    if s.starts_with("at://") {
        return s.to_string();
    }
    let stripped = s
        .strip_prefix("https://")
        .or_else(|| s.strip_prefix("http://"))
        .unwrap_or(s);
    format!("at://{stripped}")
}

async fn get_did(State(db_pool): State<Arc<db::Db>>, AxumPath(did): AxumPath<String>) -> Response {
    if !did.starts_with("did:plc:") {
        return StatusCode::NOT_FOUND.into_response();
    }

    match db::resolve(&db_pool, &did).await {
        Ok(Some(Resolution::Document(document))) => {
            // plc.directory serves did documents as did+ld+json
            let mut response = axum::Json(document).into_response();
            response.headers_mut().insert(
                CONTENT_TYPE,
                HeaderValue::from_static("application/did+ld+json"),
            );
            response
        }
        // plc.directory answers tombstoned dids with 404 ("DID not available")
        Ok(Some(Resolution::Tombstone) | None) => StatusCode::NOT_FOUND.into_response(),
        Err(error) => {
            tracing::error!(%error, "failed to resolve DID");
            StatusCode::INTERNAL_SERVER_ERROR.into_response()
        }
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "plc_mirror=info".into()),
        )
        .init();

    let args = Args::parse();
    tracing::info!(db = %args.db.display(), "opening plox database read-only");
    let db_pool = db::open_database(&args.db, args.db_pool_size as usize)?;

    let app = Router::new()
        .route(
            "/",
            get(|| async {
                "plc-mirror: a read-only mirror of plc.directory.\n\
                 \n\
                 GET /{did}  resolve a did:plc DID to its DID document\n"
            }),
        )
        .route("/{did}", get(get_did))
        .with_state(db_pool);
    let listener = tokio::net::TcpListener::bind(args.listen).await?;
    tracing::info!(listen = %args.listen, "serving PLC DID documents");
    axum::serve(listener, app)
        .with_graceful_shutdown(async {
            let _ = tokio::signal::ctrl_c().await;
        })
        .await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db::open_database;
    use rusqlite::{Connection, params};

    const REGULAR_OPERATION: &str = r#"{
        "type": "plc_operation",
        "rotationKeys": ["did:key:zRotation"],
        "verificationMethods": {
            "atproto": "did:key:zQ3shSigning",
            "backup": "did:key:zQ3shBackup"
        },
        "alsoKnownAs": ["at://alice.example"],
        "services": {
            "atproto_pds": {
                "type": "AtprotoPersonalDataServer",
                "endpoint": "https://pds.example"
            }
        },
        "prev": null
    }"#;

    #[test]
    fn converts_legacy_genesis() {
        let Resolution::Document(document) = resolution_from_operation(
            "did:plc:example",
            serde_json::from_str(
                r#"{
                "type": "create",
                "signingKey": "did:key:zDnaeSigning",
                "recoveryKey": "did:key:zRecovery",
                "handle": "alice.example",
                "service": "pds.example"
            }"#,
            )
            .unwrap(),
        )
        .unwrap() else {
            panic!("legacy genesis should produce a document")
        };

        assert_eq!(document.also_known_as, ["at://alice.example"]);
        assert_eq!(
            document.context,
            [
                "https://www.w3.org/ns/did/v1",
                "https://w3id.org/security/multikey/v1",
                "https://w3id.org/security/suites/ecdsa-2019/v1"
            ]
        );
        assert_eq!(
            document.verification_method[0].id,
            "did:plc:example#atproto"
        );
        assert_eq!(
            document.verification_method[0].public_key_multibase,
            "zDnaeSigning"
        );
        assert_eq!(document.service[0].id, "#atproto_pds");
        assert_eq!(document.service[0].service_endpoint, "https://pds.example");
    }

    #[test]
    fn serializes_regular_operation_as_did_document() {
        let Resolution::Document(document) = resolution_from_operation(
            "did:plc:example",
            serde_json::from_str(REGULAR_OPERATION).unwrap(),
        )
        .unwrap() else {
            panic!("regular operation should produce a document")
        };

        assert_eq!(
            serde_json::to_value(document).unwrap(),
            serde_json::json!({
                "@context": [
                    "https://www.w3.org/ns/did/v1",
                    "https://w3id.org/security/multikey/v1",
                    "https://w3id.org/security/suites/secp256k1-2019/v1"
                ],
                "id": "did:plc:example",
                "alsoKnownAs": ["at://alice.example"],
                "verificationMethod": [
                    {
                        "id": "did:plc:example#atproto",
                        "type": "Multikey",
                        "controller": "did:plc:example",
                        "publicKeyMultibase": "zQ3shSigning"
                    },
                    {
                        "id": "did:plc:example#backup",
                        "type": "Multikey",
                        "controller": "did:plc:example",
                        "publicKeyMultibase": "zQ3shBackup"
                    }
                ],
                "service": [{
                    "id": "#atproto_pds",
                    "type": "AtprotoPersonalDataServer",
                    "serviceEndpoint": "https://pds.example"
                }]
            })
        );
    }

    #[test]
    fn recognizes_tombstones() {
        assert!(matches!(
            resolution_from_operation(
                "did:plc:example",
                serde_json::json!({"type":"plc_tombstone","prev":"previous","sig":"signature"})
            ),
            Ok(Resolution::Tombstone)
        ));
    }

    fn database() -> (tempfile::NamedTempFile, Connection) {
        let database = tempfile::NamedTempFile::new().unwrap();
        let connection = Connection::open(database.path()).unwrap();
        connection
            .execute_batch(
                "CREATE TABLE ops (
                    seq INTEGER PRIMARY KEY CHECK (seq > 0),
                    did BLOB NOT NULL,
                    operation BLOB NOT NULL,
                    created_at BLOB NOT NULL DEFAULT X'00'
                ) STRICT;
                CREATE INDEX ops_did ON ops(did);
                CREATE TABLE strings (id INTEGER PRIMARY KEY, value TEXT NOT NULL UNIQUE) STRICT;
                INSERT INTO strings VALUES (0, 'https://pds.example');
                PRAGMA user_version = 1;",
            )
            .unwrap();
        (database, connection)
    }

    #[tokio::test]
    async fn serves_the_latest_database_operation() {
        let (database, connection) = database();
        connection
            .execute(
                "INSERT INTO ops (seq, did, operation) VALUES (?1, ?2, ?3)",
                params![
                    1,
                    opcodec::encode_did("did:plc:ragtjsm2j2vknwkz3zp4oxrd"),
                    opcodec::encode(&serde_json::json!({"type":"plc_tombstone"}), |_| None)
                        .unwrap()
                ],
            )
            .unwrap();
        connection
            .execute(
                "INSERT INTO ops (seq, did, operation) VALUES (?1, ?2, ?3)",
                params![
                    2,
                    opcodec::encode_did("did:plc:ragtjsm2j2vknwkz3zp4oxrd"),
                    opcodec::encode(&serde_json::from_str(REGULAR_OPERATION).unwrap(), |s| (s
                        == "https://pds.example")
                        .then_some(0))
                    .unwrap()
                ],
            )
            .unwrap();

        let pool = open_database(database.path(), 1).unwrap();
        let response = get_did(
            State(pool.clone()),
            AxumPath("did:plc:ragtjsm2j2vknwkz3zp4oxrd".into()),
        )
        .await;
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response.headers().get(CONTENT_TYPE).unwrap(),
            "application/did+ld+json"
        );
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let document: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(
            document["alsoKnownAs"],
            serde_json::json!(["at://alice.example"])
        );

        assert_eq!(
            get_did(State(pool.clone()), AxumPath("did:plc:missing".into()))
                .await
                .status(),
            StatusCode::NOT_FOUND
        );
        assert_eq!(
            get_did(State(pool), AxumPath("did:web:example.com".into()))
                .await
                .status(),
            StatusCode::NOT_FOUND
        );
    }

    #[tokio::test]
    async fn reads_dictionary_additions_while_serving() {
        let (database, mut connection) = database();
        let pool = open_database(database.path(), 2).unwrap();
        let did = "did:plc:ragtjsm2j2vknwkz3zp4oxrd";
        for (id, endpoint) in [(0, "https://pds.example"), (1, "https://new-pds.example")] {
            let tx = connection.transaction().unwrap();
            tx.execute(
                "INSERT OR IGNORE INTO strings VALUES (?1, ?2)",
                params![id, endpoint],
            )
            .unwrap();
            // Version 1 fixture: interned endpoint, handle, literal key, signature, Bluesky genesis.
            let mut program = b"\x01\x00\x0a\x05alice\x00\x14did:key:zQ3shSigning\x03".to_vec();
            program[1] = id as u8;
            program.extend([0; 64]);
            program.push(7);
            tx.execute(
                "INSERT INTO ops (seq, did, operation) VALUES (?1, ?2, ?3)",
                params![id + 1, opcodec::encode_did(did), program],
            )
            .unwrap();
            tx.commit().unwrap();

            let response = get_did(State(pool.clone()), AxumPath(did.into())).await;
            assert_eq!(response.status(), StatusCode::OK);
            let body = axum::body::to_bytes(response.into_body(), usize::MAX)
                .await
                .unwrap();
            let document: serde_json::Value = serde_json::from_slice(&body).unwrap();
            assert_eq!(document["id"], did);
            assert_eq!(
                document["alsoKnownAs"],
                serde_json::json!(["at://alice.bsky.social"])
            );
            assert_eq!(
                document["verificationMethod"][0]["publicKeyMultibase"],
                "zQ3shSigning"
            );
            assert_eq!(document["service"][0]["serviceEndpoint"], endpoint);
        }
    }

    #[tokio::test]
    async fn serves_compact_bluesky_create() {
        let (database, connection) = database();
        let did = "did:plc:ragtjsm2j2vknwkz3zp4oxrd";
        // Fixed BSKY_CREATE fixture, independent of the encoder.
        let mut program =
            b"\x01\x00\x00\x11alice.bsky.social\x00\x14did:key:zQ3shSigning\x03".to_vec();
        program.extend([0; 64]);
        program.push(11);
        connection
            .execute(
                "INSERT INTO ops (seq, did, operation) VALUES (1, ?1, ?2)",
                params![opcodec::encode_did(did), program],
            )
            .unwrap();

        let pool = open_database(database.path(), 1).unwrap();
        let response = get_did(State(pool), AxumPath(did.into())).await;
        assert_eq!(response.status(), StatusCode::OK);
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let document: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(
            document,
            serde_json::json!({
                "@context": [
                    "https://www.w3.org/ns/did/v1",
                    "https://w3id.org/security/multikey/v1",
                    "https://w3id.org/security/suites/secp256k1-2019/v1"
                ],
                "id": did,
                "alsoKnownAs": ["at://alice.bsky.social"],
                "verificationMethod": [{
                    "id": format!("{did}#atproto"),
                    "type": "Multikey",
                    "controller": did,
                    "publicKeyMultibase": "zQ3shSigning"
                }],
                "service": [{
                    "id": "#atproto_pds",
                    "type": "AtprotoPersonalDataServer",
                    "serviceEndpoint": "https://pds.example"
                }]
            })
        );
    }

    #[tokio::test]
    async fn serves_legacy_operations_tombstones_and_decode_errors() {
        let (database, connection) = database();
        let pool = open_database(database.path(), 1).unwrap();
        let legacy = serde_json::json!({
            "type": "create", "signingKey": "did:key:zDnaeSigning",
            "handle": "alice.example", "service": "https://pds.example"
        });
        for (program, status) in [
            (
                opcodec::encode(&legacy, |s| (s == "https://pds.example").then_some(0)).unwrap(),
                StatusCode::OK,
            ),
            (
                opcodec::encode(&serde_json::json!({"type": "plc_tombstone"}), |_| None).unwrap(),
                StatusCode::NOT_FOUND,
            ),
            (vec![255], StatusCode::INTERNAL_SERVER_ERROR),
            (vec![1, 99], StatusCode::INTERNAL_SERVER_ERROR),
        ] {
            connection
                .execute(
                    "INSERT OR REPLACE INTO ops (seq, did, operation) VALUES (1, ?1, ?2)",
                    params![b"\x00did:plc:example".as_slice(), program],
                )
                .unwrap();
            assert_eq!(
                get_did(State(pool.clone()), AxumPath("did:plc:example".into()))
                    .await
                    .status(),
                status,
            );
        }
    }

    #[test]
    fn rejects_unsupported_or_unfinished_databases() {
        let (database, connection) = database();
        for version in [0, 2] {
            connection
                .pragma_update(None, "user_version", version)
                .unwrap();
            let error = open_database(database.path(), 1).err().unwrap();
            assert!(
                error
                    .to_string()
                    .contains("unsupported plox database version")
            );
        }
        connection
            .execute_batch(
                "PRAGMA user_version = 1;
             CREATE TABLE compaction (singleton INTEGER PRIMARY KEY, complete INTEGER NOT NULL);
             INSERT INTO compaction VALUES (1, 0);",
            )
            .unwrap();
        let error = open_database(database.path(), 1).err().unwrap();
        assert!(error.to_string().contains("compaction is unfinished"));
        connection
            .execute("UPDATE compaction SET complete = 1", [])
            .unwrap();
        assert!(open_database(database.path(), 1).is_ok());

        let missing = database.path().with_extension("missing");
        assert!(open_database(&missing, 1).is_err());
        assert!(!missing.exists());
    }
}