use crate::guardian::GuardianApprovalRequest;
use crate::guardian::GuardianNetworkAccessTrigger;
use crate::guardian::guardian_rejection_message;
use crate::guardian::guardian_timeout_message;
use crate::guardian::new_guardian_review_id;
use crate::guardian::review_approval_request;
use crate::guardian::routes_approval_to_guardian;
use crate::hook_runtime::run_permission_request_hooks;
use crate::network_policy_decision::denied_network_policy_message;
use crate::session::session::Session;
use crate::tools::sandboxing::PermissionRequestPayload;
use crate::tools::sandboxing::ToolError;
use codex_hooks::PermissionRequestDecision;
use codex_network_proxy::BlockedRequest;
use codex_network_proxy::BlockedRequestObserver;
use codex_network_proxy::NetworkDecision;
use codex_network_proxy::NetworkPolicyDecider;
use codex_network_proxy::NetworkPolicyRequest;
use codex_network_proxy::NetworkProtocol;
use codex_network_proxy::NetworkProxy;
use codex_protocol::approvals::NetworkApprovalContext;
use codex_protocol::approvals::NetworkApprovalProtocol;
use codex_protocol::approvals::NetworkPolicyRuleAction;
use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ReviewDecision;
use codex_protocol::protocol::WarningEvent;
use indexmap::IndexMap;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::sync::Notify;
use tokio::sync::OnceCell;
use tokio::sync::RwLock;
use tokio_util::sync::CancellationToken;
use tracing::warn;
use uuid::Uuid;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum NetworkApprovalMode {
    Immediate,
    Deferred,
}

#[derive(Clone, Debug)]
pub(crate) struct NetworkApprovalSpec {
    pub network: Option<NetworkProxy>,
    pub mode: NetworkApprovalMode,
    pub trigger: GuardianNetworkAccessTrigger,
    pub command: String,
}

#[derive(Clone, Debug)]
pub(crate) struct DeferredNetworkApproval {
    registration_id: String,
    cancellation_token: CancellationToken,
    finish_outcome: Arc<OnceCell<Option<NetworkApprovalOutcome>>>,
}

impl DeferredNetworkApproval {
    pub(crate) fn registration_id(&self) -> &str {
        &self.registration_id
    }

    pub(crate) fn cancellation_token(&self) -> CancellationToken {
        self.cancellation_token.clone()
    }

    pub(crate) fn is_cancelled(&self) -> bool {
        self.cancellation_token.is_cancelled()
    }

    async fn finish(&self, service: &NetworkApprovalService) -> Result<(), ToolError> {
        let outcome = self
            .finish_outcome
            .get_or_init(|| async { service.finish_call_outcome(&self.registration_id).await })
            .await
            .clone();
        network_approval_outcome_to_result(outcome)
    }
}

#[derive(Debug)]
pub(crate) struct ActiveNetworkApproval {
    registration_id: Option<String>,
    mode: NetworkApprovalMode,
    cancellation_token: CancellationToken,
}

impl ActiveNetworkApproval {
    pub(crate) fn mode(&self) -> NetworkApprovalMode {
        self.mode
    }

    pub(crate) fn cancellation_token(&self) -> CancellationToken {
        self.cancellation_token.clone()
    }

    pub(crate) fn into_deferred(self) -> Option<DeferredNetworkApproval> {
        let ActiveNetworkApproval {
            registration_id,
            mode,
            cancellation_token,
        } = self;
        match (mode, registration_id) {
            (NetworkApprovalMode::Deferred, Some(registration_id)) => {
                Some(DeferredNetworkApproval {
                    registration_id,
                    cancellation_token,
                    finish_outcome: Arc::new(OnceCell::new()),
                })
            }
            _ => None,
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct HostApprovalKey {
    host: String,
    protocol: &'static str,
    port: u16,
}

impl HostApprovalKey {
    fn from_request(request: &NetworkPolicyRequest, protocol: NetworkApprovalProtocol) -> Self {
        Self {
            host: request.host.to_ascii_lowercase(),
            protocol: protocol_key_label(protocol),
            port: request.port,
        }
    }
}

fn protocol_key_label(protocol: NetworkApprovalProtocol) -> &'static str {
    match protocol {
        NetworkApprovalProtocol::Http => "http",
        NetworkApprovalProtocol::Https => "https",
        NetworkApprovalProtocol::Socks5Tcp => "socks5-tcp",
        NetworkApprovalProtocol::Socks5Udp => "socks5-udp",
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum PendingApprovalDecision {
    AllowOnce,
    AllowForSession,
    Deny,
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum NetworkApprovalOutcome {
    DeniedByUser,
    DeniedByPolicy(String),
}

fn network_approval_outcome_to_result(
    outcome: Option<NetworkApprovalOutcome>,
) -> Result<(), ToolError> {
    match outcome {
        Some(NetworkApprovalOutcome::DeniedByUser) => {
            Err(ToolError::Rejected("rejected by user".to_string()))
        }
        Some(NetworkApprovalOutcome::DeniedByPolicy(message)) => Err(ToolError::Rejected(message)),
        None => Ok(()),
    }
}

/// Whether an allowlist miss may be reviewed instead of hard-denied.
fn allows_network_approval_flow(policy: AskForApproval) -> bool {
    !matches!(policy, AskForApproval::Never)
}

fn permission_profile_allows_network_approval_flow(permission_profile: &PermissionProfile) -> bool {
    matches!(permission_profile, PermissionProfile::Managed { .. })
}

impl PendingApprovalDecision {
    fn to_network_decision(self) -> NetworkDecision {
        match self {
            Self::AllowOnce | Self::AllowForSession => NetworkDecision::Allow,
            Self::Deny => NetworkDecision::deny("not_allowed"),
        }
    }
}

struct PendingHostApproval {
    decision: Mutex<Option<PendingApprovalDecision>>,
    notify: Notify,
}

impl PendingHostApproval {
    fn new() -> Self {
        Self {
            decision: Mutex::new(None),
            notify: Notify::new(),
        }
    }

    async fn wait_for_decision(&self) -> PendingApprovalDecision {
        loop {
            let notified = self.notify.notified();
            if let Some(decision) = *self.decision.lock().await {
                return decision;
            }
            notified.await;
        }
    }

    async fn set_decision(&self, decision: PendingApprovalDecision) {
        {
            let mut current = self.decision.lock().await;
            *current = Some(decision);
        }
        self.notify.notify_waiters();
    }
}

struct ActiveNetworkApprovalCall {
    registration_id: String,
    turn_id: String,
    trigger: GuardianNetworkAccessTrigger,
    command: String,
    cancellation_token: CancellationToken,
}

#[derive(Default)]
struct NetworkApprovalCallState {
    active_calls: IndexMap<String, Arc<ActiveNetworkApprovalCall>>,
    call_outcomes: HashMap<String, NetworkApprovalOutcome>,
}

pub(crate) struct NetworkApprovalService {
    calls: Mutex<NetworkApprovalCallState>,
    pending_host_approvals: Mutex<HashMap<HostApprovalKey, Arc<PendingHostApproval>>>,
    session_approved_hosts: Mutex<HashSet<HostApprovalKey>>,
    session_denied_hosts: Mutex<HashSet<HostApprovalKey>>,
}

impl Default for NetworkApprovalService {
    fn default() -> Self {
        Self {
            calls: Mutex::new(NetworkApprovalCallState::default()),
            pending_host_approvals: Mutex::new(HashMap::new()),
            session_approved_hosts: Mutex::new(HashSet::new()),
            session_denied_hosts: Mutex::new(HashSet::new()),
        }
    }
}

impl NetworkApprovalService {
    /// Replace the target session's approval cache with the source session's
    /// currently approved hosts.
    pub(crate) async fn sync_session_approved_hosts_to(&self, other: &Self) {
        let approved_hosts = self.session_approved_hosts.lock().await.clone();
        let mut other_approved_hosts = other.session_approved_hosts.lock().await;
        other_approved_hosts.clear();
        other_approved_hosts.extend(approved_hosts.iter().cloned());
    }

    async fn register_call(
        &self,
        registration_id: String,
        turn_id: String,
        trigger: GuardianNetworkAccessTrigger,
        command: String,
        cancellation_token: CancellationToken,
    ) {
        let mut calls = self.calls.lock().await;
        let key = registration_id.clone();
        calls.active_calls.insert(
            key,
            Arc::new(ActiveNetworkApprovalCall {
                registration_id,
                turn_id,
                trigger,
                command,
                cancellation_token,
            }),
        );
    }

    pub(crate) async fn unregister_call(&self, registration_id: &str) {
        self.remove_call(registration_id).await;
    }

    async fn resolve_single_active_call(&self) -> Option<Arc<ActiveNetworkApprovalCall>> {
        let calls = self.calls.lock().await;
        // Blocked proxy requests are not attributed to a specific tool call. Only pick an owner
        // when there is exactly one candidate; with concurrent calls, canceling one would be a guess.
        // TODO: Carry blocked-request attribution so concurrent active calls can be handled safely.
        if calls.active_calls.len() == 1 {
            return calls.active_calls.values().next().cloned();
        }

        None
    }

    async fn get_or_create_pending_approval(
        &self,
        key: HostApprovalKey,
    ) -> (Arc<PendingHostApproval>, bool) {
        let mut pending = self.pending_host_approvals.lock().await;
        if let Some(existing) = pending.get(&key).cloned() {
            return (existing, false);
        }

        let created = Arc::new(PendingHostApproval::new());
        pending.insert(key, Arc::clone(&created));
        (created, true)
    }

    async fn record_outcome_for_single_active_call(&self, outcome: NetworkApprovalOutcome) {
        let Some(owner_call) = self.resolve_single_active_call().await else {
            return;
        };
        self.record_call_outcome(&owner_call.registration_id, outcome)
            .await;
    }

    #[cfg(test)]
    async fn take_call_outcome(&self, registration_id: &str) -> Option<NetworkApprovalOutcome> {
        let mut calls = self.calls.lock().await;
        calls.call_outcomes.remove(registration_id)
    }

    async fn record_call_outcome(&self, registration_id: &str, outcome: NetworkApprovalOutcome) {
        let mut calls = self.calls.lock().await;
        let Some(call) = calls.active_calls.get(registration_id).cloned() else {
            return;
        };
        if matches!(
            calls.call_outcomes.get(registration_id),
            Some(NetworkApprovalOutcome::DeniedByUser)
        ) {
            return;
        }
        calls
            .call_outcomes
            .insert(registration_id.to_string(), outcome);

        drop(calls);
        call.cancellation_token.cancel();
    }

    async fn remove_call(&self, registration_id: &str) -> Option<NetworkApprovalOutcome> {
        let mut calls = self.calls.lock().await;
        calls.active_calls.shift_remove(registration_id);
        calls.call_outcomes.remove(registration_id)
    }

    async fn finish_call_outcome(&self, registration_id: &str) -> Option<NetworkApprovalOutcome> {
        self.remove_call(registration_id).await
    }

    async fn finish_call(&self, registration_id: &str) -> Result<(), ToolError> {
        network_approval_outcome_to_result(self.finish_call_outcome(registration_id).await)
    }

    pub(crate) async fn record_blocked_request(&self, blocked: BlockedRequest) {
        let Some(message) = denied_network_policy_message(&blocked) else {
            return;
        };

        self.record_outcome_for_single_active_call(NetworkApprovalOutcome::DeniedByPolicy(message))
            .await;
    }

    async fn active_turn_context(
        session: &Session,
    ) -> Option<Arc<crate::session::turn_context::TurnContext>> {
        let active_turn = session.active_turn.lock().await;
        active_turn
            .as_ref()
            .and_then(|turn| turn.tasks.first())
            .map(|(_, task)| Arc::clone(&task.turn_context))
    }

    fn format_network_target(protocol: &str, host: &str, port: u16) -> String {
        format!("{protocol}://{host}:{port}")
    }

    fn approval_id_for_key(key: &HostApprovalKey) -> String {
        format!("network#{}#{}#{}", key.protocol, key.host, key.port)
    }

    pub(crate) async fn handle_inline_policy_request(
        &self,
        session: Arc<Session>,
        request: NetworkPolicyRequest,
    ) -> NetworkDecision {
        const REASON_NOT_ALLOWED: &str = "not_allowed";

        let protocol = match request.protocol {
            NetworkProtocol::Http => NetworkApprovalProtocol::Http,
            NetworkProtocol::HttpsConnect => NetworkApprovalProtocol::Https,
            NetworkProtocol::Socks5Tcp => NetworkApprovalProtocol::Socks5Tcp,
            NetworkProtocol::Socks5Udp => NetworkApprovalProtocol::Socks5Udp,
        };
        let key = HostApprovalKey::from_request(&request, protocol);

        {
            let denied_hosts = self.session_denied_hosts.lock().await;
            if denied_hosts.contains(&key) {
                return NetworkDecision::deny(REASON_NOT_ALLOWED);
            }
        }

        {
            let approved_hosts = self.session_approved_hosts.lock().await;
            if approved_hosts.contains(&key) {
                return NetworkDecision::Allow;
            }
        }

        let (pending, is_owner) = self.get_or_create_pending_approval(key.clone()).await;
        if !is_owner {
            return pending.wait_for_decision().await.to_network_decision();
        }

        let target = Self::format_network_target(key.protocol, request.host.as_str(), key.port);
        let policy_denial_message =
            format!("Network access to \"{target}\" was blocked by policy.");
        let prompt_reason = format!("{} is not in the allowed_domains", request.host);

        let Some(turn_context) = Self::active_turn_context(session.as_ref()).await else {
            pending.set_decision(PendingApprovalDecision::Deny).await;
            self.pending_host_approvals.lock().await.remove(&key);
            self.record_outcome_for_single_active_call(NetworkApprovalOutcome::DeniedByPolicy(
                policy_denial_message,
            ))
            .await;
            return NetworkDecision::deny(REASON_NOT_ALLOWED);
        };
        if !permission_profile_allows_network_approval_flow(&turn_context.permission_profile()) {
            pending.set_decision(PendingApprovalDecision::Deny).await;
            self.pending_host_approvals.lock().await.remove(&key);
            self.record_outcome_for_single_active_call(NetworkApprovalOutcome::DeniedByPolicy(
                policy_denial_message,
            ))
            .await;
            return NetworkDecision::deny(REASON_NOT_ALLOWED);
        }
        if !allows_network_approval_flow(turn_context.approval_policy.value()) {
            pending.set_decision(PendingApprovalDecision::Deny).await;
            self.pending_host_approvals.lock().await.remove(&key);
            self.record_outcome_for_single_active_call(NetworkApprovalOutcome::DeniedByPolicy(
                policy_denial_message,
            ))
            .await;
            return NetworkDecision::deny(REASON_NOT_ALLOWED);
        }

        let owner_call = self.resolve_single_active_call().await;
        let network_approval_context = NetworkApprovalContext {
            host: request.host.clone(),
            protocol,
        };
        let guardian_approval_id = Self::approval_id_for_key(&key);
        let prompt_command = vec!["network-access".to_string(), target.clone()];
        let command = owner_call
            .as_ref()
            .map_or_else(|| prompt_command.join(" "), |call| call.command.clone());
        if let Some(permission_request_decision) = run_permission_request_hooks(
            &session,
            &turn_context,
            &guardian_approval_id,
            PermissionRequestPayload::bash(command, Some(format!("network-access {target}"))),
        )
        .await
        {
            match permission_request_decision {
                PermissionRequestDecision::Allow => {
                    pending
                        .set_decision(PendingApprovalDecision::AllowOnce)
                        .await;
                    let mut pending_approvals = self.pending_host_approvals.lock().await;
                    pending_approvals.remove(&key);
                    return NetworkDecision::Allow;
                }
                PermissionRequestDecision::Deny { message } => {
                    if let Some(owner_call) = owner_call.as_ref() {
                        self.record_call_outcome(
                            &owner_call.registration_id,
                            NetworkApprovalOutcome::DeniedByPolicy(message),
                        )
                        .await;
                    }
                    pending.set_decision(PendingApprovalDecision::Deny).await;
                    let mut pending_approvals = self.pending_host_approvals.lock().await;
                    pending_approvals.remove(&key);
                    return NetworkDecision::deny(REASON_NOT_ALLOWED);
                }
            }
        }
        let use_guardian = routes_approval_to_guardian(&turn_context);
        let guardian_review_id = use_guardian.then(new_guardian_review_id);
        let approval_decision = if let Some(review_id) = guardian_review_id.clone() {
            review_approval_request(
                &session,
                &turn_context,
                review_id,
                GuardianApprovalRequest::NetworkAccess {
                    id: guardian_approval_id.clone(),
                    turn_id: owner_call
                        .as_ref()
                        .map_or_else(|| turn_context.sub_id.clone(), |call| call.turn_id.clone()),
                    target,
                    host: request.host,
                    protocol,
                    port: key.port,
                    trigger: owner_call.as_ref().map(|call| call.trigger.clone()),
                },
                Some(policy_denial_message.clone()),
            )
            .await
        } else {
            let available_decisions = None;
            session
                .request_command_approval(
                    turn_context.as_ref(),
                    guardian_approval_id,
                    /*approval_id*/ None,
                    prompt_command,
                    turn_context.cwd.clone(),
                    Some(prompt_reason),
                    Some(network_approval_context.clone()),
                    /*proposed_execpolicy_amendment*/ None,
                    /*additional_permissions*/ None,
                    available_decisions,
                )
                .await
        };

        let mut cache_session_deny = false;
        let resolved = match approval_decision {
            ReviewDecision::Approved | ReviewDecision::ApprovedExecpolicyAmendment { .. } => {
                PendingApprovalDecision::AllowOnce
            }
            ReviewDecision::ApprovedForSession => PendingApprovalDecision::AllowForSession,
            ReviewDecision::NetworkPolicyAmendment {
                network_policy_amendment,
            } => match network_policy_amendment.action {
                NetworkPolicyRuleAction::Allow => {
                    match session
                        .persist_network_policy_amendment(
                            &network_policy_amendment,
                            &network_approval_context,
                        )
                        .await
                    {
                        Ok(()) => {
                            session
                                .record_network_policy_amendment_message(
                                    &turn_context.sub_id,
                                    &network_policy_amendment,
                                )
                                .await;
                        }
                        Err(err) => {
                            let message =
                                format!("Failed to apply network policy amendment: {err}");
                            warn!("{message}");
                            session
                                .send_event_raw(Event {
                                    id: turn_context.sub_id.clone(),
                                    msg: EventMsg::Warning(WarningEvent { message }),
                                })
                                .await;
                        }
                    }
                    PendingApprovalDecision::AllowForSession
                }
                NetworkPolicyRuleAction::Deny => {
                    match session
                        .persist_network_policy_amendment(
                            &network_policy_amendment,
                            &network_approval_context,
                        )
                        .await
                    {
                        Ok(()) => {
                            session
                                .record_network_policy_amendment_message(
                                    &turn_context.sub_id,
                                    &network_policy_amendment,
                                )
                                .await;
                        }
                        Err(err) => {
                            let message =
                                format!("Failed to apply network policy amendment: {err}");
                            warn!("{message}");
                            session
                                .send_event_raw(Event {
                                    id: turn_context.sub_id.clone(),
                                    msg: EventMsg::Warning(WarningEvent { message }),
                                })
                                .await;
                        }
                    }
                    if let Some(owner_call) = owner_call.as_ref() {
                        self.record_call_outcome(
                            &owner_call.registration_id,
                            NetworkApprovalOutcome::DeniedByUser,
                        )
                        .await;
                    }
                    cache_session_deny = true;
                    PendingApprovalDecision::Deny
                }
            },
            ReviewDecision::Denied | ReviewDecision::Abort => {
                if let Some(review_id) = guardian_review_id.as_deref() {
                    if let Some(owner_call) = owner_call.as_ref() {
                        let message = guardian_rejection_message(session.as_ref(), review_id).await;
                        self.record_call_outcome(
                            &owner_call.registration_id,
                            NetworkApprovalOutcome::DeniedByPolicy(message),
                        )
                        .await;
                    }
                } else if let Some(owner_call) = owner_call.as_ref() {
                    self.record_call_outcome(
                        &owner_call.registration_id,
                        NetworkApprovalOutcome::DeniedByUser,
                    )
                    .await;
                }
                PendingApprovalDecision::Deny
            }
            ReviewDecision::TimedOut => {
                if let Some(owner_call) = owner_call.as_ref() {
                    self.record_call_outcome(
                        &owner_call.registration_id,
                        NetworkApprovalOutcome::DeniedByPolicy(guardian_timeout_message()),
                    )
                    .await;
                }
                PendingApprovalDecision::Deny
            }
        };

        if matches!(resolved, PendingApprovalDecision::AllowForSession) {
            {
                let mut denied_hosts = self.session_denied_hosts.lock().await;
                denied_hosts.remove(&key);
            }
            let mut approved_hosts = self.session_approved_hosts.lock().await;
            approved_hosts.insert(key.clone());
        }

        if cache_session_deny {
            {
                let mut approved_hosts = self.session_approved_hosts.lock().await;
                approved_hosts.remove(&key);
            }
            let mut denied_hosts = self.session_denied_hosts.lock().await;
            denied_hosts.insert(key.clone());
        }

        pending.set_decision(resolved).await;
        let mut pending_approvals = self.pending_host_approvals.lock().await;
        pending_approvals.remove(&key);

        resolved.to_network_decision()
    }
}

pub(crate) fn build_blocked_request_observer(
    network_approval: Arc<NetworkApprovalService>,
) -> Arc<dyn BlockedRequestObserver> {
    Arc::new(move |blocked: BlockedRequest| {
        let network_approval = Arc::clone(&network_approval);
        async move {
            network_approval.record_blocked_request(blocked).await;
        }
    })
}

pub(crate) fn build_network_policy_decider(
    network_approval: Arc<NetworkApprovalService>,
    network_policy_decider_session: Arc<RwLock<std::sync::Weak<Session>>>,
) -> Arc<dyn NetworkPolicyDecider> {
    Arc::new(move |request: NetworkPolicyRequest| {
        let network_approval = Arc::clone(&network_approval);
        let network_policy_decider_session = Arc::clone(&network_policy_decider_session);
        async move {
            let Some(session) = network_policy_decider_session.read().await.upgrade() else {
                return NetworkDecision::ask("not_allowed");
            };
            network_approval
                .handle_inline_policy_request(session, request)
                .await
        }
    })
}

pub(crate) async fn begin_network_approval(
    session: &Session,
    turn_id: &str,
    managed_network_active: bool,
    spec: Option<NetworkApprovalSpec>,
) -> Option<ActiveNetworkApproval> {
    let NetworkApprovalSpec {
        network,
        mode,
        trigger,
        command,
    } = spec?;
    if !managed_network_active || network.is_none() {
        return None;
    }

    let registration_id = Uuid::new_v4().to_string();
    let cancellation_token = CancellationToken::new();
    session
        .services
        .network_approval
        .register_call(
            registration_id.clone(),
            turn_id.to_string(),
            trigger,
            command,
            cancellation_token.clone(),
        )
        .await;

    Some(ActiveNetworkApproval {
        registration_id: Some(registration_id),
        mode,
        cancellation_token,
    })
}

pub(crate) async fn finish_immediate_network_approval(
    session: &Session,
    active: ActiveNetworkApproval,
) -> Result<(), ToolError> {
    let Some(registration_id) = active.registration_id.as_deref() else {
        return Ok(());
    };

    session
        .services
        .network_approval
        .finish_call(registration_id)
        .await
}

pub(crate) async fn finish_deferred_network_approval(
    session: &Session,
    deferred: Option<DeferredNetworkApproval>,
) -> Result<(), ToolError> {
    let Some(deferred) = deferred else {
        return Ok(());
    };
    deferred.finish(&session.services.network_approval).await
}

#[cfg(test)]
#[path = "network_approval_tests.rs"]
mod tests;
