use std::collections::HashMap;
use std::io::{self, Write};
use std::sync::Arc;

use anyhow::{Context, Result, bail};
use clap::Parser;
use reqwest::header::HeaderMap;
use serde::{Deserialize, Serialize};

use atproto_client::client::{AppPasswordAuth, post_apppassword_json};
use atproto_client::com::atproto::server::create_session;
use atproto_identity::key::{KeyType, identify_key};
use atproto_identity::model::{Document, VerificationMethod};
use atproto_identity::resolve::{HickoryDnsResolver, InnerIdentityResolver};
use atproto_plc::{Operation as PlcOperation, VerifyingKey as PlcVerifyingKey};

const PLC_DIRECTORY: &str = "plc.directory";
const PLC_DIRECTORY_URL: &str = "https://plc.directory";

#[derive(Parser, Debug)]
#[command(name = "plc-key")]
#[command(about = "Add a verification method to a did:plc document")]
struct Args {
    /// Handle or DID of the identity
    #[arg(short, long)]
    subject: String,

    /// Account password (can also be set via ATPROTO_PASSWORD env var)
    #[arg(short, long, env = "ATPROTO_PASSWORD")]
    password: String,

    /// did:key to add as verification method
    #[arg(short, long)]
    key: String,

    /// Key identifier fragment (without #)
    #[arg(short = 'i', long)]
    key_id: String,

    /// PLC operation token (skips email request if provided)
    #[arg(short, long, env = "ATPROTO_PLC_TOKEN")]
    token: Option<String>,
}

/// Response from com.atproto.identity.signPlcOperation
#[derive(Debug, Deserialize)]
struct SignPlcOperationResponse {
    operation: serde_json::Value,
}

/// Request for com.atproto.identity.signPlcOperation
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct SignPlcOperationRequest {
    token: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    rotation_keys: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    also_known_as: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    verification_methods: Option<HashMap<String, String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    services: Option<HashMap<String, ServiceRecord>>,
}

#[derive(Debug, Serialize)]
struct ServiceRecord {
    #[serde(rename = "type")]
    service_type: String,
    endpoint: String,
}

/// Validate that the key is a P256 or K256 public key
fn validate_key(key: &str) -> Result<()> {
    let key_data = identify_key(key).with_context(|| format!("Failed to parse key: {}", key))?;

    match key_data.key_type() {
        KeyType::P256Public | KeyType::K256Public => Ok(()),
        other => bail!(
            "Unsupported key type: {:?}. Only P256 and K256 public keys are supported.",
            other
        ),
    }
}

/// Extract the multibase key from a did:key URI
fn extract_multibase_key(did_key: &str) -> &str {
    did_key.strip_prefix("did:key:").unwrap_or(did_key)
}

/// Get the ID from a verification method
fn get_verification_method_id(vm: &VerificationMethod) -> Option<&str> {
    match vm {
        VerificationMethod::Multikey { id, .. } => Some(id.as_str()),
        VerificationMethod::Other { .. } => None,
    }
}

/// Get the public key multibase from a verification method
fn get_verification_method_key(vm: &VerificationMethod) -> Option<&str> {
    match vm {
        VerificationMethod::Multikey {
            public_key_multibase,
            ..
        } => Some(public_key_multibase.as_str()),
        VerificationMethod::Other { .. } => None,
    }
}

/// Validate that the key_id is not already in use
fn validate_key_id(doc: &Document, key_id: &str) -> Result<()> {
    let full_id = format!("{}#{}", doc.id, key_id);
    let fragment_id = format!("#{}", key_id);

    for vm in &doc.verification_method {
        if let Some(id) = get_verification_method_id(vm) {
            if id == full_id || id == fragment_id || id == key_id {
                bail!("Key ID '{}' is already in use in the DID document", key_id);
            }
        }
    }
    Ok(())
}

/// Extract the PDS host from a DID document
fn extract_pds_host(doc: &Document) -> Result<String> {
    for service in &doc.service {
        if service.id.ends_with("#atproto_pds")
            || service.id == "atproto_pds"
            || service.r#type == "AtprotoPersonalDataServer"
        {
            return Ok(service.service_endpoint.clone());
        }
    }
    bail!("No PDS service endpoint found in DID document")
}

/// Create an authenticated session
async fn create_auth_session(
    http_client: &reqwest::Client,
    pds_host: &str,
    identifier: &str,
    password: &str,
) -> Result<(AppPasswordAuth, String)> {
    let session = create_session(http_client, pds_host, identifier, password, None)
        .await
        .with_context(|| "Failed to create session")?;

    let auth = AppPasswordAuth {
        access_token: session.access_jwt.clone(),
    };

    Ok((auth, session.did))
}

/// Request a PLC operation signature token (sent via email)
async fn request_plc_signature(
    http_client: &reqwest::Client,
    pds_host: &str,
    auth: &AppPasswordAuth,
) -> Result<()> {
    let url = format!(
        "{}/xrpc/com.atproto.identity.requestPlcOperationSignature",
        pds_host
    );

    let mut headers = HeaderMap::default();
    headers.insert(
        reqwest::header::AUTHORIZATION,
        reqwest::header::HeaderValue::from_str(&format!("Bearer {}", auth.access_token))?,
    );

    let http_response = http_client
        .post(url)
        .headers(headers)
        .send()
        .await
        .with_context(|| "HTTP Request failed")?;

    print!("HTTP response {}", http_response.status());

    Ok(())
}

/// Prompt the user to enter the PLC signing token
fn prompt_for_token() -> Result<String> {
    print!("Enter the PLC signing token from email: ");
    io::stdout().flush()?;

    let mut token = String::new();
    io::stdin().read_line(&mut token)?;

    let token = token.trim().to_string();
    if token.is_empty() {
        bail!("Token cannot be empty");
    }

    Ok(token)
}

/// Sign a PLC operation with the updated verification methods
async fn sign_plc_operation(
    http_client: &reqwest::Client,
    pds_host: &str,
    auth: &AppPasswordAuth,
    token: &str,
    new_key: &str,
    key_id: &str,
    current_doc: &Document,
) -> Result<SignPlcOperationResponse> {
    // Build verification_methods map from current doc + new key
    let mut verification_methods: HashMap<String, String> = HashMap::new();

    // Add existing verification methods
    for vm in &current_doc.verification_method {
        if let Some(id) = get_verification_method_id(vm) {
            if let Some(key) = get_verification_method_key(vm) {
                // Extract just the fragment part
                let id_fragment = id.split('#').last().unwrap_or(id);
                verification_methods.insert(id_fragment.to_string(), key.to_string());
            }
        }
    }

    // Add the new key
    verification_methods.insert(
        key_id.to_string(),
        extract_multibase_key(new_key).to_string(),
    );

    println!(
        "DEBUG: verification_methods being sent: {:?}",
        verification_methods
    );

    let request = SignPlcOperationRequest {
        token: token.to_string(),
        rotation_keys: None,
        also_known_as: None,
        verification_methods: Some(verification_methods),
        services: None,
    };

    let request_json = serde_json::to_value(&request)?;
    println!(
        "DEBUG: Full request JSON: {}",
        serde_json::to_string_pretty(&request_json)?
    );

    let url = format!("{}/xrpc/com.atproto.identity.signPlcOperation", pds_host);

    let response_value = post_apppassword_json(http_client, auth, &url, request_json)
        .await
        .with_context(|| "Failed to sign PLC operation")?;

    println!(
        "DEBUG: Response from signPlcOperation: {}",
        serde_json::to_string_pretty(&response_value)?
    );

    let result: SignPlcOperationResponse = serde_json::from_value(response_value)
        .with_context(|| "Failed to parse signPlcOperation response")?;

    Ok(result)
}

/// Submit a signed PLC operation via the PDS
async fn submit_plc_operation(
    http_client: &reqwest::Client,
    pds_host: &str,
    auth: &AppPasswordAuth,
    operation: &serde_json::Value,
) -> Result<()> {
    let url = format!("{}/xrpc/com.atproto.identity.submitPlcOperation", pds_host);

    let response = post_apppassword_json(
        http_client,
        auth,
        &url,
        serde_json::json!({ "operation": operation }),
    )
    .await
    .with_context(|| "Failed to submit PLC operation")?;

    // Check for error response
    if let Some(error) = response.get("error") {
        let error_type = error.as_str().unwrap_or("Unknown");
        let message = response
            .get("message")
            .and_then(|m| m.as_str())
            .unwrap_or("No message");
        bail!("PLC operation failed: {error_type} - {message} - {response:?}");
    }

    Ok(())
}

/// Fetch the audit log for a DID from plc.directory
async fn fetch_audit_log(http_client: &reqwest::Client, did: &str) -> Result<Vec<PlcOperation>> {
    let url = format!("{}/{}/log", PLC_DIRECTORY_URL, did);

    let response = http_client
        .get(&url)
        .send()
        .await
        .with_context(|| format!("Failed to fetch audit log from {}", url))?;

    if !response.status().is_success() {
        bail!("Failed to fetch audit log: HTTP {}", response.status());
    }

    let operations: Vec<PlcOperation> = response
        .json()
        .await
        .with_context(|| "Failed to parse audit log response")?;

    Ok(operations)
}

/// Verify a signed operation against the rotation keys from the audit log
fn verify_signed_operation(
    signed_operation: &serde_json::Value,
    audit_log: &[PlcOperation],
) -> Result<()> {
    // Parse the signed operation
    let operation: PlcOperation = serde_json::from_value(signed_operation.clone())
        .with_context(|| "Failed to parse signed operation")?;

    // Get rotation keys from the last operation in the audit log
    let last_op = audit_log
        .last()
        .ok_or_else(|| anyhow::anyhow!("Audit log is empty"))?;

    let rotation_keys = last_op
        .rotation_keys()
        .ok_or_else(|| anyhow::anyhow!("Last operation has no rotation keys"))?;

    // Convert rotation keys to verifying keys
    let verifying_keys: Vec<PlcVerifyingKey> = rotation_keys
        .iter()
        .map(|k| PlcVerifyingKey::from_did_key(k))
        .collect::<std::result::Result<Vec<_>, _>>()
        .with_context(|| "Failed to parse rotation keys")?;

    // Verify the signature
    operation
        .verify(&verifying_keys)
        .with_context(|| "Operation signature verification failed")?;

    Ok(())
}

/// Build a proposed document showing what the result will look like
fn build_proposed_document(doc: &Document, new_key: &str, key_id: &str) -> serde_json::Value {
    let mut proposed = serde_json::to_value(doc).unwrap();

    if let Some(vms) = proposed.get_mut("verificationMethod") {
        if let Some(arr) = vms.as_array_mut() {
            arr.push(serde_json::json!({
                "id": format!("{}#{}", doc.id, key_id),
                "type": "Multikey",
                "controller": doc.id,
                "publicKeyMultibase": extract_multibase_key(new_key)
            }));
        }
    }

    proposed
}

#[tokio::main]
async fn main() -> Result<()> {
    let args = Args::parse();

    // Create HTTP client
    let http_client = reqwest::Client::new();

    // Create identity resolver
    let resolver = InnerIdentityResolver {
        dns_resolver: Arc::new(HickoryDnsResolver::create_resolver(&[])),
        http_client: http_client.clone(),
        plc_hostname: PLC_DIRECTORY.to_string(),
    };

    // 1. Resolve subject to DID document
    println!("Resolving identity: {}", args.subject);
    let doc = resolver
        .resolve(&args.subject)
        .await
        .with_context(|| format!("Failed to resolve identity: {}", args.subject))?;
    let did = &doc.id;
    println!("Resolved DID: {}", did);
    println!("\nCurrent DID Document:");
    println!("{}", serde_json::to_string_pretty(&doc)?);

    // 2. Validate the key
    println!("\nValidating key: {}", args.key);
    validate_key(&args.key)?;
    println!("Key is valid (P256 or K256)");

    // 3. Validate the key_id
    validate_key_id(&doc, &args.key_id)?;
    println!("Key ID '{}' is available", args.key_id);

    // 4. Extract PDS host from document
    let pds_host = extract_pds_host(&doc)?;
    println!("\nPDS: {}", pds_host);

    // 5. Create session
    println!("\nCreating session...");
    let (auth, session_did) =
        create_auth_session(&http_client, &pds_host, did, &args.password).await?;
    println!("Session created for: {}", session_did);

    // 6. Show proposed changes
    println!("\nProposed DID Document (with new key):");
    let proposed = build_proposed_document(&doc, &args.key, &args.key_id);
    println!("{}", serde_json::to_string_pretty(&proposed)?);

    // 7. Get PLC operation token
    let token = match args.token {
        Some(token) => {
            println!("\nUsing provided PLC operation token.");
            token
        }
        None => {
            println!("\nRequesting PLC operation signature...");
            println!("Check your email for the confirmation token.");
            request_plc_signature(&http_client, &pds_host, &auth).await?;
            prompt_for_token()?
        }
    };

    // 8. Sign PLC operation
    println!("\nSigning PLC operation...");
    let response = sign_plc_operation(
        &http_client,
        &pds_host,
        &auth,
        &token,
        &args.key,
        &args.key_id,
        &doc,
    )
    .await?;

    // 9. Verify the signed operation against the audit log
    println!("\nVerifying signed operation...");
    let audit_log = fetch_audit_log(&http_client, did).await?;
    verify_signed_operation(&response.operation, &audit_log)?;
    println!("Operation signature verified against rotation keys");

    // 10. Submit PLC operation via PDS
    println!("\nSubmitting PLC operation...");
    submit_plc_operation(&http_client, &pds_host, &auth, &response.operation).await?;

    // 11. Verify the change
    println!("\nVerifying change...");
    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
    let updated_doc = resolver
        .resolve(did)
        .await
        .with_context(|| "Failed to verify DID document update")?;
    println!("\nUpdated DID Document:");
    println!("{}", serde_json::to_string_pretty(&updated_doc)?);

    println!(
        "\nSuccess! Key '{}' has been added to the DID document.",
        args.key_id
    );

    Ok(())
}