//! DQN checkpoint implementation — serialization/deserialization of DQN agent state.

use std::collections::HashMap;

use async_trait::async_trait;
use ml_core::{Checkpointable, MLError, ModelType};
use rust_decimal::Decimal;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tracing::{debug, warn};

use crate::agent::DQNAgent;
use crate::DQNConfig;

/// Serializable state for DQN model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DQNCheckpointState {
    pub config: DQNConfig,
    pub epoch: Option<u64>,
    pub step: Option<u64>,
    pub total_episodes: u64,
    pub total_steps: u64,
    pub q_network_weights: Vec<u8>,
    pub target_network_weights: Vec<u8>,
    pub replay_buffer_size: usize,
    pub replay_buffer_capacity: usize,
    pub average_reward: f64,
    pub epsilon: f64,
    pub loss_history: Vec<f64>,
    pub total_inferences: u64,
    pub avg_inference_time_us: f64,
}

#[async_trait]
impl Checkpointable for DQNAgent {
    fn model_type(&self) -> ModelType {
        ModelType::DQN
    }

    fn model_name(&self) -> &str {
        "dqn_agent"
    }

    fn model_version(&self) -> &str {
        "1.0.0"
    }

    async fn serialize_state(&self) -> Result<Vec<u8>, MLError> {
        let state = DQNCheckpointState {
            config: self.config.clone(),
            epoch: None,
            step: None,
            total_episodes: self.metrics.total_episodes,
            total_steps: 0,
            q_network_weights: self
                .extract_network_weights("q_network")
                .unwrap_or_default(),
            target_network_weights: vec![],
            replay_buffer_size: self.get_replay_buffer_size(),
            replay_buffer_capacity: self.config.replay_buffer_capacity,
            average_reward: self.get_average_reward(),
            epsilon: self.config.epsilon_start as f64,
            loss_history: vec![],
            total_inferences: 0,
            avg_inference_time_us: 0.0,
        };

        let serialized = serde_json::to_vec(&state)
            .map_err(|e| MLError::ModelError(format!("DQN serialization failed: {}", e)))?;

        debug!("Serialized DQN state: {} bytes", serialized.len());
        Ok(serialized)
    }

    async fn deserialize_state(&mut self, data: &[u8]) -> Result<(), MLError> {
        let state: DQNCheckpointState = serde_json::from_slice(data)
            .map_err(|e| MLError::ModelError(format!("DQN deserialization failed: {}", e)))?;

        self.config = state.config;
        if !state.q_network_weights.is_empty() {
            if let Err(e) = self.restore_network_weights("q_network", &state.q_network_weights) {
                warn!("Failed to restore Q-network weights: {}", e);
            }
        }

        if !state.target_network_weights.is_empty() {
            if let Err(e) =
                self.restore_network_weights("target_network", &state.target_network_weights)
            {
                warn!("Failed to restore target network weights: {}", e);
            }
        }

        self.metrics.total_episodes = state.total_episodes;
        self.metrics.avg_reward = Decimal::try_from(state.average_reward).unwrap_or(Decimal::ZERO);

        debug!("Deserialized DQN state from {} bytes", data.len());
        Ok(())
    }

    fn get_training_state(&self) -> (Option<u64>, Option<u64>, Option<f64>, Option<f64>) {
        let current_epoch = None;
        let total_episodes = Some(self.metrics.total_episodes);
        let average_reward =
            Some(TryInto::<f64>::try_into(self.metrics.avg_reward).unwrap_or(0.0));
        let loss = Some(TryInto::<f64>::try_into(self.metrics.current_loss).unwrap_or(0.0));
        (current_epoch, total_episodes, average_reward, loss)
    }

    fn get_hyperparameters(&self) -> HashMap<String, Value> {
        let mut params = HashMap::new();
        params.insert(
            "learning_rate".to_owned(),
            Value::from(self.config.learning_rate),
        );
        params.insert(
            "discount_factor".to_owned(),
            Value::from(self.config.gamma),
        );
        params.insert(
            "epsilon".to_owned(),
            Value::from(self.config.epsilon_start),
        );
        params.insert(
            "epsilon_decay".to_owned(),
            Value::from(self.config.epsilon_decay),
        );
        params.insert(
            "epsilon_min".to_owned(),
            Value::from(self.config.epsilon_end),
        );
        params.insert(
            "replay_buffer_size".to_owned(),
            Value::from(self.config.replay_buffer_capacity),
        );
        params.insert(
            "batch_size".to_owned(),
            Value::from(self.config.batch_size),
        );
        params.insert(
            "target_update_frequency".to_owned(),
            Value::from(self.config.target_update_freq),
        );
        params
    }

    fn get_metrics(&self) -> HashMap<String, f64> {
        let mut metrics = HashMap::new();
        metrics.insert(
            "total_episodes".to_owned(),
            self.metrics.total_episodes as f64,
        );
        metrics.insert(
            "average_reward".to_owned(),
            TryInto::<f64>::try_into(self.metrics.avg_reward).unwrap_or(0.0),
        );
        metrics.insert(
            "current_loss".to_owned(),
            TryInto::<f64>::try_into(self.metrics.current_loss).unwrap_or(0.0),
        );
        metrics.insert("epsilon".to_owned(), self.config.epsilon_start as f64);
        metrics.insert(
            "replay_buffer_size".to_owned(),
            self.get_replay_buffer_size() as f64,
        );
        metrics.insert(
            "exploration_rate".to_owned(),
            self.config.epsilon_start as f64,
        );
        metrics
    }

    fn get_architecture_info(&self) -> HashMap<String, Value> {
        let mut info = HashMap::new();
        info.insert("network_type".to_owned(), Value::from("DQN"));
        info.insert("input_size".to_owned(), Value::from(self.config.state_dim));
        info.insert(
            "hidden_size".to_owned(),
            Value::from(self.config.hidden_dims.len()),
        );
        info.insert(
            "output_size".to_owned(),
            Value::from(self.config.num_actions),
        );
        info.insert("num_hidden_layers".to_owned(), Value::from(2));
        info.insert("activation".to_owned(), Value::from("ReLU"));
        info
    }
}
