use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use shared_kernel::normalize_required_string;
#[cfg(feature = "spacetime-types")]
use spacetimedb::SpacetimeType;

use crate::{
    AiResultReferenceKind, AiTaskFieldError, AiTaskKind, AiTaskStageBlueprint, AiTaskStageKind,
};

#[cfg_attr(feature = "spacetime-types", derive(SpacetimeType))]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AiTaskCreateInput {
    pub task_id: String,
    pub task_kind: AiTaskKind,
    pub owner_user_id: String,
    pub request_label: String,
    pub source_module: String,
    pub source_entity_id: Option<String>,
    pub request_payload_json: Option<String>,
    pub stages: Vec<AiTaskStageBlueprint>,
    pub created_at_micros: i64,
}

#[cfg_attr(feature = "spacetime-types", derive(SpacetimeType))]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AiTaskStartInput {
    pub task_id: String,
    pub started_at_micros: i64,
}

#[cfg_attr(feature = "spacetime-types", derive(SpacetimeType))]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AiTaskStageStartInput {
    pub task_id: String,
    pub stage_kind: AiTaskStageKind,
    pub started_at_micros: i64,
}

#[cfg_attr(feature = "spacetime-types", derive(SpacetimeType))]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AiTextChunkAppendInput {
    pub task_id: String,
    pub stage_kind: AiTaskStageKind,
    pub sequence: u32,
    pub delta_text: String,
    pub created_at_micros: i64,
}

#[cfg_attr(feature = "spacetime-types", derive(SpacetimeType))]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AiStageCompletionInput {
    pub task_id: String,
    pub stage_kind: AiTaskStageKind,
    pub text_output: Option<String>,
    pub structured_payload_json: Option<String>,
    pub warning_messages: Vec<String>,
    pub completed_at_micros: i64,
}

#[cfg_attr(feature = "spacetime-types", derive(SpacetimeType))]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AiResultReferenceInput {
    pub task_id: String,
    pub reference_kind: AiResultReferenceKind,
    pub reference_id: String,
    pub label: Option<String>,
    pub created_at_micros: i64,
}

#[cfg_attr(feature = "spacetime-types", derive(SpacetimeType))]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AiTaskFinishInput {
    pub task_id: String,
    pub completed_at_micros: i64,
}

#[cfg_attr(feature = "spacetime-types", derive(SpacetimeType))]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AiTaskCancelInput {
    pub task_id: String,
    pub completed_at_micros: i64,
}

#[cfg_attr(feature = "spacetime-types", derive(SpacetimeType))]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AiTaskFailureInput {
    pub task_id: String,
    pub failure_message: String,
    pub completed_at_micros: i64,
}

pub fn validate_task_create_input(input: &AiTaskCreateInput) -> Result<(), AiTaskFieldError> {
    if normalize_required_string(&input.task_id).is_none() {
        return Err(AiTaskFieldError::MissingTaskId);
    }
    if normalize_required_string(&input.owner_user_id).is_none() {
        return Err(AiTaskFieldError::MissingOwnerUserId);
    }
    if normalize_required_string(&input.request_label).is_none() {
        return Err(AiTaskFieldError::MissingRequestLabel);
    }
    if normalize_required_string(&input.source_module).is_none() {
        return Err(AiTaskFieldError::MissingSourceModule);
    }
    if input.stages.is_empty() {
        return Err(AiTaskFieldError::MissingStageBlueprints);
    }

    let mut seen = HashMap::new();
    for stage in &input.stages {
        if normalize_required_string(&stage.label).is_none()
            || normalize_required_string(&stage.detail).is_none()
        {
            return Err(AiTaskFieldError::MissingStageBlueprints);
        }

        if seen.insert(stage.stage_kind, true).is_some() {
            return Err(AiTaskFieldError::DuplicateStageBlueprint);
        }
    }

    Ok(())
}