// fedproxy: com.fedproxy.sshPublicKey records, and the public host each service is served at.
//
// Entries are step protocol handlers:
//   list_keys   -> listRecords (paged) -> keys grouped by service, with hosts and ssh commands
//   create_key  -> createRecord        -> uri and host of the new key

use abi::repo::{create, list_records};
use abi::{Call, Datetime, Did, Event, Handle, LexiconType, Run, entries};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};

const COLLECTION: &str = "com.fedproxy.sshPublicKey";
const DOMAIN: &str = "fedproxy.com";
const PAGE_LIMIT: usize = 100;
const KEY_TYPES: [&str; 3] = ["ssh-", "ecdsa-sha2-", "sk-"];

#[derive(Deserialize)]
struct Args {
    did: Did,
    now: String,
    #[serde(default)]
    handle: Option<String>,
    #[serde(default)]
    name: Option<String>,
    #[serde(default)]
    service: Option<String>,
    #[serde(default)]
    key: Option<String>,
}

#[derive(Serialize, PartialEq)]
struct SshPublicKey {
    key: String,
    name: String,
    service: String,
    #[serde(rename = "createdAt")]
    created_at: Datetime,
}

impl LexiconType for SshPublicKey {
    fn lexicon_type() -> &'static str {
        COLLECTION
    }
}

// The service name and the handle fold into a SINGLE DNS label, dots becoming dashes,
// joined by "--". Every service then sits one level under fedproxy.com, covered by the
// one "*.fedproxy.com" wildcard certificate, so no per-service certificate is issued.
//   "app"   alice.bsky.social -> app--alice-bsky-social.fedproxy.com
//   "*.app" alice.bsky.social -> *.app--alice-bsky-social.fedproxy.com (any sub-host)
//   "*"     means the key is valid for ALL services: an authorization wildcard, no host.
fn service_host(service: &str, handle: &str) -> Option<String> {
    let flat = |s: &str| s.replace('.', "-");
    match service {
        "*" => None,
        _ => match service.strip_prefix("*.") {
            Some(rest) => Some(format!("*.{}--{}.{DOMAIN}", flat(rest), flat(handle))),
            None => Some(format!("{}--{}.{DOMAIN}", flat(service), flat(handle))),
        },
    }
}

// Line breaks are dropped, as a pasted key often wraps.
fn field(value: Option<&str>, label: &str) -> Result<String, String> {
    let value: String = value.unwrap_or_default().chars().filter(|c| *c != '\r' && *c != '\n').collect();
    let value = value.trim().to_string();
    if value.is_empty() {
        return Err(format!("{label} is required"));
    }
    if value.chars().any(char::is_control) {
        return Err(format!("{label} must not contain control characters"));
    }
    Ok(value)
}

fn validate_service(service: &str) -> Result<(), String> {
    let name = service.strip_prefix("*.").unwrap_or(service);
    if service == "*" {
        return Ok(());
    }
    let label_ok = |c: char| c.is_ascii_alphanumeric() || c == '-' || c == '.';
    if name.is_empty() || !name.chars().all(label_ok) || name.starts_with(['-', '.']) || name.ends_with(['-', '.']) {
        return Err("service must be a name, * (all), or *.name (wildcard) using letters, digits, - and .".into());
    }
    Ok(())
}

fn validate_key(key: &str) -> Result<(), String> {
    let mut parts = key.split_whitespace();
    let (kind, body) = (parts.next().unwrap_or_default(), parts.next().unwrap_or_default());
    let base64 = |c: char| c.is_ascii_alphanumeric() || c == '+' || c == '/' || c == '=';
    if !KEY_TYPES.iter().any(|prefix| kind.starts_with(prefix)) || body.is_empty() || !body.chars().all(base64) {
        return Err("key must be an OpenSSH public key, like ssh-ed25519 AAAA... comment".into());
    }
    Ok(())
}

// The tunnel to run for a service; 8080 stands for whichever local port serves it. A "*" key
// works for any service, so its example keeps a placeholder name. A wildcard bind is quoted,
// or the shell would expand its glob.
fn ssh_command(service: &str, handle: &str) -> String {
    let forward = match service {
        "*" => "my-cool-service:80:127.0.0.1:8080".to_string(),
        _ if service.starts_with("*.") => format!("'{service}:80:127.0.0.1:8080'"),
        _ => format!("{service}:80:127.0.0.1:8080"),
    };
    format!(
        "ssh -NnT -p 2222 \\\n  -o UserKnownHostsFile=/dev/null \\\n  -o StrictHostKeyChecking=no \\\n  -o PasswordAuthentication=no \\\n  -R {forward} \\\n  {handle}@{DOMAIN}"
    )
}

fn handle(value: Option<&str>) -> Result<String, String> {
    let handle = value.unwrap_or_default();
    if handle.is_empty() || !handle.contains('.') {
        return Err("handle is required to name hosts".into());
    }
    Ok(handle.to_string())
}

fn list_page(did: &Did, cursor: Option<&str>) -> Call {
    list_records(did.as_str(), COLLECTION, PAGE_LIMIT, cursor).id("page")
}

// Each page lands here; ask for the next while the PDS hands back a cursor.
fn list_keys_handler(run: Run<Args>) -> Result<Handle, String> {
    let Run { event, state, args } = run;
    Ok(match event {
        Event::Start => {
            let handle = handle(args.handle.as_deref())?;
            Handle::new().call(list_page(&args.did, None)).state(json!({ "handle": handle, "records": [] }))
        }
        Event::Result { result, .. } => {
            let output = result.output()?;
            let mut records = state["records"].as_array().cloned().unwrap_or_default();
            records.extend(output["records"].as_array().cloned().unwrap_or_default());
            let handle = handle(state["handle"].as_str())?;
            if let Some(cursor) = output["cursor"].as_str().filter(|cursor| !cursor.is_empty()) {
                return Ok(Handle::new().call(list_page(&args.did, Some(cursor))).state(json!({ "handle": handle, "records": records })));
            }
            Handle::new().result(json!({ "services": group_by_service(&records, &handle) })).done()
        }
        _ => Handle::new(),
    })
}

// Services keep the order their first key was listed in.
fn group_by_service(records: &[Value], handle: &str) -> Vec<Value> {
    let mut services: Vec<(String, Vec<Value>)> = Vec::new();
    for record in records {
        let value = &record["value"];
        let Some(service) = value["service"].as_str() else { continue };
        let key = json!({
            "uri": record["uri"],
            "name": value["name"],
            "key": value["key"],
            "createdAt": value["createdAt"],
        });
        match services.iter_mut().find(|(name, _)| name == service) {
            Some((_, keys)) => keys.push(key),
            None => services.push((service.to_string(), vec![key])),
        }
    }
    services
        .into_iter()
        .map(|(service, keys)| {
            json!({
                "service": service,
                "host": service_host(&service, handle),
                "command": ssh_command(&service, handle),
                "keys": keys,
            })
        })
        .collect()
}

fn create_key_handler(run: Run<Args>) -> Result<Handle, String> {
    let Run { event, state, args } = run;
    Ok(match event {
        Event::Start => {
            let handle = handle(args.handle.as_deref())?;
            let name = field(args.name.as_deref(), "name")?;
            let service = field(args.service.as_deref(), "service")?;
            let key = field(args.key.as_deref(), "key")?;
            validate_service(&service)?;
            validate_key(&key)?;
            let now = Datetime::new(args.now).map_err(|e| format!("invalid now: {e}"))?;
            let created_at = Datetime::try_from(now.to_utc()).map_err(|e| format!("invalid now: {e}"))?;
            let host = service_host(&service, &handle);
            let record = SshPublicKey { key, name, service: service.clone(), created_at };
            Handle::new()
                .call(create(args.did.as_str(), record).id("create"))
                .state(json!({ "service": service, "host": host }))
        }
        Event::Result { result, .. } => {
            let created = result.output()?;
            Handle::new().result(json!({ "uri": created["uri"], "service": state["service"], "host": state["host"] })).done()
        }
        _ => Handle::new(),
    })
}

entries!(
    list_keys => list_keys_handler,
    create_key => create_key_handler
);