//! Core data models for AI decision tracking and observability
//!
//! This module contains the fundamental data structures used throughout the Briefcase AI system:
//! - [`Input`] and [`Output`] for capturing AI function parameters and results
//! - [`DecisionSnapshot`] for recording complete AI decision context
//! - [`ModelParameters`] for tracking AI model configuration
//! - [`ExecutionContext`] for environmental reproducibility
//! - [`Snapshot`] for grouping related decisions
//!
//! ## Example Usage
//!
//! ```rust
//! use briefcase_core::models::*;
//! use serde_json::json;
//!
//! // Create input and output data
//! let input = Input::new("prompt", json!("What is AI?"), "string");
//! let output = Output::new("response", json!("AI is..."), "string")
//!     .with_confidence(0.92);
//!
//! // Create a decision snapshot
//! let decision = DecisionSnapshot::new("gpt_query")
//!     .add_input(input)
//!     .add_output(output)
//!     .with_execution_time(234.5);
//! ```

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use uuid::Uuid;

/// Input parameter to an AI decision point
///
/// Represents a single input parameter passed to an AI function or model.
/// Inputs are captured with their name, value, and type information for
/// complete reproducibility.
///
/// # Examples
///
/// ```rust
/// use briefcase_core::models::Input;
/// use serde_json::json;
///
/// let input = Input::new("user_query", json!("Hello world"), "string");
/// assert_eq!(input.name, "user_query");
/// assert_eq!(input.data_type, "string");
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Input {
    pub name: String,
    pub value: serde_json::Value,
    pub data_type: String,
    #[serde(default = "default_schema_version")]
    pub schema_version: String,
}

impl Input {
    pub fn new(
        name: impl Into<String>,
        value: serde_json::Value,
        data_type: impl Into<String>,
    ) -> Self {
        Self {
            name: name.into(),
            value,
            data_type: data_type.into(),
            schema_version: default_schema_version(),
        }
    }
}

/// Output result from an AI decision point
///
/// Represents the result produced by an AI function or model.
/// Outputs can include confidence scores and are captured with
/// type information for analysis and replay.
///
/// # Examples
///
/// ```rust
/// use briefcase_core::models::Output;
/// use serde_json::json;
///
/// let output = Output::new("response", json!("Hello back!"), "string")
///     .with_confidence(0.95);
/// assert_eq!(output.confidence, Some(0.95));
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Output {
    pub name: String,
    pub value: serde_json::Value,
    pub data_type: String,
    pub confidence: Option<f64>, // 0.0 - 1.0
    #[serde(default = "default_schema_version")]
    pub schema_version: String,
}

impl Output {
    pub fn new(
        name: impl Into<String>,
        value: serde_json::Value,
        data_type: impl Into<String>,
    ) -> Self {
        Self {
            name: name.into(),
            value,
            data_type: data_type.into(),
            confidence: None,
            schema_version: default_schema_version(),
        }
    }

    pub fn with_confidence(mut self, confidence: f64) -> Self {
        self.confidence = Some(confidence.clamp(0.0, 1.0));
        self
    }
}

/// AI model parameters for reproducibility and tracking
///
/// Captures the configuration of an AI model including its name, version,
/// provider, and all parameters needed for reproducible execution.
///
/// # Examples
///
/// ```rust
/// use briefcase_core::models::ModelParameters;
/// use serde_json::json;
///
/// let params = ModelParameters::new("gpt-4")
///     .with_provider("openai")
///     .with_parameter("temperature", json!(0.7))
///     .with_hyperparameter("max_tokens", json!(1000));
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ModelParameters {
    pub model_name: String,
    pub model_version: Option<String>,
    pub provider: Option<String>, // openai, anthropic, etc.
    #[serde(default)]
    pub parameters: HashMap<String, serde_json::Value>,
    #[serde(default)]
    pub hyperparameters: HashMap<String, serde_json::Value>,
    pub weights_hash: Option<String>, // For reproducibility verification
}

impl ModelParameters {
    pub fn new(model_name: impl Into<String>) -> Self {
        Self {
            model_name: model_name.into(),
            model_version: None,
            provider: None,
            parameters: HashMap::new(),
            hyperparameters: HashMap::new(),
            weights_hash: None,
        }
    }

    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
        self.provider = Some(provider.into());
        self
    }

    pub fn with_version(mut self, version: impl Into<String>) -> Self {
        self.model_version = Some(version.into());
        self
    }

    pub fn with_parameter(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
        self.parameters.insert(key.into(), value);
        self
    }

    pub fn with_hyperparameter(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
        self.hyperparameters.insert(key.into(), value);
        self
    }

    pub fn with_weights_hash(mut self, hash: impl Into<String>) -> Self {
        self.weights_hash = Some(hash.into());
        self
    }
}

/// Execution context for deterministic replay
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct ExecutionContext {
    pub runtime_version: Option<String>, // Python 3.11, Node 20, etc.
    #[serde(default)]
    pub dependencies: HashMap<String, String>,
    pub random_seed: Option<i64>,
    #[serde(default)]
    pub environment_variables: HashMap<String, String>,
    #[serde(default)]
    pub hardware_info: HashMap<String, serde_json::Value>,
}

impl ExecutionContext {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_runtime_version(mut self, version: impl Into<String>) -> Self {
        self.runtime_version = Some(version.into());
        self
    }

    pub fn with_dependency(mut self, name: impl Into<String>, version: impl Into<String>) -> Self {
        self.dependencies.insert(name.into(), version.into());
        self
    }

    pub fn with_random_seed(mut self, seed: i64) -> Self {
        self.random_seed = Some(seed);
        self
    }

    pub fn with_env_var(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.environment_variables.insert(key.into(), value.into());
        self
    }
}

/// Metadata for a snapshot
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct SnapshotMetadata {
    pub snapshot_id: Uuid,
    pub timestamp: DateTime<Utc>,
    #[serde(default = "default_schema_version")]
    pub schema_version: String,
    pub sdk_version: String,
    pub created_by: Option<String>,
    pub checksum: Option<String>,
}

impl SnapshotMetadata {
    pub fn new() -> Self {
        Self {
            snapshot_id: Uuid::new_v4(),
            timestamp: Utc::now(),
            schema_version: default_schema_version(),
            sdk_version: env!("CARGO_PKG_VERSION").to_string(),
            created_by: None,
            checksum: None,
        }
    }

    pub fn with_created_by(mut self, created_by: impl Into<String>) -> Self {
        self.created_by = Some(created_by.into());
        self
    }

    pub fn compute_checksum(&mut self, data: &[u8]) {
        let mut hasher = Sha256::new();
        hasher.update(data);
        let result = hasher.finalize();
        self.checksum = Some(format!("{:x}", result));
    }
}

impl Default for SnapshotMetadata {
    fn default() -> Self {
        Self::new()
    }
}

/// A single AI decision capture
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct DecisionSnapshot {
    pub metadata: SnapshotMetadata,
    pub context: ExecutionContext,
    pub function_name: String,
    pub module_name: Option<String>,
    pub inputs: Vec<Input>,
    pub outputs: Vec<Output>,
    pub model_parameters: Option<ModelParameters>,
    pub execution_time_ms: Option<f64>,
    pub error: Option<String>,
    pub error_type: Option<String>,
    #[serde(default)]
    pub tags: HashMap<String, String>,
    #[serde(default)]
    pub custom_data: HashMap<String, serde_json::Value>,
}

impl DecisionSnapshot {
    pub fn new(function_name: impl Into<String>) -> Self {
        let metadata = SnapshotMetadata::new();
        let snapshot = Self {
            metadata,
            context: ExecutionContext::new(),
            function_name: function_name.into(),
            module_name: None,
            inputs: Vec::new(),
            outputs: Vec::new(),
            model_parameters: None,
            execution_time_ms: None,
            error: None,
            error_type: None,
            tags: HashMap::new(),
            custom_data: HashMap::new(),
        };

        // Update checksum after creation
        let mut result = snapshot;
        result.update_checksum();
        result
    }

    pub fn with_module(mut self, module_name: impl Into<String>) -> Self {
        self.module_name = Some(module_name.into());
        self.update_checksum();
        self
    }

    pub fn with_context(mut self, context: ExecutionContext) -> Self {
        self.context = context;
        self.update_checksum();
        self
    }

    pub fn add_input(mut self, input: Input) -> Self {
        self.inputs.push(input);
        self.update_checksum();
        self
    }

    pub fn add_output(mut self, output: Output) -> Self {
        self.outputs.push(output);
        self.update_checksum();
        self
    }

    pub fn with_model_parameters(mut self, params: ModelParameters) -> Self {
        self.model_parameters = Some(params);
        self.update_checksum();
        self
    }

    pub fn with_execution_time(mut self, time_ms: f64) -> Self {
        self.execution_time_ms = Some(time_ms);
        self.update_checksum();
        self
    }

    pub fn with_error(mut self, error: impl Into<String>, error_type: Option<String>) -> Self {
        self.error = Some(error.into());
        self.error_type = error_type;
        self.update_checksum();
        self
    }

    pub fn add_tag(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.tags.insert(key.into(), value.into());
        self.update_checksum();
        self
    }

    pub fn add_custom_data(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
        self.custom_data.insert(key.into(), value);
        self.update_checksum();
        self
    }

    fn update_checksum(&mut self) {
        if let Ok(json_bytes) = serde_json::to_vec(self) {
            self.metadata.compute_checksum(&json_bytes);
        }
    }

    /// Serialize to canonical JSON for consistent checksums
    pub fn to_canonical_json(&self) -> Result<String, serde_json::Error> {
        // Create a copy without the checksum for canonical representation
        let mut copy = self.clone();
        copy.metadata.checksum = None;

        // Use sorted keys for consistent JSON
        let value = serde_json::to_value(&copy)?;
        serde_json::to_string(&value)
    }
}

/// Root snapshot containing multiple decisions (e.g., a session)
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Snapshot {
    pub metadata: SnapshotMetadata,
    pub decisions: Vec<DecisionSnapshot>,
    pub snapshot_type: SnapshotType,
}

impl Snapshot {
    pub fn new(snapshot_type: SnapshotType) -> Self {
        let metadata = SnapshotMetadata::new();
        let snapshot = Self {
            metadata,
            decisions: Vec::new(),
            snapshot_type,
        };

        let mut result = snapshot;
        result.update_checksum();
        result
    }

    pub fn add_decision(&mut self, decision: DecisionSnapshot) {
        self.decisions.push(decision);
        self.update_checksum();
    }

    pub fn with_created_by(mut self, created_by: impl Into<String>) -> Self {
        self.metadata.created_by = Some(created_by.into());
        self.update_checksum();
        self
    }

    fn update_checksum(&mut self) {
        if let Ok(json_bytes) = serde_json::to_vec(self) {
            self.metadata.compute_checksum(&json_bytes);
        }
    }

    /// Serialize to canonical JSON for consistent checksums
    pub fn to_canonical_json(&self) -> Result<String, serde_json::Error> {
        let mut copy = self.clone();
        copy.metadata.checksum = None;

        let value = serde_json::to_value(&copy)?;
        serde_json::to_string(&value)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum SnapshotType {
    Decision,
    Batch,
    Session,
}

fn default_schema_version() -> String {
    "1.0".to_string()
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn test_input_creation() {
        let input = Input::new("test", json!("value"), "string");
        assert_eq!(input.name, "test");
        assert_eq!(input.value, json!("value"));
        assert_eq!(input.data_type, "string");
        assert_eq!(input.schema_version, "1.0");
    }

    #[test]
    fn test_output_with_confidence() {
        let output = Output::new("result", json!(42), "number").with_confidence(0.95);
        assert_eq!(output.confidence, Some(0.95));
    }

    #[test]
    fn test_output_confidence_clamping() {
        let output = Output::new("result", json!(42), "number").with_confidence(1.5);
        assert_eq!(output.confidence, Some(1.0));

        let output = Output::new("result", json!(42), "number").with_confidence(-0.5);
        assert_eq!(output.confidence, Some(0.0));
    }

    #[test]
    fn test_model_parameters_builder() {
        let params = ModelParameters::new("gpt-4")
            .with_provider("openai")
            .with_version("1.0")
            .with_parameter("temperature", json!(0.7))
            .with_hyperparameter("max_tokens", json!(1000));

        assert_eq!(params.model_name, "gpt-4");
        assert_eq!(params.provider, Some("openai".to_string()));
        assert_eq!(params.model_version, Some("1.0".to_string()));
        assert_eq!(params.parameters.get("temperature"), Some(&json!(0.7)));
        assert_eq!(params.hyperparameters.get("max_tokens"), Some(&json!(1000)));
    }

    #[test]
    fn test_execution_context_builder() {
        let context = ExecutionContext::new()
            .with_runtime_version("Python 3.11")
            .with_dependency("numpy", "1.24.0")
            .with_random_seed(42)
            .with_env_var("DEBUG", "true");

        assert_eq!(context.runtime_version, Some("Python 3.11".to_string()));
        assert_eq!(
            context.dependencies.get("numpy"),
            Some(&"1.24.0".to_string())
        );
        assert_eq!(context.random_seed, Some(42));
        assert_eq!(
            context.environment_variables.get("DEBUG"),
            Some(&"true".to_string())
        );
    }

    #[test]
    fn test_decision_snapshot_creation() {
        let snapshot = DecisionSnapshot::new("test_function");
        assert_eq!(snapshot.function_name, "test_function");
        assert!(snapshot.metadata.checksum.is_some());
    }

    #[test]
    fn test_decision_snapshot_builder() {
        let input = Input::new("x", json!(1), "number");
        let output = Output::new("result", json!(2), "number");
        let params = ModelParameters::new("gpt-4");

        let snapshot = DecisionSnapshot::new("add")
            .with_module("math")
            .add_input(input)
            .add_output(output)
            .with_model_parameters(params)
            .with_execution_time(100.5)
            .add_tag("version", "1.0");

        assert_eq!(snapshot.function_name, "add");
        assert_eq!(snapshot.module_name, Some("math".to_string()));
        assert_eq!(snapshot.inputs.len(), 1);
        assert_eq!(snapshot.outputs.len(), 1);
        assert!(snapshot.model_parameters.is_some());
        assert_eq!(snapshot.execution_time_ms, Some(100.5));
        assert_eq!(snapshot.tags.get("version"), Some(&"1.0".to_string()));
    }

    #[test]
    fn test_snapshot_creation() {
        let mut snapshot = Snapshot::new(SnapshotType::Session);
        let decision = DecisionSnapshot::new("test");

        snapshot.add_decision(decision);

        assert_eq!(snapshot.snapshot_type, SnapshotType::Session);
        assert_eq!(snapshot.decisions.len(), 1);
        assert!(snapshot.metadata.checksum.is_some());
    }

    #[test]
    fn test_json_serialization_roundtrip() {
        let input = Input::new("test", json!({"key": "value"}), "object");
        let output = Output::new("result", json!([1, 2, 3]), "array");
        let params = ModelParameters::new("gpt-4")
            .with_provider("openai")
            .with_parameter("temperature", json!(0.8));

        let snapshot = DecisionSnapshot::new("process")
            .add_input(input)
            .add_output(output)
            .with_model_parameters(params);

        let json = serde_json::to_string(&snapshot).unwrap();
        let deserialized: DecisionSnapshot = serde_json::from_str(&json).unwrap();

        assert_eq!(snapshot, deserialized);
    }

    #[test]
    fn test_canonical_json_consistency() {
        let snapshot = DecisionSnapshot::new("test")
            .add_tag("key1", "value1")
            .add_tag("key2", "value2");

        let json1 = snapshot.to_canonical_json().unwrap();
        let json2 = snapshot.to_canonical_json().unwrap();

        assert_eq!(json1, json2);
    }

    #[test]
    fn test_checksum_updates() {
        let mut snapshot = DecisionSnapshot::new("test");
        let initial_checksum = snapshot.metadata.checksum.clone();

        snapshot = snapshot.add_tag("new", "tag");
        let updated_checksum = snapshot.metadata.checksum.clone();

        assert_ne!(initial_checksum, updated_checksum);
    }
}
