use std::{collections::HashMap, net::SocketAddr};

use playit_agent_proto::PortProto;
use playit_api_client::api::{AgentRunData, PortType, ProxyProtocol};
use tokio::sync::RwLock;

#[derive(Default)]
pub struct OriginLookup {
    map: RwLock<HashMap<Key, OriginResource>>,
}

impl OriginLookup {
    pub async fn update_from_run_data(&self, run_data: &AgentRunData) {
        self.update(run_data.tunnels.iter().map(|tunn| {
            OriginResource {
                tunnel_id: tunn.internal_id,
                proto: match tunn.proto {
                    PortType::Tcp => PortProto::Tcp,
                    PortType::Udp => PortProto::Udp,
                    PortType::Both => PortProto::Both,
                },
                local_addr: SocketAddr::new(tunn.local_ip, tunn.local_port),
                port_count: tunn.port.to - tunn.port.from,
                proxy_protocol: tunn.proxy_protocol,
            }
        })).await;
    }

    pub async fn update<I: Iterator<Item = OriginResource>>(&self, resources: I) {
        let mut lock = self.map.write().await;
        lock.clear();

        for res in resources {
            match res.proto {
                PortProto::Tcp => {
                    lock.insert(Key { tunnel_id: res.tunnel_id, is_tcp: true }, res);
                }
                PortProto::Udp => {
                    lock.insert(Key { tunnel_id: res.tunnel_id, is_tcp: false }, res);
                }
                PortProto::Both => {
                    lock.insert(Key { tunnel_id: res.tunnel_id, is_tcp: true }, res.clone());
                    lock.insert(Key { tunnel_id: res.tunnel_id, is_tcp: false }, res);
                }
            }
        }
    }

    pub async fn lookup(&self, tunnel_id: u64, is_tcp: bool) -> Option<OriginResource> {
        self.map.read().await.get(&Key { tunnel_id, is_tcp }).cloned()
    }
}

#[derive(Debug, PartialEq, Eq, Hash)]
struct Key {
    tunnel_id: u64,
    is_tcp: bool,
}

#[derive(Debug, Clone)]
pub struct OriginResource {
    pub tunnel_id: u64,
    pub proto: PortProto,
    pub local_addr: SocketAddr,
    pub port_count: u16,
    pub proxy_protocol: Option<ProxyProtocol>,
}

impl OriginResource {
    pub fn resolve_local(&self, port_offset: u16) -> Option<SocketAddr> {
        if port_offset == 0 {
            Some(self.local_addr)
        } else if port_offset < self.port_count {
            Some(SocketAddr::new(self.local_addr.ip(), self.local_addr.port() + port_offset))
        } else {
            None
        }
    }
}
