use std::collections::BTreeMap;

use libmcp::{Fault, FaultClass, FaultCode, Generation, RecoveryDirective, ToolErrorDetail};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};

#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
pub(crate) struct FaultContext {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) consult: Option<ConsultFaultContext>,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub(crate) struct ConsultFaultContext {
    pub(crate) cwd: String,
    pub(crate) context_mode: String,
    pub(crate) planned_session_id: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) reused_session_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) observed_session_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) resume_session_id: Option<String>,
    #[serde(default, skip_serializing_if = "is_false")]
    pub(crate) quota_limited: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) quota_reset_hint: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) retry_hint: Option<String>,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub(crate) enum FaultStage {
    Host,
    Worker,
    Claude,
    Transport,
    Protocol,
    Rollout,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub(crate) struct FaultRecord {
    pub(crate) fault: Fault,
    pub(crate) stage: FaultStage,
    pub(crate) operation: String,
    pub(crate) jsonrpc_code: i64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) context: Option<Box<FaultContext>>,
    pub(crate) retryable: bool,
    pub(crate) retried: bool,
}

impl FaultRecord {
    pub(crate) fn invalid_input(
        generation: Generation,
        stage: FaultStage,
        operation: impl Into<String>,
        detail: impl Into<String>,
    ) -> Self {
        Self::new(
            generation,
            FaultClass::Protocol,
            "invalid_input",
            RecoveryDirective::AbortRequest,
            stage,
            operation,
            detail,
            -32602,
        )
    }

    pub(crate) fn not_initialized(
        generation: Generation,
        stage: FaultStage,
        operation: impl Into<String>,
        detail: impl Into<String>,
    ) -> Self {
        Self::new(
            generation,
            FaultClass::Protocol,
            "not_initialized",
            RecoveryDirective::AbortRequest,
            stage,
            operation,
            detail,
            -32002,
        )
    }

    pub(crate) fn transport(
        generation: Generation,
        stage: FaultStage,
        operation: impl Into<String>,
        detail: impl Into<String>,
    ) -> Self {
        Self::new(
            generation,
            FaultClass::Transport,
            "transport_failure",
            RecoveryDirective::RestartAndReplay,
            stage,
            operation,
            detail,
            -32603,
        )
    }

    pub(crate) fn process(
        generation: Generation,
        stage: FaultStage,
        operation: impl Into<String>,
        detail: impl Into<String>,
    ) -> Self {
        Self::new(
            generation,
            FaultClass::Process,
            "process_failure",
            RecoveryDirective::RestartAndReplay,
            stage,
            operation,
            detail,
            -32603,
        )
    }

    pub(crate) fn downstream(
        generation: Generation,
        stage: FaultStage,
        operation: impl Into<String>,
        detail: impl Into<String>,
    ) -> Self {
        Self::new(
            generation,
            FaultClass::Downstream,
            "downstream_failure",
            RecoveryDirective::AbortRequest,
            stage,
            operation,
            detail,
            -32603,
        )
    }

    pub(crate) fn internal(
        generation: Generation,
        stage: FaultStage,
        operation: impl Into<String>,
        detail: impl Into<String>,
    ) -> Self {
        Self::new(
            generation,
            FaultClass::Invariant,
            "internal_failure",
            RecoveryDirective::AbortRequest,
            stage,
            operation,
            detail,
            -32603,
        )
    }

    pub(crate) fn rollout(
        generation: Generation,
        operation: impl Into<String>,
        detail: impl Into<String>,
    ) -> Self {
        Self::new(
            generation,
            FaultClass::Rollout,
            "rollout_failure",
            RecoveryDirective::RestartAndReplay,
            FaultStage::Rollout,
            operation,
            detail,
            -32603,
        )
    }

    pub(crate) fn mark_retried(mut self) -> Self {
        self.retried = true;
        self
    }

    pub(crate) fn with_context(mut self, context: FaultContext) -> Self {
        self.context = Some(Box::new(context));
        self
    }

    pub(crate) fn message(&self) -> &str {
        self.fault.detail.as_str()
    }

    fn rendered_message(&self) -> String {
        let mut lines = vec![self.message().to_owned()];
        let Some(consult) = self
            .context
            .as_ref()
            .and_then(|context| context.consult.as_ref())
        else {
            return lines.join("\n");
        };
        let mut fields: BTreeMap<String, String> = BTreeMap::from([
            ("cwd".to_owned(), consult.cwd.clone()),
            ("context_mode".to_owned(), consult.context_mode.clone()),
            (
                "planned_session".to_owned(),
                consult.planned_session_id.clone(),
            ),
        ]);
        if let Some(session_id) = consult.reused_session_id.as_ref() {
            let _ = fields.insert("reused_session".to_owned(), session_id.clone());
        }
        if let Some(session_id) = consult.observed_session_id.as_ref() {
            let _ = fields.insert("observed_session".to_owned(), session_id.clone());
        }
        if let Some(session_id) = consult.resume_session_id.as_ref() {
            let _ = fields.insert("resume_session".to_owned(), session_id.clone());
        }
        if consult.quota_limited {
            let _ = fields.insert("quota_limited".to_owned(), "true".to_owned());
        }
        if let Some(reset_hint) = consult.quota_reset_hint.as_ref() {
            let _ = fields.insert("quota_reset".to_owned(), reset_hint.clone());
        }
        if let Some(retry_hint) = consult.retry_hint.as_ref() {
            let _ = fields.insert("retry_hint".to_owned(), retry_hint.clone());
        }
        lines.extend(
            fields
                .into_iter()
                .map(|(key, value)| format!("{key}: {value}")),
        );
        lines.join("\n")
    }

    pub(crate) fn error_detail(&self) -> ToolErrorDetail {
        ToolErrorDetail {
            code: Some(self.jsonrpc_code),
            kind: Some(self.fault.code.as_str().to_owned()),
            message: Some(self.rendered_message()),
        }
    }

    pub(crate) fn into_jsonrpc_error(self) -> Value {
        json!({
            "code": self.jsonrpc_code,
            "message": self.message(),
            "data": self,
        })
    }

    pub(crate) fn into_tool_result(self) -> Value {
        let rendered_message = self.rendered_message();
        json!({
            "content": [{
                "type": "text",
                "text": rendered_message,
            }],
            "structuredContent": self,
            "isError": true,
        })
    }

    fn new(
        generation: Generation,
        class: FaultClass,
        code: &'static str,
        directive: RecoveryDirective,
        stage: FaultStage,
        operation: impl Into<String>,
        detail: impl Into<String>,
        jsonrpc_code: i64,
    ) -> Self {
        let fault = Fault::new(generation, class, fault_code(code), directive, detail);
        Self {
            retryable: directive != RecoveryDirective::AbortRequest,
            fault,
            stage,
            operation: operation.into(),
            jsonrpc_code,
            context: None,
            retried: false,
        }
    }
}

const fn is_false(value: &bool) -> bool {
    !*value
}

fn fault_code(code: &'static str) -> FaultCode {
    match FaultCode::try_new(code.to_owned()) {
        Ok(value) => value,
        Err(_) => std::process::abort(),
    }
}