use std::{collections::BTreeSet, fs};

use remora_contracts::{
    ActivationResponseDto, AttemptActivateRequestDto, AttemptCheckpointRequestDto,
    AttemptOutcomeRequestDto, AttemptProgressRequestDto, AttemptRenewRequestDto,
    AttemptResolvedSourceRequestDto, AttemptStagingMultipartRequestDto,
    AttemptStagingObjectRequestDto, AttemptStagingWriteIntentRequestDto, AttemptStopRequestDto,
    ClaimRequestDto, ClaimResponseDto, JobPhaseProjectionDto, LeaseProjectionDto,
};
use serde::{Serialize, de::DeserializeOwned};
use serde_json::Value;

#[test]
fn api_contract_documents_parse_and_local_references_resolve() {
    let openapi_source = fs::read_to_string("../../docs/api/openapi.yaml").unwrap();
    let openapi: Value = serde_yaml_ng::from_str(&openapi_source).unwrap();
    assert_eq!(
        openapi.pointer("/openapi").and_then(Value::as_str),
        Some("3.2.0")
    );
    assert_local_references_resolve(&openapi);

    let events_source = fs::read_to_string("../../docs/api/events.schema.json").unwrap();
    let events: Value = serde_json::from_str(&events_source).unwrap();
    assert_eq!(
        events.pointer("/$schema").and_then(Value::as_str),
        Some("https://json-schema.org/draft/2020-12/schema")
    );
    assert_local_references_resolve(&events);
    jsonschema::meta::validate(&events).unwrap();
    let event_validator = jsonschema::draft202012::options()
        .should_validate_formats(true)
        .build(&events)
        .unwrap();

    let corpus: Value =
        serde_json::from_str(include_str!("../../../contracts/golden/domain.json")).unwrap();
    for frame in corpus["event_frames"].as_array().unwrap() {
        event_validator.validate(frame).unwrap();
    }

    let component_schema = openapi_component_schema(&openapi);
    jsonschema::meta::validate(&component_schema).unwrap();
    validate_openapi_corpus(
        &component_schema,
        &corpus,
        "JobSubmission",
        "job_submissions",
    );
    validate_openapi_corpus(
        &component_schema,
        &corpus,
        "AttemptOutcomeRequest",
        "attempt_outcomes",
    );
    validate_openapi_corpus(&component_schema, &corpus, "JobFailure", "job_failures");
}

#[test]
fn worker_routes_and_dtos_conform_to_the_mutable_openapi_contract() {
    let openapi: Value =
        serde_yaml_ng::from_str(&fs::read_to_string("../../docs/api/openapi.yaml").unwrap())
            .unwrap();
    let component_schema = openapi_component_schema(&openapi);
    let mutation = serde_json::json!({
        "operation_id": "operation_0001",
        "fence": 1,
        "worker_incarnation": "worker_0001"
    });
    let descriptor = serde_json::json!({
        "id": "descriptor_0001",
        "kind": "resolved_source",
        "schema_version": 1,
        "byte_length": 64,
        "item_count": 1,
        "checksum": format!("sha256:{}", "22".repeat(32))
    });
    let execution = serde_json::json!({
        "id": "grant_0001",
        "job_id": "job_00000001",
        "attempt_id": "attempt_0001",
        "fence": 1,
        "worker_incarnation": "worker_0001",
        "source": {
            "kind": "http",
            "display_url": "https://example.invalid/artifact",
            "result_path": "artifact.bin"
        },
        "sensitive_source_ref": "source_00001",
        "credential_version": null,
        "resolved_source_ref": null,
        "target": {
            "store": "primary",
            "bucket": "remora-results",
            "prefix": "users/alice"
        },
        "staging_target": {
            "store": "primary",
            "bucket": "remora-results",
            "prefix": "users/alice/_remora/a/attempt_0001/"
        },
        "execution_grant_ref": "grant_0001",
        "checkpoint_ref": null
    });
    let authority = serde_json::json!({
        "attempt_id": "attempt_0001",
        "fence": 1,
        "worker_incarnation": "worker_0001"
    });
    let lease = serde_json::json!({
        "expires_at": "2026-07-17T00:01:00Z",
        "renew_after": "2026-07-17T00:00:30Z"
    });

    let claim_request = serde_json::json!({
        "operation_id": "operation_0001",
        "worker_incarnation": "worker_0001"
    });
    assert_dto_and_schema::<ClaimRequestDto>(&component_schema, "ClaimRequest", claim_request);
    let claim_response = serde_json::json!({
        "authority": authority,
        "execution": execution
    });
    assert_dto_and_schema::<ClaimResponseDto>(
        &component_schema,
        "ClaimResponse",
        claim_response.clone(),
    );
    assert_dto_and_schema::<AttemptRenewRequestDto>(
        &component_schema,
        "AttemptRenewRequest",
        mutation.clone(),
    );
    assert_dto_and_schema::<AttemptActivateRequestDto>(
        &component_schema,
        "AttemptActivateRequest",
        mutation.clone(),
    );
    assert_dto_and_schema::<AttemptStopRequestDto>(
        &component_schema,
        "AttemptStopRequest",
        serde_json::json!({
            "operation_id": "operation_0001",
            "fence": 1,
            "worker_incarnation": "worker_0001"
        }),
    );
    assert_dto_and_schema::<AttemptStagingObjectRequestDto>(
        &component_schema,
        "AttemptStagingObjectRequest",
        serde_json::json!({
            "operation_id": "operation_0001",
            "fence": 1,
            "worker_incarnation": "worker_0001",
            "object": {
                "store": "primary",
                "bucket": "remora-results",
                "key": "users/alice/_remora/a/attempt_0001/o/object_0001"
            }
        }),
    );
    assert_dto_and_schema::<AttemptStagingWriteIntentRequestDto>(
        &component_schema,
        "AttemptStagingWriteIntentRequest",
        serde_json::json!({
            "operation_id": "operation_0001",
            "fence": 1,
            "worker_incarnation": "worker_0001",
            "intent": {
                "store": "primary",
                "bucket": "remora-results",
                "key": "users/alice/_remora/a/attempt_0001/o/object_0001"
            }
        }),
    );
    assert_dto_and_schema::<AttemptStagingMultipartRequestDto>(
        &component_schema,
        "AttemptStagingMultipartRequest",
        serde_json::json!({
            "operation_id": "operation_0001",
            "fence": 1,
            "worker_incarnation": "worker_0001",
            "multipart": {
                "store": "primary",
                "bucket": "remora-results",
                "key": "users/alice/_remora/a/attempt_0001/o/object_0001",
                "upload_id": "upload_0001"
            }
        }),
    );
    assert_dto_and_schema::<ActivationResponseDto>(
        &component_schema,
        "ActivationResponse",
        serde_json::json!({
            "authority": claim_response["authority"],
            "execution": claim_response["execution"],
            "lease": lease.clone()
        }),
    );
    assert_dto_and_schema::<LeaseProjectionDto>(
        &component_schema,
        "LeaseProjection",
        lease.clone(),
    );
    assert_dto_and_schema::<AttemptResolvedSourceRequestDto>(
        &component_schema,
        "AttemptResolvedSourceRequest",
        merged(
            &mutation,
            serde_json::json!({ "resolved_source_ref": descriptor.clone() }),
        ),
    );
    let mut checkpoint_descriptor = descriptor.clone();
    checkpoint_descriptor["kind"] = Value::String("checkpoint".into());
    assert_dto_and_schema::<AttemptCheckpointRequestDto>(
        &component_schema,
        "AttemptCheckpointRequest",
        merged(
            &mutation,
            serde_json::json!({
                "checkpoint_ref": checkpoint_descriptor,
                "resolved_source_id": "descriptor_0001",
                "validated_bytes": 1024,
                "validated_members": 1
            }),
        ),
    );
    assert_dto_and_schema::<AttemptProgressRequestDto>(
        &component_schema,
        "AttemptProgressRequest",
        merged(
            &mutation,
            serde_json::json!({
                "progress": {
                    "completed_bytes": 1024,
                    "total_bytes": 4096,
                    "completed_members": 1,
                    "total_members": 1
                }
            }),
        ),
    );
    assert_dto_and_schema::<AttemptOutcomeRequestDto>(
        &component_schema,
        "AttemptOutcomeRequest",
        merged(
            &mutation,
            serde_json::json!({
                "outcome": {
                    "kind": "failure",
                    "failure": {
                        "class": "transient_source",
                        "message": "temporarily unavailable"
                    }
                }
            }),
        ),
    );
    assert_dto_and_schema::<JobPhaseProjectionDto>(
        &component_schema,
        "JobPhaseProjection",
        serde_json::json!({ "phase": "RetryWait" }),
    );

    let routes = [
        (
            "/internal/v1/jobs/{job_id}/claim",
            "ClaimRequest",
            "200",
            Some("ClaimResponse"),
        ),
        (
            "/internal/v1/attempts/{attempt_id}/renew",
            "AttemptRenewRequest",
            "200",
            Some("LeaseProjection"),
        ),
        (
            "/internal/v1/attempts/{attempt_id}/activate",
            "AttemptActivateRequest",
            "200",
            Some("ActivationResponse"),
        ),
        (
            "/internal/v1/attempts/{attempt_id}/stop",
            "AttemptStopRequest",
            "204",
            None,
        ),
        (
            "/internal/v1/attempts/{attempt_id}/staging/write-intent",
            "AttemptStagingWriteIntentRequest",
            "204",
            None,
        ),
        (
            "/internal/v1/attempts/{attempt_id}/staging/object",
            "AttemptStagingObjectRequest",
            "204",
            None,
        ),
        (
            "/internal/v1/attempts/{attempt_id}/staging/multipart",
            "AttemptStagingMultipartRequest",
            "204",
            None,
        ),
        (
            "/internal/v1/attempts/{attempt_id}/resolved-source",
            "AttemptResolvedSourceRequest",
            "204",
            None,
        ),
        (
            "/internal/v1/attempts/{attempt_id}/checkpoint",
            "AttemptCheckpointRequest",
            "204",
            None,
        ),
        (
            "/internal/v1/attempts/{attempt_id}/progress",
            "AttemptProgressRequest",
            "204",
            None,
        ),
        (
            "/internal/v1/attempts/{attempt_id}/outcome",
            "AttemptOutcomeRequest",
            "200",
            Some("JobPhaseProjection"),
        ),
    ];
    for (path, request, status, response) in routes {
        assert_worker_route(&openapi, path, request, status, response);
    }
}

fn assert_dto_and_schema<T>(component_schema: &Value, definition: &str, value: Value)
where
    T: DeserializeOwned + Serialize,
{
    let dto: T = serde_json::from_value(value.clone()).unwrap();
    let serialized = serde_json::to_value(&dto).unwrap();
    assert_eq!(serialized, value);
    if matches!(
        definition,
        "ClaimRequest"
            | "AttemptRenewRequest"
            | "AttemptActivateRequest"
            | "AttemptStopRequest"
            | "AttemptStagingObjectRequest"
            | "AttemptStagingWriteIntentRequest"
            | "AttemptStagingMultipartRequest"
            | "AttemptResolvedSourceRequest"
            | "AttemptCheckpointRequest"
            | "AttemptProgressRequest"
            | "AttemptOutcomeRequest"
    ) {
        assert!(remora_contracts::request_is_valid(definition, &dto));
    }
    validate_openapi_value(component_schema, definition, &value);
}

fn validate_openapi_value(component_schema: &Value, definition: &str, value: &Value) {
    let mut schema = component_schema.clone();
    schema.as_object_mut().unwrap().insert(
        "$ref".to_owned(),
        Value::String(format!("#/$defs/{definition}")),
    );
    jsonschema::draft202012::options()
        .should_validate_formats(true)
        .build(&schema)
        .unwrap()
        .validate(value)
        .unwrap();
}

fn merged(base: &Value, extension: Value) -> Value {
    let mut value = base.clone();
    value
        .as_object_mut()
        .unwrap()
        .extend(extension.as_object().unwrap().clone());
    value
}

fn assert_worker_route(
    openapi: &Value,
    path: &str,
    request_schema: &str,
    success_status: &str,
    response_schema: Option<&str>,
) {
    let operation = &openapi["paths"][path]["post"];
    assert_eq!(
        operation.pointer("/security/0/serviceAuth"),
        Some(&serde_json::json!([]))
    );
    assert_eq!(
        operation.pointer("/requestBody/content/application~1json/schema/$ref"),
        Some(&Value::String(format!(
            "#/components/schemas/{request_schema}"
        )))
    );
    assert!(operation["responses"].get("409").is_some());
    assert!(operation["responses"].get("default").is_some());
    let success = &operation["responses"][success_status];
    match response_schema {
        Some(response_schema) => assert_eq!(
            success.pointer("/content/application~1json/schema/$ref"),
            Some(&Value::String(format!(
                "#/components/schemas/{response_schema}"
            )))
        ),
        None => assert!(success.get("content").is_none()),
    }
}

fn openapi_component_schema(openapi: &Value) -> Value {
    let mut definitions = openapi
        .pointer("/components/schemas")
        .and_then(Value::as_object)
        .unwrap()
        .clone();
    for schema in definitions.values_mut() {
        rewrite_component_references(schema);
    }
    serde_json::json!({
        "$schema": "https://json-schema.org/draft/2020-12/schema",
        "$defs": definitions
    })
}

fn rewrite_component_references(value: &mut Value) {
    match value {
        Value::Array(values) => {
            for value in values {
                rewrite_component_references(value);
            }
        }
        Value::Object(object) => {
            if let Some(Value::String(reference)) = object.get_mut("$ref")
                && let Some(name) = reference.strip_prefix("#/components/schemas/")
            {
                *reference = format!("#/$defs/{name}");
            }
            object.remove("discriminator");
            for value in object.values_mut() {
                rewrite_component_references(value);
            }
        }
        _ => {}
    }
}

fn validate_openapi_corpus(
    component_schema: &Value,
    corpus: &Value,
    definition: &str,
    corpus_key: &str,
) {
    let mut schema = component_schema.clone();
    schema.as_object_mut().unwrap().insert(
        "$ref".to_owned(),
        Value::String(format!("#/$defs/{definition}")),
    );
    let validator = jsonschema::draft202012::options()
        .should_validate_formats(true)
        .build(&schema)
        .unwrap();
    for value in corpus[corpus_key].as_array().unwrap() {
        validator.validate(value).unwrap();
    }
}

fn assert_local_references_resolve(document: &Value) {
    let mut references = BTreeSet::new();
    collect_local_references(document, &mut references);
    assert!(!references.is_empty());
    for reference in references {
        let pointer = reference.strip_prefix('#').unwrap();
        assert!(
            document.pointer(pointer).is_some(),
            "unresolved local reference {reference}"
        );
    }
}

fn collect_local_references(value: &Value, references: &mut BTreeSet<String>) {
    match value {
        Value::Array(values) => {
            for value in values {
                collect_local_references(value, references);
            }
        }
        Value::Object(object) => {
            if let Some(reference) = object.get("$ref").and_then(Value::as_str)
                && reference.starts_with("#/")
            {
                references.insert(reference.to_owned());
            }
            for value in object.values() {
                collect_local_references(value, references);
            }
        }
        _ => {}
    }
}