#include "whetstone_rsa/diagnosis_engine.h"
#include "whetstone_rsa/probes/confidence_threshold_probe.h"
#include "whetstone_rsa/suitability_assessor.h"

#include <iostream>

using whetstone::rsa::ConfidenceThresholdProbe;
using whetstone::rsa::DiagnosisEngine;
using whetstone::rsa::FailureClass;
using whetstone::rsa::GateDefinition;
using whetstone::rsa::GateEvidence;
using whetstone::rsa::Intervention;
using whetstone::rsa::LabelStability;
using whetstone::rsa::PolicyContext;
using whetstone::rsa::ProbeRegistry;
using whetstone::rsa::RiskTier;
using whetstone::rsa::SuitabilityAssessor;
using whetstone::rsa::to_string;

namespace {

void print_diagnosis(const whetstone::rsa::GateDiagnosis& diagnosis) {
    std::cout << "gate_id: " << diagnosis.gate_id << '\n';
    std::cout << "primary_failure_class: " << to_string(diagnosis.primary_failure_class) << '\n';
    std::cout << "confidence: " << diagnosis.confidence << '\n';
    std::cout << "recommended_policy: " << to_string(diagnosis.recommended_policy.deploy_mode) << '\n';
    std::cout << "recommended_interventions:\n";
    for (const auto intervention : diagnosis.recommended_interventions) {
        std::cout << "  - " << to_string(intervention) << '\n';
    }
    std::cout << "supporting_signals:\n";
    for (const auto& signal : diagnosis.supporting_signals) {
        std::cout << "  - " << signal << '\n';
    }
}

void print_suitability(const whetstone::rsa::GateSuitability& suitability) {
    std::cout << "suitability:\n";
    std::cout << "  recommended: " << (suitability.recommended ? "yes" : "no") << '\n';
    std::cout << "  target_layer: " << to_string(suitability.target_layer) << '\n';
    std::cout << "  bounded_output_confidence: " << suitability.bounded_output_confidence << '\n';
    std::cout << "  suggested_gate_shape: " << suitability.suggested_gate_shape << '\n';
    std::cout << "  rationale: " << suitability.rationale << '\n';
}

}  // namespace

int main() {
    ProbeRegistry probe_registry;
    probe_registry.register_probe(std::make_shared<ConfidenceThresholdProbe>());

    DiagnosisEngine diagnosis_engine(std::move(probe_registry));

    const GateDefinition gate_definition{
        .gate_id = "verification_type",
        .task_family = "classification",
        .labels = {"unit", "integration", "schema", "smoke", "docs"},
        .input_structure = whetstone::rsa::InputStructure::Unstructured,
        .input_modality = whetstone::rsa::InputModality::Textual,
        .output_topology = whetstone::rsa::OutputTopology::Discrete,
        .output_constraint = whetstone::rsa::OutputConstraint::Bounded,
        .risk_tier = RiskTier::Medium,
        .label_stability = LabelStability::Stable,
        .deterministic_baseline_available = false,
        .supports_abstain = true,
        .candidate_factorizations = {},
        .metadata = {{"case_study", "whetstone_dsl"}}
    };

    const GateEvidence gate_evidence{
        .gate_id = "verification_type",
        .model_tier = "tiny",
        .raw_accuracy = 0.849,
        .accept_rate = 0.62,
        .accuracy_on_accepted = 0.94,
        .abstain_rate = 0.38,
        .retry_recovery_rate = 0.30,
        .escalation_rate = 0.12,
        .silent_error_rate = 0.05,
        .expected_calibration_error = 0.08,
        .latency_ms = 0.8,
        .compute_cost = 1.0,
        .confusion_summary = {},
        .notes = {"demo evidence"}
    };

    const PolicyContext policy_context{
        .max_silent_error_rate = 0.08,
        .min_accuracy_on_accepted = 0.92,
        .max_latency_ms = 10.0,
        .deterministic_checks_exist = true,
        .available_escalation_targets = {"larger_specialist", "llm"}
    };

    SuitabilityAssessor suitability_assessor;
    print_suitability(suitability_assessor.assess(gate_definition));
    std::cout << '\n';
    print_diagnosis(diagnosis_engine.diagnose(gate_definition, gate_evidence, policy_context));
    return 0;
}