use std::collections::{BTreeMap, HashMap, HashSet};
use std::fmt::Debug;
use std::num::NonZeroU16;

use lazaret_core::codec::{CodecError, MAX_SCHEMA_RECORD_BYTES, decode, encode};
use lazaret_core::configuration::parse_trusted_configuration;
use lazaret_core::identity::{
    BootId, BundleDigest, CapabilityId, CertificateDigest, ClientId, ConfigurationFingerprint,
    ContainerId, CoordinatorId, DerivedPolicyDigest, DeviceId, DomainId, FileDescriptor,
    GenerationId, GuestProcessGroupId, GuestProcessId, GuestProcessStartTimeTicks, HostPath,
    HostPidLocator, IdMapDigest, InodeId, MountId, MountIdentity, NamespaceIdentity,
    NetworkSidecarId, PidfsDeviceId, PidfsInodeId, PidfsProcessAuthorityV1, PidfsProcessIdentityV1,
    PolicyId, PrepareOperationId, PreparedArtifactId, ProcessStartTimeTicks, ProtocolCredentialId,
    ReadinessFingerprint, RequestDigest, RequestId, SandboxId, SelectedManifestDigest,
    SessionAnchorId, SessionId, SourceIdentity, StreamDigest, TrustManifestDigest,
};
use lazaret_core::network::readiness_fingerprint as derive_readiness_fingerprint;
use lazaret_core::records::*;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};

#[derive(Deserialize, Serialize)]
struct RawParticipant {
    generation_id: GenerationId,
    role: ParticipantRole,
    state: ParticipantState,
    ownership: Option<ParticipantOwnership>,
}

#[derive(Deserialize, Serialize)]
struct RawTlsCredential {
    generation_identity: GenerationIdentity,
    role: TlsRole,
    subject_id: TlsSubjectId,
    certificate_der: Vec<u8>,
    certificate_sha256: CertificateDigest,
}

#[derive(Deserialize, Serialize)]
struct RawProtocolCredential {
    protocol_version: ProtocolVersion,
    credential_id: ProtocolCredentialId,
    generation_identity: GenerationIdentity,
    role: ProtocolCredentialRole,
    subject_id: ProtocolSubjectId,
    public_key: Vec<u8>,
    issuer_signature: Vec<u8>,
}

#[derive(Deserialize, Serialize)]
struct RawSandboxCredentialBundle {
    protocol_version: ProtocolVersion,
    generation_identity: GenerationIdentity,
    trust_manifest: TrustManifest,
    coordinator_protocol_credential: ProtocolCredential,
    sandbox_protocol_credential: ProtocolCredential,
    sandbox_private_key: Vec<u8>,
}

#[derive(Deserialize, Serialize)]
struct RawMetadata {
    schema_id: String,
    schema_version: u32,
}

#[derive(Serialize)]
struct MissingVersion {
    schema_id: String,
}

#[derive(Serialize)]
struct RawDomainStore {
    metadata: StoreMetadata,
    domain: Domain,
    generations: Vec<Generation>,
    participants: Vec<Participant>,
    sessions: Vec<Session>,
    session_terminals: Vec<SessionTerminal>,
    mailboxes: Vec<Mailbox>,
}

#[derive(Serialize)]
struct RawStorageCatalog {
    metadata: StoreMetadata,
    prepared_artifacts: Vec<PreparedArtifact>,
    artifact_leases: Vec<ArtifactLease>,
    prepare_operations: Vec<PrepareOperation>,
}

#[derive(Serialize)]
struct RawTerminalResize {
    session_id: SessionId,
    rows: u16,
    columns: u16,
    pixel_width: u16,
    pixel_height: u16,
}

#[derive(Serialize)]
struct RawSessionTerminal {
    generation_id: GenerationId,
    session_id: SessionId,
    state: SessionTerminalState,
    container_id: Option<ContainerId>,
    host_coordinator: PersistedHostProcessOwnership,
    pty_master: Option<PtyIdentity>,
    exec_client: PersistedHostProcessOwnership,
    guest_process: Option<GuestProcessIdentity>,
}

#[derive(Deserialize, Serialize)]
struct InvalidSessionExit {
    session_id: SessionId,
    outcome: SessionExitOutcome,
    exit_code: Option<i32>,
}

#[derive(Serialize)]
struct RawDerivedNetworkPolicy {
    schema_id: String,
    address_families: Vec<AddressFamilyDisposition>,
    default_actions: Vec<DefaultAction>,
    prefix_rules: Vec<PrefixRule>,
    resolvers: Vec<Resolver>,
}

#[derive(Serialize)]
struct RawNetworkReadinessAttestation {
    schema_id: String,
    domain_id: DomainId,
    generation_id: GenerationId,
    configuration_fingerprint: ConfigurationFingerprint,
    network_sidecar_id: NetworkSidecarId,
    derived_policy_digest: DerivedPolicyDigest,
    enforcement_namespace: NamespaceIdentity,
    runtime_supervisor: PidfsProcessIdentityV1,
    connectivity_helper: PidfsProcessIdentityV1,
    observation: NetworkObservation,
}

#[derive(Serialize)]
struct RawNetworkReadinessIdentity {
    schema_id: String,
    domain_id: DomainId,
    generation_id: GenerationId,
    configuration_fingerprint: ConfigurationFingerprint,
    network_sidecar_id: NetworkSidecarId,
    derived_policy_digest: DerivedPolicyDigest,
    enforcement_namespace: NamespaceIdentity,
}

#[derive(Serialize)]
struct RawMailboxRequest {
    protocol_version: ProtocolVersion,
    session_binding: SessionBinding,
    request_id: RequestId,
    session_capability: SessionRequestCapability,
    sandbox_credential_id: ProtocolCredentialId,
    command_argv: Vec<Vec<u8>>,
    justification: String,
}

#[derive(Serialize)]
struct RawTrustedHostConfigurationV1 {
    schema_id: String,
    storage: StorageConfigurationV1,
    network: NetworkConfigurationV1,
    profiles: BTreeMap<String, ProfileConfigurationV1>,
}

fn opaque(byte: u8) -> [u8; 32] {
    [byte; 32]
}

fn domain_id() -> DomainId {
    DomainId::new(opaque(1))
}

fn generation_id() -> GenerationId {
    GenerationId::new(7)
}

fn coordinator_id() -> CoordinatorId {
    CoordinatorId::new(opaque(2))
}

fn sidecar_id() -> NetworkSidecarId {
    NetworkSidecarId::new(opaque(3))
}

fn sandbox_id() -> SandboxId {
    SandboxId::new(opaque(4))
}

fn session_id() -> SessionId {
    SessionId::new(opaque(5))
}

fn process(seed: u32) -> PidfsProcessIdentityV1 {
    PidfsProcessIdentityV1::new(
        PidfsProcessAuthorityV1 {
            boot_id: BootId::new([0x42; 16]),
            pidfs_st_dev: PidfsDeviceId::new(7),
            pidfs_st_ino: PidfsInodeId::new(u64::from(seed)),
        },
        HostPidLocator::new(seed),
        ProcessStartTimeTicks::new(u64::from(seed) * 10),
    )
}

fn owned_process(seed: u32) -> PersistedHostProcessOwnership {
    PersistedHostProcessOwnership::committed(process(seed))
}

fn namespace() -> NamespaceIdentity {
    NamespaceIdentity {
        device: DeviceId::new(8),
        inode: InodeId::new(9),
    }
}

fn mount() -> MountIdentity {
    MountIdentity {
        mount_id: MountId::new(10),
        device: DeviceId::new(11),
        inode: InodeId::new(12),
    }
}

fn generation_identity() -> GenerationIdentity {
    GenerationIdentity {
        agent_profile: "trusted".into(),
        domain_id: domain_id(),
        generation_id: generation_id(),
        configuration_fingerprint: ConfigurationFingerprint::new(opaque(6)),
        coordinator_id: coordinator_id(),
        sandbox_id: sandbox_id(),
    }
}

fn trust_manifest() -> TrustManifest {
    TrustManifest {
        protocol_version: ProtocolVersion::V1,
        generation_identity: generation_identity(),
        tls_ca_sha256: CertificateDigest::new(opaque(7)),
        protocol_issuer_public_key: vec![8; 32],
        issuer_signature: vec![9; 64],
    }
}

fn protocol_credential(role: ProtocolCredentialRole) -> ProtocolCredential {
    match role {
        ProtocolCredentialRole::SessionCoordinator => ProtocolCredential::session_coordinator(
            ProtocolCredentialId::new(opaque(10)),
            generation_identity(),
            coordinator_id(),
            vec![11; 32],
            vec![12; 64],
        ),
        ProtocolCredentialRole::AgentSandbox => ProtocolCredential::agent_sandbox(
            ProtocolCredentialId::new(opaque(10)),
            generation_identity(),
            sandbox_id(),
            vec![11; 32],
            vec![12; 64],
        ),
    }
}

fn session_binding() -> SessionBinding {
    SessionBinding {
        generation_identity: generation_identity(),
        cli_id: ClientId::new(opaque(13)),
        session_id: session_id(),
        session_anchor_public_key: vec![14; 32],
    }
}

fn session_capability() -> SessionRequestCapability {
    SessionRequestCapability {
        protocol_version: ProtocolVersion::V1,
        capability_id: CapabilityId::new(opaque(15)),
        session_binding: session_binding(),
        coordinator_credential_id: ProtocolCredentialId::new(opaque(16)),
        sandbox_credential_id: ProtocolCredentialId::new(opaque(17)),
        coordinator_signature: vec![18; 64],
    }
}

fn network_observation() -> NetworkObservation {
    NetworkObservation {
        policy_delivery: PolicyDelivery::SealedMemfd,
        helper_ready: true,
        observed_routes: vec![IpNetwork {
            address: IpAddress::V4([10, 0, 0, 0]),
            prefix_length: 8,
        }],
        routes_attested: true,
    }
}

fn readiness_attestation() -> NetworkReadinessAttestation {
    NetworkReadinessAttestation::v1(NetworkReadinessAttestationFields {
        domain_id: domain_id(),
        generation_id: generation_id(),
        configuration_fingerprint: ConfigurationFingerprint::new(opaque(19)),
        network_sidecar_id: sidecar_id(),
        derived_policy_digest: DerivedPolicyDigest::new(opaque(20)),
        enforcement_namespace: namespace(),
        runtime_supervisor: process(21),
        connectivity_helper: process(22),
        observation: network_observation(),
    })
}

fn readiness_identity() -> NetworkReadinessIdentity {
    NetworkReadinessIdentity::v1(NetworkReadinessIdentityFields {
        domain_id: domain_id(),
        generation_id: generation_id(),
        configuration_fingerprint: ConfigurationFingerprint::new(opaque(19)),
        network_sidecar_id: sidecar_id(),
        derived_policy_digest: DerivedPolicyDigest::new(opaque(20)),
        enforcement_namespace: namespace(),
    })
}

fn reachability_probe() -> NetworkReachabilityProbe {
    NetworkReachabilityProbe {
        domain_id: domain_id(),
        generation_id: generation_id(),
        enforcement_namespace: namespace(),
        denied_destinations: vec![DeniedDestinationObservation {
            destination: IpAddress::V4([10, 1, 2, 3]),
            unreachable: true,
        }],
        permitted_destinations: vec![PermittedDestinationObservation {
            destination: IpAddress::V4([1, 1, 1, 1]),
            reachable: true,
        }],
        dns_observations: vec![DnsObservation {
            resolver: IpAddress::V4([9, 9, 9, 9]),
            query_name: "example.invalid".into(),
            resolved_addresses: vec![IpAddress::V6([1; 16])],
        }],
    }
}

fn readiness_evidence() -> NetworkReadinessEvidence {
    let readiness_identity = readiness_identity();
    NetworkReadinessEvidence {
        domain_id: domain_id(),
        generation_id: generation_id(),
        network_sidecar_id: sidecar_id(),
        readiness_fingerprint: derive_readiness_fingerprint(&readiness_identity).unwrap(),
        attestation: readiness_attestation(),
        readiness_identity,
        reachability_probe: reachability_probe(),
    }
}

fn prepared_ownership() -> PreparedLowerOwnership {
    PreparedLowerOwnership {
        prepared_artifact_id: Some(PreparedArtifactId::new(opaque(24))),
        source_identity: Some(SourceIdentity::new(vec![25, 26])),
        lease_state: Some(ParticipantState::LeaseHeld),
        path: Some(HostPath::new(vec![b'/', 0xff])),
        device: Some(DeviceId::new(27)),
        inode: Some(InodeId::new(28)),
    }
}

fn network_ownership() -> NetworkOwnership {
    let readiness_evidence = readiness_evidence();
    NetworkOwnership {
        derived_policy_digest: Some(DerivedPolicyDigest::new(opaque(20))),
        readiness_fingerprint: Some(readiness_evidence.readiness_fingerprint),
        readiness_evidence: Some(readiness_evidence),
        namespace: Some(namespace()),
        runtime_supervisor: owned_process(21),
    }
}

fn mutate_network_ownership(mutate: impl FnOnce(&mut NetworkOwnership)) -> NetworkOwnership {
    let mut ownership = network_ownership();
    mutate(&mut ownership);
    ownership
}

fn state_anchor_ownership(
    daemon: PersistedHostProcessOwnership,
    mount_identity: Option<MountIdentity>,
) -> StateAnchorOwnership {
    StateAnchorOwnership {
        lower_artifact_id: Some(PreparedArtifactId::new(opaque(24))),
        upper_path: Some(HostPath::new(b"upper".to_vec())),
        work_path: Some(HostPath::new(b"work".to_vec())),
        merged_path: Some(HostPath::new(b"merged".to_vec())),
        daemon,
        user_namespace: Some(namespace()),
        mount_namespace: Some(namespace()),
        mount: mount_identity,
    }
}

fn pty() -> PtyIdentity {
    PtyIdentity {
        master_fd: FileDescriptor::new(31),
        device: DeviceId::new(32),
        inode: InodeId::new(33),
    }
}

fn guest_process() -> GuestProcessIdentity {
    GuestProcessIdentity {
        leader_pid: GuestProcessId::new(35),
        leader_start_time: GuestProcessStartTimeTicks::new(36),
        process_group_id: GuestProcessGroupId::new(37),
    }
}

fn active_state(role: ParticipantRole) -> ParticipantState {
    match role {
        ParticipantRole::PreparedLower => ParticipantState::LeaseHeld,
        ParticipantRole::StateAnchor | ParticipantRole::GitRouting => ParticipantState::Prepared,
        ParticipantRole::Network
        | ParticipantRole::Helper
        | ParticipantRole::Sandbox
        | ParticipantRole::Publication => ParticipantState::Committed,
    }
}

fn complete_ownership(role: ParticipantRole, state: ParticipantState) -> ParticipantOwnership {
    match role {
        ParticipantRole::PreparedLower => {
            let mut ownership = prepared_ownership();
            ownership.lease_state = Some(state);
            ParticipantOwnership::PreparedLower(ownership)
        }
        ParticipantRole::StateAnchor => ParticipantOwnership::StateAnchor(state_anchor_ownership(
            owned_process(41),
            Some(mount()),
        )),
        ParticipantRole::GitRouting => ParticipantOwnership::GitRouting(GitRoutingOwnership {
            mountpoint: Some(HostPath::new(b"workspace".to_vec())),
            host_root: Some(HostPath::new(b"host".to_vec())),
            isolated_root: Some(HostPath::new(b"isolated".to_vec())),
            daemon: owned_process(42),
            mount: Some(mount()),
        }),
        ParticipantRole::Network => ParticipantOwnership::Network(Box::new(network_ownership())),
        ParticipantRole::Helper => ParticipantOwnership::Helper(HelperOwnership {
            network_sidecar_id: Some(sidecar_id()),
            connectivity_helper: owned_process(22),
            runtime_supervisor: owned_process(21),
        }),
        ParticipantRole::Sandbox => ParticipantOwnership::Sandbox(SandboxOwnership {
            container_id: Some(ContainerId::new(opaque(29))),
            runtime_root: Some(HostPath::new(b"runtime".to_vec())),
            bundle_digest: Some(BundleDigest::new(opaque(43))),
            host_coordinator: owned_process(30),
            runtime_supervisor: owned_process(21),
            attached_runsc: owned_process(44),
        }),
        ParticipantRole::Publication => ParticipantOwnership::Publication(PublicationOwnership {
            socket_path: Some(HostPath::new(b"socket".to_vec())),
            identity_path: Some(HostPath::new(b"identity".to_vec())),
        }),
    }
}

fn release_ownership(role: ParticipantRole) -> ParticipantOwnership {
    let state = if role == ParticipantRole::PreparedLower {
        ParticipantState::ReleaseIntent
    } else {
        active_state(role)
    };
    complete_ownership(role, state)
}

fn terminal_snapshot(
    state: SessionTerminalState,
    host_coordinator: PersistedHostProcessOwnership,
    exec_client: PersistedHostProcessOwnership,
    has_physical_evidence: bool,
) -> Result<SessionTerminal, RecordInvariantError> {
    SessionTerminal::snapshot(SessionTerminalFields {
        generation_id: generation_id(),
        session_id: session_id(),
        state,
        container_id: has_physical_evidence.then(|| ContainerId::new(opaque(29))),
        host_coordinator,
        pty_master: has_physical_evidence.then_some(pty()),
        exec_client,
        guest_process: has_physical_evidence.then_some(guest_process()),
    })
}

#[derive(Clone, Copy, Debug)]
enum ParticipantProcessEdge {
    StateAnchorDaemon,
    GitRoutingDaemon,
    NetworkRuntimeSupervisor,
    HelperConnectivityHelper,
    HelperRuntimeSupervisor,
    SandboxHostCoordinator,
    SandboxRuntimeSupervisor,
    SandboxAttachedRunsc,
}

fn participant_edge_snapshot(
    edge: ParticipantProcessEdge,
    process_ownership: PersistedHostProcessOwnership,
) -> Participant {
    let (role, ownership) = match edge {
        ParticipantProcessEdge::StateAnchorDaemon => (
            ParticipantRole::StateAnchor,
            ParticipantOwnership::StateAnchor(state_anchor_ownership(process_ownership, None)),
        ),
        ParticipantProcessEdge::GitRoutingDaemon => (
            ParticipantRole::GitRouting,
            ParticipantOwnership::GitRouting(GitRoutingOwnership {
                mountpoint: Some(HostPath::new(b"workspace".to_vec())),
                host_root: Some(HostPath::new(b"host".to_vec())),
                isolated_root: Some(HostPath::new(b"isolated".to_vec())),
                daemon: process_ownership,
                mount: None,
            }),
        ),
        ParticipantProcessEdge::NetworkRuntimeSupervisor => (
            ParticipantRole::Network,
            ParticipantOwnership::Network(Box::new(NetworkOwnership {
                derived_policy_digest: Some(DerivedPolicyDigest::new(opaque(20))),
                readiness_fingerprint: None,
                readiness_evidence: None,
                namespace: None,
                runtime_supervisor: process_ownership,
            })),
        ),
        ParticipantProcessEdge::HelperConnectivityHelper => (
            ParticipantRole::Helper,
            ParticipantOwnership::Helper(HelperOwnership {
                network_sidecar_id: Some(sidecar_id()),
                connectivity_helper: process_ownership,
                runtime_supervisor: owned_process(121),
            }),
        ),
        ParticipantProcessEdge::HelperRuntimeSupervisor => (
            ParticipantRole::Helper,
            ParticipantOwnership::Helper(HelperOwnership {
                network_sidecar_id: Some(sidecar_id()),
                connectivity_helper: PersistedHostProcessOwnership::unassigned(),
                runtime_supervisor: process_ownership,
            }),
        ),
        ParticipantProcessEdge::SandboxHostCoordinator => (
            ParticipantRole::Sandbox,
            ParticipantOwnership::Sandbox(SandboxOwnership {
                container_id: Some(ContainerId::new(opaque(29))),
                runtime_root: None,
                bundle_digest: None,
                host_coordinator: process_ownership,
                runtime_supervisor: PersistedHostProcessOwnership::unassigned(),
                attached_runsc: PersistedHostProcessOwnership::unassigned(),
            }),
        ),
        ParticipantProcessEdge::SandboxRuntimeSupervisor => (
            ParticipantRole::Sandbox,
            ParticipantOwnership::Sandbox(SandboxOwnership {
                container_id: Some(ContainerId::new(opaque(29))),
                runtime_root: None,
                bundle_digest: None,
                host_coordinator: owned_process(122),
                runtime_supervisor: process_ownership,
                attached_runsc: PersistedHostProcessOwnership::unassigned(),
            }),
        ),
        ParticipantProcessEdge::SandboxAttachedRunsc => (
            ParticipantRole::Sandbox,
            ParticipantOwnership::Sandbox(SandboxOwnership {
                container_id: Some(ContainerId::new(opaque(29))),
                runtime_root: None,
                bundle_digest: None,
                host_coordinator: owned_process(122),
                runtime_supervisor: owned_process(123),
                attached_runsc: process_ownership,
            }),
        ),
    };
    Participant::snapshot(
        generation_id(),
        role,
        ParticipantState::Intent,
        Some(ownership),
    )
    .unwrap()
}

fn verify_shape<T>(record: T, seen: &mut HashSet<&'static str>)
where
    T: Clone + Debug + DeserializeOwned + Eq + SchemaType + Serialize,
{
    assert!(seen.insert(T::NAME), "duplicate test for {}", T::NAME);
    let encoded = encode(&record).expect("representative record encodes");
    assert_eq!(encode(&record).unwrap(), encoded);
    assert_eq!(decode::<T>(&encoded).unwrap(), record);

    let mut truncated = encoded.clone();
    truncated.pop();
    assert!(decode::<T>(&truncated).is_err());

    let mut trailing = encoded;
    trailing.push(0);
    assert!(matches!(
        decode::<T>(&trailing),
        Err(CodecError::TrailingPayload)
    ));
    drop(record);
}

#[path = "schema_records/credentials.rs"]
mod credentials;
#[path = "schema_records/documents.rs"]
mod documents;
#[path = "schema_records/encoding.rs"]
mod encoding;
#[path = "schema_records/participant_transitions.rs"]
mod participant_transitions;
#[path = "schema_records/process_evidence.rs"]
mod process_evidence;
#[path = "schema_records/terminal.rs"]
mod terminal;