use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use serde_json::Value;

use crate::pipe::protocol::{
    Action, AgentMessage, BrowserMessage, ExecutionSurfaceMetadata, SecurityFields, Timing,
};
use crate::pipe::{PipeError, Transport};
use crate::security::{sign_command, MacPolicy};

const LOCAL_DASHBOARD_EXPECTED_DOMAIN: &str = "__sgclaw_local_dashboard__";

#[derive(Debug, Clone, PartialEq)]
pub struct CommandOutput {
    pub seq: u64,
    pub success: bool,
    pub data: Value,
    pub aom_snapshot: Vec<Value>,
    pub timing: Timing,
}

pub struct BrowserPipeTool<T: Transport> {
    transport: Arc<T>,
    mac_policy: MacPolicy,
    session_key: Vec<u8>,
    next_seq: Arc<AtomicU64>,
    response_timeout: Duration,
}

impl<T: Transport> Clone for BrowserPipeTool<T> {
    fn clone(&self) -> Self {
        Self {
            transport: self.transport.clone(),
            mac_policy: self.mac_policy.clone(),
            session_key: self.session_key.clone(),
            next_seq: self.next_seq.clone(),
            response_timeout: self.response_timeout,
        }
    }
}

impl<T: Transport> BrowserPipeTool<T> {
    pub fn new(transport: Arc<T>, mac_policy: MacPolicy, session_key: Vec<u8>) -> Self {
        Self {
            transport,
            mac_policy,
            session_key,
            next_seq: Arc::new(AtomicU64::new(1)),
            response_timeout: Duration::from_secs(30),
        }
    }

    pub fn with_response_timeout(mut self, response_timeout: Duration) -> Self {
        self.response_timeout = response_timeout;
        self
    }

    pub fn response_timeout(&self) -> Duration {
        self.response_timeout
    }

    pub fn mac_policy(&self) -> &MacPolicy {
        &self.mac_policy
    }

    pub fn surface_metadata(&self) -> ExecutionSurfaceMetadata {
        self.mac_policy.privileged_surface_metadata()
    }

    pub fn supports_eval(&self) -> bool {
        self.mac_policy.supports_pipe_action(&Action::Eval)
    }

    pub fn invoke(
        &self,
        action: Action,
        params: Value,
        expected_domain: &str,
    ) -> Result<CommandOutput, PipeError> {
        if let Some((presentation_url, output_path)) = approved_local_dashboard_request(
            &action,
            &params,
            expected_domain,
        ) {
            self.mac_policy
                .validate_local_dashboard_presentation(
                    &action,
                    expected_domain,
                    &presentation_url,
                    &output_path,
                )
                .map_err(PipeError::Security)?;
        } else {
            self.mac_policy
                .validate(&action, expected_domain)
                .map_err(PipeError::Security)?;
        }

        let seq = self.next_seq.fetch_add(1, Ordering::Relaxed);
        let hmac = sign_command(&self.session_key, seq, &action, &params, expected_domain)?;
        let command = AgentMessage::Command {
            seq,
            action,
            params,
            security: SecurityFields {
                expected_domain: expected_domain.to_string(),
                hmac,
            },
        };

        self.transport.send(&command)?;

        let started = Instant::now();
        loop {
            let Some(remaining) = self.response_timeout.checked_sub(started.elapsed()) else {
                return Err(PipeError::Timeout);
            };

            match self.transport.recv_timeout(remaining)? {
                BrowserMessage::Response {
                    seq: response_seq,
                    success,
                    data,
                    aom_snapshot,
                    timing,
                } if response_seq == seq => {
                    return Ok(CommandOutput {
                        seq: response_seq,
                        success,
                        data,
                        aom_snapshot,
                        timing,
                    });
                }
                BrowserMessage::Response {
                    seq: response_seq, ..
                } => {
                    return Err(PipeError::Protocol(format!(
                        "received response seq {response_seq} while waiting for {seq}"
                    )));
                }
                BrowserMessage::Init { .. } => {
                    return Err(PipeError::UnexpectedMessage(
                        "received duplicate init after handshake".to_string(),
                    ));
                }
                BrowserMessage::Connect => {
                    return Err(PipeError::UnexpectedMessage(
                        "received connect while waiting for response".to_string(),
                    ));
                }
                BrowserMessage::Start => {
                    return Err(PipeError::UnexpectedMessage(
                        "received start while waiting for response".to_string(),
                    ));
                }
                BrowserMessage::Stop => {
                    return Err(PipeError::UnexpectedMessage(
                        "received stop while waiting for response".to_string(),
                    ));
                }
                BrowserMessage::SubmitTask { .. } => {
                    return Err(PipeError::UnexpectedMessage(
                        "received submit_task while waiting for response".to_string(),
                    ));
                }
            }
        }
    }
}

fn approved_local_dashboard_request(
    action: &Action,
    params: &Value,
    expected_domain: &str,
) -> Option<(String, String)> {
    if action != &Action::Navigate || expected_domain != LOCAL_DASHBOARD_EXPECTED_DOMAIN {
        return None;
    }

    let presentation_url = params.get("url")?.as_str()?.trim();
    let marker = params.get("sgclaw_local_dashboard_open")?.as_object()?;
    let source = marker.get("source")?.as_str()?.trim();
    let kind = marker.get("kind")?.as_str()?.trim();
    let output_path = marker.get("output_path")?.as_str()?.trim();
    let marker_presentation_url = marker.get("presentation_url")?.as_str()?.trim();

    if source != "compat.workflow_executor"
        || kind != "zhihu_hotlist_screen"
        || output_path.is_empty()
        || presentation_url.is_empty()
        || marker_presentation_url != presentation_url
    {
        return None;
    }

    Some((presentation_url.to_string(), output_path.to_string()))
}
