use pyo3::prelude::*;
use pyo3::FromPyObject;

use frisk_core::policy_engine::PolicyEngine;
use frisk_core::policy_types::ToolPolicyRecord;
use py_dict_carrier::PyDictCarrier;

use crate::redact_dictionary::redact_dictionary as redact_dictionary_fn;
use frisk_core::auth_token::AuthToken;
use frisk_core::otel::otel_manager::OtelManager;
use frisk_core::policy_manager::PolicyManager;
use frisk_core::ProcessToolCallResult;
use opentelemetry::global;
use pyo3::exceptions::PyValueError;
use pyo3::types::PyDict;
use serde_json::Value;
use std::sync::{Arc, Mutex, OnceLock};
use tokio::runtime::{Builder, Runtime};
use tracing_opentelemetry::OpenTelemetrySpanExt;

pub mod py_dict_carrier;
mod redact_dictionary;

// todo: Map Rust exceptions to Python exceptions. https://linear.app/friskai/issue/POL-106/polish-map-rust-exceptions-to-pythonic-exceptions
static POLICY_REFRESH_INTERVAL_SECONDS: u64 = 10;

pub const SPAN_NAME_ENGINE_DECIDE_TOOL_CALL: &str = "frisk.engine.decide_tool_call";

#[pyclass]
#[derive(Debug, Clone)]
pub struct ProcessToolCallResultPy {
    #[pyo3(get)]
    decision: String,
    #[pyo3(get)]
    rules_matched_count: usize,
    #[pyo3(get)]
    reason: String,
}

impl From<ProcessToolCallResult> for ProcessToolCallResultPy {
    fn from(r: ProcessToolCallResult) -> Self {
        match r {
            ProcessToolCallResult::Allow {
                rules_matched,
                reason,
            } => Self {
                decision: "allow".into(),
                rules_matched_count: rules_matched.len(),
                reason,
            },
            ProcessToolCallResult::Deny {
                rules_matched,
                reason,
            } => Self {
                decision: "deny".into(),
                rules_matched_count: rules_matched.len(),
                reason,
            },
        }
    }
}

#[pymethods]
impl ProcessToolCallResultPy {
    pub fn __repr__(&self) -> String {
        format!(
            "<ProcessToolCallResult decision='{}' matched={} reason={:?}>",
            self.decision, self.rules_matched_count, self.reason
        )
    }
}

#[pyclass]
#[derive(Debug, Clone)]
pub struct FriskHandle {
    pub policy_engine: Arc<Mutex<PolicyEngine>>,
    #[allow(dead_code)]
    runtime: Arc<Runtime>,
    otel_manager: Arc<OtelManager>,
    #[allow(dead_code)]
    policy_manager: Arc<Mutex<PolicyManager>>,
    // Store constructor args to validate subsequent calls
    auth_token: Arc<AuthToken>,
    frisk_base_url: String,
    otel_collector_endpoint: String,
}

// Process-wide singleton holder
static FRISK_HANDLE_SINGLETON: OnceLock<FriskHandle> = OnceLock::new();

#[derive(Debug, Clone, FromPyObject)]
/// Python-side options for processing a tool call.
///
/// This struct is constructed from a Python `dict` and passed into the
/// policy engine. Its redaction-related fields accept, on the Python side,
/// either:
/// - a `bool` (e.g. `True` to enable default redaction / `False` to disable), or
/// - a `list[str]` of field names to be selectively redacted.
pub struct ProcessToolCallOptions<'py> {
    /// Redaction configuration for tool arguments.
    ///
    /// From Python this may be provided as either:
    /// - a `bool`, or
    /// - a `list[str]` specifying which argument keys should be redacted.
    #[pyo3(item)]
    pub redact_tool_args: Bound<'py, PyAny>,

    /// Redaction configuration for the agent state.
    ///
    /// From Python this may be provided as either:
    /// - a `bool`, or
    /// - a `list[str]` specifying which agent-state keys should be redacted.
    #[pyo3(item)]
    pub redact_agent_state: Bound<'py, PyAny>,
}

impl From<ProcessToolCallOptions<'_>> for frisk_core::policy_engine::ProcessToolCallOptions {
    fn from(v: ProcessToolCallOptions) -> Self {
        Self {
            redact_tool_args: crate::redact_dictionary::parse_redact_option(&v.redact_tool_args)
                .map(Some)
                .unwrap_or_else(|err| {
                    eprintln!(
                        "frisk_py: failed to parse 'redact_tool_args' option from Python; treating it as None. Error: {err}"
                    );
                    None
                }),
            redact_agent_state: crate::redact_dictionary::parse_redact_option(&v.redact_agent_state)
                .map(Some)
                .unwrap_or_else(|err| {
                    eprintln!(
                        "frisk_py: failed to parse 'redact_agent_state' option from Python; treating it as None. Error: {err}"
                    );
                    None
                }),
        }
    }
}

#[pymethods]
impl FriskHandle {
    #[new]
    fn new(token: &str, frisk_base_url: &str, otel_collector_endpoint: &str) -> PyResult<Self> {
        // If already initialized, return the existing instance if args match; otherwise error.
        if let Some(existing) = FRISK_HANDLE_SINGLETON.get() {
            if existing.frisk_base_url == frisk_base_url
                && existing.otel_collector_endpoint == otel_collector_endpoint
            {
                existing.update_auth_token(token)?;
                return Ok(existing.clone());
            } else {
                return Err(PyValueError::new_err(
                    "FriskHandle has already been initialized with different constructor arguments",
                ));
            }
        }

        let auth_token = Arc::new(AuthToken::new(token));

        // Build a dedicated Tokio runtime for this handle.
        let runtime = Arc::new(
            Builder::new_multi_thread()
                .thread_name("frisk-tokio")
                .enable_all()
                .build()
                .map_err(|e| {
                    PyValueError::new_err(format!("Failed to build Tokio runtime: {e}"))
                })?, // todo: Map to Python exception. https://linear.app/friskai/issue/POL-116/map-remaining-rust-errors-to-python-errors
        );

        // Initialize instrumentation and policy manager using the provided token.
        let otel_manager = Arc::new(
            OtelManager::create_from_token(auth_token.clone(), otel_collector_endpoint).map_err(
                |e| PyValueError::new_err(format!("Failed to initialize OtelManager: {e}")),
            )?, // todo: Map to Python exception. https://linear.app/friskai/issue/POL-116/map-remaining-rust-errors-to-python-errors
        );

        let policy_manager = Arc::new(Mutex::new(PolicyManager::new(
            auth_token.clone(),
            POLICY_REFRESH_INTERVAL_SECONDS,
            frisk_base_url.into(),
        )));

        // Initialize policy engine with empty policies.
        // WARNING: There is a window before the first policy fetch completes during which
        // tool calls may be processed with no policies. This may result in unintended behavior.
        // The engine will be updated via subscription once policies are fetched.
        let policy_engine = Arc::new(Mutex::new(PolicyEngine::new(
            vec![],
            Some(Arc::clone(&otel_manager)),
        )));

        // Start the policy refresh loop within this handle's runtime context.
        {
            let _guard = runtime.enter();
            let mut guard = policy_manager.lock().unwrap();
            guard.start();
        }

        // Subscribe to policy updates
        {
            let engine_handle = Arc::clone(&policy_engine);
            let pm_guard = policy_manager.lock().unwrap();
            pm_guard
                .on_policy_update
                .subscribe(move |policies: &Vec<ToolPolicyRecord>| {
                    let mut engine = engine_handle.lock().unwrap();
                    engine.set_policies(policies.clone());
                });
        }

        let handle = Self {
            auth_token,
            policy_engine,
            runtime,
            otel_manager,
            policy_manager,
            frisk_base_url: frisk_base_url.to_string(),
            otel_collector_endpoint: otel_collector_endpoint.to_string(),
        };

        // Attempt to set the singleton; if another thread raced and set with identical args,
        // get() above would have returned. If set fails here, return the existing if args match or error.
        match FRISK_HANDLE_SINGLETON.set(handle.clone()) {
            Ok(()) => Ok(handle),
            Err(_) => {
                // Fallback: another thread initialized while we were building.
                let existing = FRISK_HANDLE_SINGLETON
                    .get()
                    .expect("OnceLock set failed but no instance present");
                if existing.frisk_base_url == frisk_base_url
                    && existing.otel_collector_endpoint == otel_collector_endpoint
                {
                    existing.update_auth_token(token)?;
                    Ok(existing.clone())
                } else {
                    Err(PyValueError::new_err(
                        "FriskHandle has already been initialized with different constructor arguments",
                    )) // todo: Map to Python exception. https://linear.app/friskai/issue/POL-116/map-remaining-rust-errors-to-python-errors
                }
            }
        }
    }

    pub fn update_auth_token(&self, auth_token: &str) -> PyResult<()> {
        self.auth_token.set(auth_token.to_string());
        Ok(())
    }

    pub fn shutdown(&self) {
        // Allow Python callers to explicitly shutdown instrumentation for this handle.
        self.otel_manager.shutdown();
    }

    pub fn process<'py>(
        &self,
        _py: Python<'py>,
        tool_name: &str,
        tool_args_json: Option<&str>,
        agent_state_json: Option<&str>,
        tool_call_id: &str,
        options: Option<ProcessToolCallOptions>,
        trace_context_carrier: Option<Bound<'py, PyAny>>,
    ) -> PyResult<ProcessToolCallResultPy> {
        let tool_args_v: Value = match tool_args_json {
            Some(tool_args_str) => serde_json::from_str(tool_args_str)
                .map_err(|e| PyValueError::new_err(format!("Invalid tool_args JSON: {e}")))?,
            None => Value::Null,
        };
        let agent_state_v: Value = match agent_state_json {
            Some(agent_state_str) => serde_json::from_str(agent_state_str)
                .map_err(|e| PyValueError::new_err(format!("Invalid agent_state JSON: {e}")))?,
            None => Value::Null,
        };

        let context_dict: Option<Bound<'py, PyDict>> = trace_context_carrier
            .map(|obj| obj.cast_into::<PyDict>()) // try to cast to dict
            .transpose()?;

        let core_options: Option<frisk_core::policy_engine::ProcessToolCallOptions> =
            options.map(|opts| opts.into());

        let result = match context_dict {
            Some(trace_context_carrier_dict) => {
                let carrier = PyDictCarrier(&trace_context_carrier_dict);
                let parent_context =
                    global::get_text_map_propagator(|propagator| propagator.extract(&carrier));

                let span =
                    tracing::info_span!(target: "policy_engine", SPAN_NAME_ENGINE_DECIDE_TOOL_CALL);
                let _ = span.set_parent(parent_context);
                let _guard = span.enter();

                self.policy_engine.lock().unwrap().process_tool_call(
                    tool_name,
                    tool_call_id,
                    &tool_args_v,
                    &agent_state_v,
                    core_options,
                )
            }
            None => self.policy_engine.lock().unwrap().process_tool_call(
                tool_name,
                tool_call_id,
                &tool_args_v,
                &agent_state_v,
                core_options,
            ),
        };

        Ok(result.into())
    }
}

#[pymodule]
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<FriskHandle>()?;
    m.add_class::<ProcessToolCallResultPy>()?;
    m.add_function(wrap_pyfunction!(redact_dictionary_fn, m)?)?;
    Ok(())
}
