use chrono::Utc;
use surrealdb::{engine::remote::ws::Client, Surreal};

use crate::error::{A2aError, Result};
use crate::metrics::A2A_DB_OPERATIONS;
use crate::protocol::{A2aErrorObj, A2aTask, A2aTaskResult, A2aTaskStatus, TaskState};

#[derive(Clone)]
pub struct TaskManager {
    db: Surreal<Client>,
}

impl TaskManager {
    pub fn new(db: Surreal<Client>) -> Self {
        Self { db }
    }

    pub async fn create(&self, task: A2aTask) -> Result<A2aTaskStatus> {
        let now = Utc::now().to_rfc3339();

        // Pre-serialize complex fields with serde_json to avoid SurrealDB's
        // own serializer choking on adjacently-tagged enums (tag + content).
        let message_json = serde_json::to_value(&task.message)
            .map_err(|e| A2aError::InternalError(format!("Failed to serialize message: {}", e)))?;
        let metadata_json = serde_json::to_value(&task.metadata)
            .map_err(|e| A2aError::InternalError(format!("Failed to serialize metadata: {}", e)))?;

        match self
            .db
            .query(
                "INSERT INTO a2a_tasks (task_id, state, message, result, error, metadata, \
                 created_at, updated_at) VALUES ($task_id, $state, $message, NONE, NONE, \
                 $metadata, $created_at, $updated_at)",
            )
            .bind(("task_id", task.id.clone()))
            .bind(("state", TaskState::Waiting.as_str().to_string()))
            .bind(("message", message_json))
            .bind(("metadata", metadata_json))
            .bind(("created_at", now.clone()))
            .bind(("updated_at", now.clone()))
            .await
        {
            Ok(_) => {
                A2A_DB_OPERATIONS
                    .with_label_values(&["create", "success"])
                    .inc();
                Ok(A2aTaskStatus {
                    id: task.id,
                    state: TaskState::Waiting.as_str().to_string(),
                    message: Some(task.message),
                    result: None,
                    error: None,
                    created_at: now.clone(),
                    updated_at: now,
                })
            }
            Err(e) => {
                A2A_DB_OPERATIONS
                    .with_label_values(&["create", "error"])
                    .inc();
                Err(A2aError::InternalError(format!(
                    "Failed to create task: {}",
                    e
                )))
            }
        }
    }

    pub async fn get(&self, id: &str) -> Result<A2aTaskStatus> {
        // Explicit projection: exclude `id` (SurrealDB Thing) and cast datetime
        // fields to strings to avoid deserialization failures into serde_json::Value.
        let mut response = match self
            .db
            .query(
                "SELECT task_id, state, message, result, error, metadata, \
                 type::string(created_at) AS created_at, type::string(updated_at) AS updated_at \
                 FROM a2a_tasks WHERE task_id = $task_id LIMIT 1",
            )
            .bind(("task_id", id.to_string()))
            .await
        {
            Ok(r) => {
                A2A_DB_OPERATIONS
                    .with_label_values(&["query", "success"])
                    .inc();
                r
            }
            Err(e) => {
                A2A_DB_OPERATIONS
                    .with_label_values(&["query", "error"])
                    .inc();
                return Err(A2aError::InternalError(format!("Query failed: {}", e)));
            }
        };

        let records: Vec<serde_json::Value> = response
            .take(0)
            .map_err(|e| A2aError::InternalError(format!("Failed to extract result: {}", e)))?;

        let record = records
            .into_iter()
            .next()
            .ok_or_else(|| A2aError::TaskNotFound(id.to_string()))?;

        // Deserialize from DB record
        let task_id = record
            .get("task_id")
            .and_then(|v| v.as_str())
            .ok_or_else(|| A2aError::InternalError("Missing task_id".to_string()))?
            .to_string();

        let state = record
            .get("state")
            .and_then(|v| v.as_str())
            .ok_or_else(|| A2aError::InternalError("Missing state".to_string()))?
            .to_string();

        let message = record
            .get("message")
            .and_then(|v| serde_json::from_value(v.clone()).ok());

        let result = record
            .get("result")
            .filter(|v| !v.is_null())
            .and_then(|v| serde_json::from_value(v.clone()).ok());

        let error = record
            .get("error")
            .filter(|v| !v.is_null())
            .and_then(|v| serde_json::from_value(v.clone()).ok());

        let created_at = record
            .get("created_at")
            .and_then(|v| v.as_str())
            .ok_or_else(|| A2aError::InternalError("Missing created_at".to_string()))?
            .to_string();

        let updated_at = record
            .get("updated_at")
            .and_then(|v| v.as_str())
            .ok_or_else(|| A2aError::InternalError("Missing updated_at".to_string()))?
            .to_string();

        Ok(A2aTaskStatus {
            id: task_id,
            state,
            message,
            result,
            error,
            created_at,
            updated_at,
        })
    }

    pub async fn update_state(&self, id: &str, new_state: TaskState) -> Result<()> {
        let now = Utc::now().to_rfc3339();

        match self
            .db
            .query(
                "UPDATE a2a_tasks SET state = $state, updated_at = $now WHERE task_id = $task_id",
            )
            .bind(("task_id", id.to_string()))
            .bind(("state", new_state.as_str().to_string()))
            .bind(("now", now))
            .await
        {
            Ok(_) => {
                A2A_DB_OPERATIONS
                    .with_label_values(&["update", "success"])
                    .inc();
                Ok(())
            }
            Err(e) => {
                A2A_DB_OPERATIONS
                    .with_label_values(&["update", "error"])
                    .inc();
                Err(A2aError::InternalError(format!(
                    "Failed to update state: {}",
                    e
                )))
            }
        }
    }

    pub async fn complete(&self, id: &str, result: A2aTaskResult) -> Result<()> {
        let now = Utc::now().to_rfc3339();
        let result_json = serde_json::to_value(&result)
            .map_err(|e| A2aError::InternalError(format!("Failed to serialize result: {}", e)))?;

        match self
            .db
            .query(
                "UPDATE a2a_tasks SET state = $state, result = $result, updated_at = $now WHERE \
                 task_id = $task_id",
            )
            .bind(("task_id", id.to_string()))
            .bind(("state", TaskState::Completed.as_str().to_string()))
            .bind(("result", result_json))
            .bind(("now", now))
            .await
        {
            Ok(_) => {
                A2A_DB_OPERATIONS
                    .with_label_values(&["update", "success"])
                    .inc();
                Ok(())
            }
            Err(e) => {
                A2A_DB_OPERATIONS
                    .with_label_values(&["update", "error"])
                    .inc();
                Err(A2aError::InternalError(format!(
                    "Failed to complete task: {}",
                    e
                )))
            }
        }
    }

    pub async fn fail(&self, id: &str, error: A2aErrorObj) -> Result<()> {
        let now = Utc::now().to_rfc3339();
        let error_json = serde_json::to_value(&error)
            .map_err(|e| A2aError::InternalError(format!("Failed to serialize error: {}", e)))?;

        match self
            .db
            .query(
                "UPDATE a2a_tasks SET state = $state, error = $error, updated_at = $now WHERE \
                 task_id = $task_id",
            )
            .bind(("task_id", id.to_string()))
            .bind(("state", TaskState::Failed.as_str().to_string()))
            .bind(("error", error_json))
            .bind(("now", now))
            .await
        {
            Ok(_) => {
                A2A_DB_OPERATIONS
                    .with_label_values(&["update", "success"])
                    .inc();
                Ok(())
            }
            Err(e) => {
                A2A_DB_OPERATIONS
                    .with_label_values(&["update", "error"])
                    .inc();
                Err(A2aError::InternalError(format!(
                    "Failed to fail task: {}",
                    e
                )))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::protocol::{A2aMessage, A2aMessagePart};

    // Note: These tests require a running SurrealDB instance
    // Run: docker run -p 8000:8000 surrealdb/surrealdb:latest start --bind
    // 0.0.0.0:8000

    #[tokio::test]
    #[ignore] // Requires SurrealDB running
    async fn test_task_lifecycle() {
        let db = Surreal::new::<surrealdb::engine::remote::ws::Ws>("127.0.0.1:8000")
            .await
            .unwrap();

        db.signin(surrealdb::opt::auth::Root {
            username: "root".to_string(),
            password: "root".to_string(),
        })
        .await
        .unwrap();

        db.use_ns("test").use_db("vapora_a2a_test").await.unwrap();

        let manager = TaskManager::new(db);

        let task = A2aTask {
            id: "task-1".to_string(),
            message: A2aMessage {
                role: "user".to_string(),
                parts: vec![A2aMessagePart::Text("Test".to_string())],
            },
            metadata: Default::default(),
        };

        let created = manager.create(task).await.unwrap();
        assert_eq!(created.state, "waiting");

        manager
            .update_state("task-1", TaskState::Working)
            .await
            .unwrap();

        let updated = manager.get("task-1").await.unwrap();
        assert_eq!(updated.state, "working");
    }
}
