//! Arena 1 tests: Query Path, Recency Lens, Lifecycle Filtering, and Audit Trail.

use std::sync::Arc;
use stemedb_core::types::{
    ContributingAssertion, LifecycleStage, ObjectValue, QueryAudit, QueryId, QueryParams,
};
use stemedb_lens::{Lens, RecencyLens};
use stemedb_query::{Query, QueryEngine};
use stemedb_storage::{AuditStore, GenericAuditStore, KVStore};
use stemedb_wal::Journal;
use tokio::sync::Mutex;
use tracing::debug;

use crate::agent::Agent;
use crate::helpers::{wait_until_ingested, write_assertion_to_wal};
use crate::types::{ErrorKind, SimulationError, SimulationResult};

/// Test that RecencyLens correctly selects the most recent assertion.
///
/// Creates two assertions for the same subject+predicate with different timestamps,
/// then verifies that the lens resolves to the newer one.
pub(crate) async fn run_recency_lens_test<S: KVStore + 'static>(
    journal: &Arc<Mutex<Journal>>,
    store: &Arc<S>,
    agents: &[Agent],
    result: &mut SimulationResult,
    ingestion_wait_ms: u64,
) -> bool {
    let agent = &agents[0];
    let subject = "RecencyTest_Entity";
    let predicate = "test_property";

    // Create older assertion (timestamp = 1000)
    let old_assertion = agent.sign_assertion_with_options(
        subject,
        predicate,
        ObjectValue::Text("old_value".to_string()),
        LifecycleStage::Proposed,
        Some(1000),
    );

    // Create newer assertion (timestamp = 2000)
    let new_assertion = agent.sign_assertion_with_options(
        subject,
        predicate,
        ObjectValue::Text("new_value".to_string()),
        LifecycleStage::Proposed,
        Some(2000),
    );

    // Write both to WAL and track last offset
    let _old_result = match write_assertion_to_wal(journal, &old_assertion).await {
        Ok(r) => r,
        Err(e) => {
            result.errors.push(SimulationError {
                tick: 0,
                kind: ErrorKind::WriteFailure,
                message: format!("Recency test: failed to write old assertion: {}", e),
            });
            return false;
        }
    };

    let new_result = match write_assertion_to_wal(journal, &new_assertion).await {
        Ok(r) => r,
        Err(e) => {
            result.errors.push(SimulationError {
                tick: 0,
                kind: ErrorKind::WriteFailure,
                message: format!("Recency test: failed to write new assertion: {}", e),
            });
            return false;
        }
    };
    let last_offset = new_result.end_offset;

    // Wait for ingestion to reach the last offset
    if let Err(e) = wait_until_ingested(&**store, last_offset, ingestion_wait_ms).await {
        result.errors.push(e);
        return false;
    }

    // Query to get candidates
    let engine = QueryEngine::new(store.clone());
    let query = Query::builder().subject(subject).predicate(predicate).build();

    let query_result = match engine.execute(&query).await {
        Ok(r) => r,
        Err(e) => {
            result.errors.push(SimulationError {
                tick: 0,
                kind: ErrorKind::QueryFailure,
                message: format!("Recency test: query failed: {}", e),
            });
            return false;
        }
    };

    if query_result.assertions.len() < 2 {
        result.errors.push(SimulationError {
            tick: 0,
            kind: ErrorKind::QueryFailure,
            message: format!(
                "Recency test: expected 2 assertions, got {}",
                query_result.assertions.len()
            ),
        });
        return false;
    }

    // Apply RecencyLens
    let lens = RecencyLens;
    let resolution = lens.resolve(&query_result.assertions);

    match resolution.winner {
        Some(winner) => {
            // Winner should be the newer assertion (timestamp = 2000)
            if winner.timestamp == 2000 {
                if let ObjectValue::Text(ref value) = winner.object {
                    if value == "new_value" {
                        debug!("  Recency lens correctly selected newest assertion");
                        result.queries_executed += 1;
                        return true;
                    }
                }
            }
            result.errors.push(SimulationError {
                tick: 0,
                kind: ErrorKind::LensResolutionFailure,
                message: format!(
                    "Recency test: wrong winner selected (timestamp={}, expected 2000)",
                    winner.timestamp
                ),
            });
            false
        }
        None => {
            result.errors.push(SimulationError {
                tick: 0,
                kind: ErrorKind::LensResolutionFailure,
                message: "Recency test: lens returned no winner".to_string(),
            });
            false
        }
    }
}

/// Test that lifecycle filtering correctly filters assertions.
///
/// Creates a Proposed and an Approved assertion for the same subject+predicate,
/// then verifies that querying with lifecycle=Approved returns only the Approved one.
pub(crate) async fn run_lifecycle_test<S: KVStore + 'static>(
    journal: &Arc<Mutex<Journal>>,
    store: &Arc<S>,
    agents: &[Agent],
    result: &mut SimulationResult,
    ingestion_wait_ms: u64,
) -> bool {
    let agent = &agents[0];
    let subject = "LifecycleTest_Entity";
    let predicate = "jwt_algorithm";

    // Proposed assertion (the bug - using RS256)
    let proposed = agent.sign_assertion_with_options(
        subject,
        predicate,
        ObjectValue::Text("RS256".to_string()),
        LifecycleStage::Proposed,
        Some(1000),
    );

    // Approved assertion (the fix - using ES256)
    let approved = agent.sign_assertion_with_options(
        subject,
        predicate,
        ObjectValue::Text("ES256".to_string()),
        LifecycleStage::Approved,
        Some(2000),
    );

    // Write both to WAL and track last offset
    let _proposed_result = match write_assertion_to_wal(journal, &proposed).await {
        Ok(r) => r,
        Err(e) => {
            result.errors.push(SimulationError {
                tick: 0,
                kind: ErrorKind::WriteFailure,
                message: format!("Lifecycle test: failed to write proposed: {}", e),
            });
            return false;
        }
    };

    let approved_result = match write_assertion_to_wal(journal, &approved).await {
        Ok(r) => r,
        Err(e) => {
            result.errors.push(SimulationError {
                tick: 0,
                kind: ErrorKind::WriteFailure,
                message: format!("Lifecycle test: failed to write approved: {}", e),
            });
            return false;
        }
    };
    let last_offset = approved_result.end_offset;

    // Wait for ingestion to reach the last offset
    if let Err(e) = wait_until_ingested(&**store, last_offset, ingestion_wait_ms).await {
        result.errors.push(e);
        return false;
    }

    // Query with lifecycle filter
    let engine = QueryEngine::new(store.clone());
    let query = Query::builder()
        .subject(subject)
        .predicate(predicate)
        .lifecycle(LifecycleStage::Approved)
        .build();

    let query_result = match engine.execute(&query).await {
        Ok(r) => r,
        Err(e) => {
            result.errors.push(SimulationError {
                tick: 0,
                kind: ErrorKind::QueryFailure,
                message: format!("Lifecycle test: query failed: {}", e),
            });
            return false;
        }
    };

    // Should return exactly 1 assertion (the Approved one)
    if query_result.assertions.len() != 1 {
        result.errors.push(SimulationError {
            tick: 0,
            kind: ErrorKind::QueryFailure,
            message: format!(
                "Lifecycle test: expected 1 Approved assertion, got {}",
                query_result.assertions.len()
            ),
        });
        return false;
    }

    let found = &query_result.assertions[0];
    if found.lifecycle != LifecycleStage::Approved {
        result.errors.push(SimulationError {
            tick: 0,
            kind: ErrorKind::QueryFailure,
            message: format!("Lifecycle test: expected Approved, got {:?}", found.lifecycle),
        });
        return false;
    }

    // Verify it's the ES256 value (the fix)
    if let ObjectValue::Text(ref value) = found.object {
        if value == "ES256" {
            debug!("  Lifecycle filter correctly returned only Approved assertion (ES256)");
            result.queries_executed += 1;
            return true;
        }
    }

    result.errors.push(SimulationError {
        tick: 0,
        kind: ErrorKind::QueryFailure,
        message: "Lifecycle test: wrong object value returned".to_string(),
    });
    false
}

/// Test that query audits are correctly stored and retrievable.
///
/// Creates a query audit with an agent ID, stores it, then verifies it can be
/// retrieved via `get_audits_for_agent()`.
pub(crate) async fn run_audit_test<S: KVStore + 'static>(
    store: &Arc<S>,
    agents: &[Agent],
    result: &mut SimulationResult,
) -> bool {
    let audit_store = GenericAuditStore::new(store.clone());
    let agent = &agents[0];
    let agent_id = agent.public_key();

    // Create a query audit
    let timestamp = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);

    let query_id: QueryId = *blake3::hash(b"test_query").as_bytes();

    let audit = QueryAudit {
        query_id,
        agent_id: Some(agent_id),
        timestamp,
        params: QueryParams {
            subject: Some("AuditTest_Entity".to_string()),
            predicate: Some("test_property".to_string()),
            lifecycle: Some(LifecycleStage::Approved),
            epoch: None,
            lens: Some("Recency".to_string()),
        },
        result_hash: Some([1u8; 32]),
        result_confidence: 0.95,
        contributing_assertions: vec![ContributingAssertion {
            assertion_hash: [2u8; 32],
            weight: 1.0,
            source_hash: [3u8; 32],
            lifecycle: LifecycleStage::Approved,
        }],
    };

    // Store the audit
    if let Err(e) = audit_store.put_audit(&audit).await {
        result.errors.push(SimulationError {
            tick: 0,
            kind: ErrorKind::AuditFailure,
            message: format!("Audit test: failed to store audit: {}", e),
        });
        return false;
    }

    // Retrieve audits for this agent
    let audits = match audit_store.get_audits_for_agent(&agent_id, 0, None, 10).await {
        Ok(a) => a,
        Err(e) => {
            result.errors.push(SimulationError {
                tick: 0,
                kind: ErrorKind::AuditFailure,
                message: format!("Audit test: failed to retrieve audits: {}", e),
            });
            return false;
        }
    };

    // Verify we got our audit back
    if audits.is_empty() {
        result.errors.push(SimulationError {
            tick: 0,
            kind: ErrorKind::AuditFailure,
            message: "Audit test: no audits found for agent".to_string(),
        });
        return false;
    }

    let found = &audits[0];
    if found.query_id != query_id {
        result.errors.push(SimulationError {
            tick: 0,
            kind: ErrorKind::AuditFailure,
            message: "Audit test: query_id mismatch".to_string(),
        });
        return false;
    }

    if found.agent_id != Some(agent_id) {
        result.errors.push(SimulationError {
            tick: 0,
            kind: ErrorKind::AuditFailure,
            message: "Audit test: agent_id mismatch".to_string(),
        });
        return false;
    }

    // Verify contributing assertions
    if found.contributing_assertions.len() != 1 {
        result.errors.push(SimulationError {
            tick: 0,
            kind: ErrorKind::AuditFailure,
            message: format!(
                "Audit test: expected 1 contributing assertion, got {}",
                found.contributing_assertions.len()
            ),
        });
        return false;
    }

    debug!("  Audit trail correctly stored and retrieved for agent");
    true
}
