use crc32fast::{hash, Hasher};
use log::warn;
use nom::bytes::complete::{tag, take};
use nom::error::ParseError;
use nom::multi::many;
use nom::number::complete::{be_u128, be_u16, be_u32, be_u8};
use nom::{AsBytes, IResult, Parser};
use rand::{self, RngCore};
use serde::Serialize;
use std::fmt::{self, Debug};
use std::net::{IpAddr, SocketAddr};

// https://github.com/tailscale/tailscale/blob/main/net/stun/stun.go

const ATTR_NUM_SOFTWARE: u16 = 0x8022;
const ATTR_NUM_FINGERPRINT: u16 = 0x8028;
const ATTR_MAPPED_ADDRESS: u16 = 0x0001;
const ATTR_XOR_MAPPED_ADDRESS: u16 = 0x0020;
// This alternative attribute type is not
// mentioned in the RFC, but the shift into
// the "comprehension-optional" range seems
// like an easy mistake for a server to make.
// And servers appear to send it.
const ATTR_XOR_MAPPED_ADDRESS_ALT: u16 = 0x8020;
const ATTR_SOURCE_ADDRESS: u16 = 0x0004;
const ATTR_CHANGED_ADDRESS: u16 = 0x0005;
const ATTR_USERNAME: u16 = 0x0006;
const ATTR_MESSAGE_INTEGRITY: u16 = 0x0008;
const ATTR_ERROR_CODE: u16 = 0x0009;
const ATTR_UNKNOWN_ATTRIBUTES: u16 = 0x000a;
const ATTR_REALM: u16 = 0x0014;
const ATTR_NONCE: u16 = 0x0015;
const ATTR_ALTERNATE_SERVER: u16 = 0x8023;
const ATTR_RESPONSE_ORIGIN: u16 = 0x802b;
const ATTR_OTHER_ADDRESS: u16 = 0x802c;

const SOFTWARE: [u8; 8] = *b"tailnode";
const BINDING_REQUEST: [u8; 2] = [0x00, 0x01];
const MAGIC_COOKIE: [u8; 4] = [0x21, 0x12, 0xa4, 0x42];
const LEN_FINGERPRINT: u16 = 8;
const HEADER_LEN: u16 = 20;

#[derive(Debug, Clone)]
pub struct TxId([u8; 12]);

impl Default for TxId {
    fn default() -> Self {
        Self::new([0; 12])
    }
}

impl TxId {
    pub fn new(tx_id: [u8; 12]) -> Self {
        Self(tx_id)
    }

    pub fn random() -> Self {
        let mut tx_id = [0; 12];
        rand::rng().fill_bytes(&mut tx_id);
        Self::new(tx_id)
    }

    pub fn from_bytes(bytes: &[u8]) -> Self {
        let mut tx_id = [0; 12];
        tx_id.copy_from_slice(bytes);
        Self(tx_id)
    }

    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }
}

impl Serialize for TxId {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_u128(u128::from(self))
    }
}

impl From<&TxId> for u128 {
    fn from(value: &TxId) -> Self {
        let mut padded = [0u8; 16];
        padded[4..].copy_from_slice(&value.0);
        u128::from_be_bytes(padded)
    }
}

fn fingerprint(msg: &[u8]) -> u32 {
    hash(msg) ^ 0x5354554e
}

#[derive(Debug, Clone, Serialize)]
pub struct AddrPort {
    pub address: IpAddr,
    pub port: u16,
}

impl From<(IpAddr, u16)> for AddrPort {
    fn from(value: (IpAddr, u16)) -> Self {
        Self {
            address: value.0,
            port: value.1,
        }
    }
}

#[derive(Debug, Clone, Copy, Serialize)]
pub enum StunClass {
    Request = 0,
    Indication = 1,
    SuccessResponse = 2,
    ErrorResponse = 3,
}

#[derive(Debug, Clone, Copy, Serialize)]
pub enum StunMethod {
    Binding = 1,
}

#[derive(Debug, Clone, Serialize)]
pub enum StunAttribute {
    MappedAddress(SocketAddr),
    XorMappedAddress(SocketAddr),
    SourceAddress(SocketAddr),
    ChangedAddress(SocketAddr),
    Username(String),
    MessageIntegrity([u8; 20]),
    Fingerprint((u32, bool)),
    ErrorCode((u16, String)),
    Realm(String),
    Nonce(String),
    UnknownAttributes(Vec<u16>),
    Software(String),
    AlternateServer(SocketAddr),
    ResponseOrigin(SocketAddr),
    OtherAddress(SocketAddr),
    Unknown((u16, Vec<u8>)),
}

fn addr_family(addr: &IpAddr) -> &'static str {
    match addr {
        IpAddr::V4(_) => "IPv4",
        IpAddr::V6(_) => "IPv6",
    }
}

impl fmt::Display for StunAttribute {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Helper function for attributes with IP and port
        fn format_ip_port(
            f: &mut fmt::Formatter<'_>,
            label: &str,
            addr: &SocketAddr,
        ) -> fmt::Result {
            write!(f, "  {} ({}) {}", label, addr_family(&addr.ip()), addr)
        }

        match self {
            StunAttribute::MappedAddress(a) => format_ip_port(f, "MappedAddress", a),
            StunAttribute::SourceAddress(a) => format_ip_port(f, "SourceAddress", a),
            StunAttribute::ChangedAddress(a) => format_ip_port(f, "ChangedAddress", a),
            StunAttribute::XorMappedAddress(a) => format_ip_port(f, "XorMappedAddress", a),
            StunAttribute::AlternateServer(a) => format_ip_port(f, "AlternateServer", a),
            StunAttribute::ResponseOrigin(a) => format_ip_port(f, "ResponseOrigin", a),
            StunAttribute::OtherAddress(a) => format_ip_port(f, "OtherAddress", a),
            StunAttribute::Username(username) => write!(f, "  Username {}", username),
            StunAttribute::MessageIntegrity(msg_integrity) => {
                write!(f, "  MessageIntegrity {:?}", msg_integrity)
            }
            StunAttribute::Fingerprint((crc, ok)) => {
                write!(
                    f,
                    "  Fingerprint 0x{:08x} ({})",
                    crc,
                    if *ok { "OK" } else { "FAIL" }
                )
            }
            StunAttribute::ErrorCode((err_num, error)) => {
                write!(f, "  ErrorCode {} ({})", err_num, error)
            }
            StunAttribute::Realm(realm) => write!(f, "  Realm {}", realm),
            StunAttribute::Nonce(nonce) => write!(f, "  Nonce {}", nonce),
            StunAttribute::UnknownAttributes(unknown_attrs) => {
                write!(f, "  UnknownAttributes {:?}", unknown_attrs)
            }
            StunAttribute::Software(software) => write!(f, "  Software {}", software),
            StunAttribute::Unknown((attr_type, data)) => {
                write!(f, "  Unknown ({:04x}) {:?}", attr_type, data)
            }
        }
    }
}

#[derive(Debug, Serialize)]
pub struct StunMessageType {
    pub class: StunClass,
    pub method: StunMethod,
}

#[derive(Debug, Serialize)]
pub struct StunHeader {
    pub msg_type: StunMessageType,
    pub msg_length: u16,
    pub tx_id: TxId,
}

impl fmt::Display for StunHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "  Header")?;
        writeln!(
            f,
            "    MessageType class={:?} (0x{:02x}) method={:?} (0x{:04x})",
            self.msg_type.class,
            self.msg_type.class as usize,
            self.msg_type.method,
            self.msg_type.method as usize
        )
    }
}

#[derive(Debug, Serialize)]
pub struct StunAttributes(Vec<StunAttribute>);

impl fmt::Display for StunAttributes {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "  Attributes")?;
        for attr in &self.0 {
            writeln!(f, "    {}", attr)?;
        }
        Ok(())
    }
}

impl StunAttributes {
    pub fn mapped_address(&self) -> Option<&SocketAddr> {
        self.0.iter().find_map(|attr| match attr {
            StunAttribute::MappedAddress(addr) | StunAttribute::XorMappedAddress(addr) => {
                Some(addr)
            }
            _ => None,
        })
    }
}

#[derive(Debug, Serialize)]
pub struct StunMessage {
    pub header: StunHeader,
    pub attributes: StunAttributes,
}

impl fmt::Display for StunMessage {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "StunMessage")?;
        write!(f, "{}", self.header)?;
        write!(f, "{}", self.attributes)
    }
}

impl StunMessage {
    pub fn parse(bytes: &[u8]) -> Result<Self, nom::Err<nom::error::Error<&[u8]>>> {
        let (_, msg) = parse_stun_message(bytes)?;
        Ok(msg)
    }
}

pub fn rand_request() -> Vec<u8> {
    const LEN_ATTR_SOFTWARE: u16 = 4 + SOFTWARE.len() as u16;
    let tx_id = TxId::random();
    let mut buf = Vec::with_capacity((HEADER_LEN + LEN_ATTR_SOFTWARE + LEN_FINGERPRINT) as usize);
    buf.extend(&BINDING_REQUEST);
    buf.extend((LEN_ATTR_SOFTWARE + LEN_FINGERPRINT).to_be_bytes());
    buf.extend(&MAGIC_COOKIE);
    buf.extend(tx_id.as_bytes());
    buf.extend(ATTR_NUM_SOFTWARE.to_be_bytes());
    buf.extend((SOFTWARE.len() as u16).to_be_bytes());
    buf.extend(&SOFTWARE);

    let fp = fingerprint(&buf);
    buf.extend(ATTR_NUM_FINGERPRINT.to_be_bytes());
    buf.extend(4_u16.to_be_bytes());
    buf.extend(&fp.to_be_bytes());

    buf
}

fn take_txid<I, E: ParseError<I>>(input: I) -> IResult<I, TxId, E>
where
    I: nom::Input<Item = u8> + AsBytes,
{
    let (input, tx_id) = take(12usize)(input)?;
    Ok((input, TxId::from_bytes(tx_id.as_bytes())))
}

fn parse_stun_message<'a, I, E: ParseError<I>>(input: I) -> IResult<I, StunMessage, E>
where
    I: nom::Input<Item = u8>
        + nom::Compare<I>
        + nom::Compare<&'a [u8]>
        + nom::Offset
        + AsBytes
        + Debug,
{
    let mut hasher = Some(Hasher::new());
    let input_start = input.clone();

    let (input, h) = parse_stun_header(input)?;
    hasher
        .as_mut()
        .unwrap()
        .update(input_start.take(input_start.offset(&input)).as_bytes());

    let (residual, input) = take(h.msg_length)(input)?;
    if residual.input_len() != 0 {
        warn!("Trailing bytes in STUN message: {:?}", residual);
    }

    let mut input = input;
    let mut attributes = Vec::new();

    while let Ok((new_input, attr)) = parse_stun_attribute::<I, E>(&h.tx_id)(input.clone()) {
        let attr = if let Some(StunAttribute::Fingerprint(fingerprint)) = attr {
            let crc = hasher.unwrap().finalize() ^ 0x5354554e;
            hasher = None;
            if crc == fingerprint.0 {
                Some(StunAttribute::Fingerprint((crc, true)))
            } else {
                warn!(
                    "Fingerprint mismatch: expected 0x{:08x}, got 0x{:08x}",
                    crc, fingerprint.0
                );
                attr
            }
        } else {
            if let Some(hasher) = hasher.as_mut() {
                hasher.update(input.take(input.offset(&new_input)).as_bytes());
            } else {
                warn!("Received attributes after FINGERPRINT");
            }
            attr
        };

        attributes.push(attr);
        input = new_input;
    }
    if input.input_len() != 0 {
        warn!("Trailing bytes in STUN message attributes: {:?}", input);
    }
    let attributes = StunAttributes(attributes.iter().filter_map(|i| i.clone()).collect());

    Ok((
        residual,
        StunMessage {
            header: h,
            attributes,
        },
    ))
}

fn parse_stun_message_type<I: nom::Input<Item = u8>, E: ParseError<I>>(
    input: I,
) -> IResult<I, StunMessageType, E> {
    let (input, msg_type_raw) = be_u16(input)?;
    if msg_type_raw & 0b11000000 != 0 {
        panic!("Invalid STUN message type");
    }
    let class_raw = ((msg_type_raw & (1 << 4)) >> 4) | ((msg_type_raw & (1 << 8)) >> 7);
    let class = match class_raw {
        0 => StunClass::Request,
        1 => StunClass::Indication,
        2 => StunClass::SuccessResponse,
        3 => StunClass::ErrorResponse,
        _ => panic!("Invalid STUN message class"),
    };
    let method_raw = msg_type_raw & 0x0f | msg_type_raw & 0xe0 >> 1 | msg_type_raw & 0x3e >> 2;
    let method = match method_raw {
        1 => StunMethod::Binding,
        _ => panic!("Invalid STUN message method"),
    };
    Ok((input, StunMessageType { class, method }))
}

fn parse_stun_address<I: nom::Input<Item = u8>, E: ParseError<I>>(
    input: I,
) -> IResult<I, SocketAddr, E> {
    let (input, _) = take(1usize)(input)?;
    let (input, family) = be_u8(input)?;
    let (input, port) = be_u16(input)?;
    let (input, addr) = match family {
        0x01 => {
            let (input, val) = be_u32(input)?;
            (input, (IpAddr::V4(val.into()), port).into())
        }
        0x02 => {
            let (input, val) = be_u128(input)?;
            (input, (IpAddr::V6(val.into()), port).into())
        }
        _ => panic!("Invalid address family"),
    };
    Ok((input, addr))
}

fn parse_stun_xor_address<I, E: ParseError<I>>(input: I, tx_id: &TxId) -> IResult<I, SocketAddr, E>
where
    I: nom::Input<Item = u8>,
{
    let (input, addr) = parse_stun_address(input)?;
    let xor_port = addr.port() ^ 0x2112;
    let xor_addr = match addr {
        SocketAddr::V4(v4) => {
            let v4 = v4.ip().to_bits();
            let xor_v4 = v4 ^ 0x2112a442;
            IpAddr::V4(xor_v4.into())
        }
        SocketAddr::V6(v6) => {
            let v6 = v6.ip().to_bits();
            let xor_v6: u128 = v6 ^ (0x2112a442 << 96 | u128::from(tx_id));
            IpAddr::V6(xor_v6.into())
        }
    };
    Ok((input, (xor_addr, xor_port).into()))
}

fn parse_stun_attribute<I, E: ParseError<I>>(
    tx_id: &TxId,
) -> impl Fn(I) -> IResult<I, Option<StunAttribute>, E>
where
    I: nom::Input<Item = u8> + nom::Compare<I> + AsBytes,
{
    let tx_id = tx_id.clone();
    move |input| parse_stun_attribute_impl(input, &tx_id)
}

fn parse_stun_attribute_impl<I, E: ParseError<I>>(
    input: I,
    tx_id: &TxId,
) -> IResult<I, Option<StunAttribute>, E>
where
    I: nom::Input<Item = u8> + nom::Compare<I> + AsBytes,
{
    let (input, attr_type) = be_u16(input)?;
    let (input, attr_len) = be_u16(input)?;
    let (input, attr_data) = take(attr_len)(input)?;

    if attr_len == 0 {
        return Ok((input, None));
    }

    let (input, _padding) = take(attr_len % 4)(input)?;

    let attr = match attr_type {
        ATTR_MAPPED_ADDRESS => {
            let (_residual, addr) = parse_stun_address(attr_data)?;
            StunAttribute::MappedAddress(addr)
        }
        ATTR_SOURCE_ADDRESS => {
            let (_residual, addr) = parse_stun_address(attr_data)?;
            StunAttribute::SourceAddress(addr)
        }
        ATTR_CHANGED_ADDRESS => {
            let (_residual, addr) = parse_stun_address(attr_data)?;
            StunAttribute::ChangedAddress(addr)
        }
        ATTR_XOR_MAPPED_ADDRESS | ATTR_XOR_MAPPED_ADDRESS_ALT => {
            let (_residual, addr) = parse_stun_xor_address(attr_data, tx_id)?;

            StunAttribute::XorMappedAddress(addr)
        }
        ATTR_USERNAME => {
            let username = String::from_iter(attr_data.iter_elements().map(|b| b as char));
            StunAttribute::Username(username)
        }
        ATTR_MESSAGE_INTEGRITY => {
            let mut msg_integrity = [0u8; 20];
            msg_integrity.copy_from_slice(attr_data.as_bytes());
            StunAttribute::MessageIntegrity(msg_integrity)
        }
        ATTR_ERROR_CODE => {
            let (attr_data, zeros) = take(2usize)(attr_data)?;
            if zeros.iter_elements().any(|b| b != 0) {
                panic!("Invalid STUN error code");
            }
            let (attr_data, err_num) = be_u16(attr_data)?;
            if err_num & 0b1111100000000000 != 0 {
                panic!("Invalid STUN error code");
            }
            let error = String::from_iter(attr_data.iter_elements().map(|b| b as char));
            StunAttribute::ErrorCode((err_num, error))
        }
        ATTR_UNKNOWN_ATTRIBUTES => {
            let (_residual, unknown_attrs) =
                many((attr_len / 2) as usize, be_u16).parse(attr_data)?;
            StunAttribute::UnknownAttributes(unknown_attrs)
        }
        ATTR_NUM_FINGERPRINT => {
            let (_residual, fingerprint) = be_u32(attr_data)?;
            StunAttribute::Fingerprint((fingerprint, false))
        }
        ATTR_REALM => {
            let realm = String::from_iter(attr_data.iter_elements().map(|b| b as char));
            StunAttribute::Realm(realm)
        }
        ATTR_NONCE => {
            let nonce = String::from_iter(attr_data.iter_elements().map(|b| b as char));
            StunAttribute::Nonce(nonce)
        }
        ATTR_NUM_SOFTWARE => {
            let software = String::from_iter(attr_data.iter_elements().map(|b| b as char));
            StunAttribute::Software(software)
        }
        ATTR_ALTERNATE_SERVER => {
            let (_residual, addr) = parse_stun_address(attr_data)?;
            StunAttribute::AlternateServer(addr)
        }
        ATTR_RESPONSE_ORIGIN => {
            let (_residual, addr) = parse_stun_address(attr_data)?;
            StunAttribute::ResponseOrigin(addr)
        }
        ATTR_OTHER_ADDRESS => {
            let (_residual, addr) = parse_stun_address(attr_data)?;
            StunAttribute::OtherAddress(addr)
        }

        t => {
            warn!("Unknown STUN attribute type {t}");
            StunAttribute::Unknown((t, attr_data.iter_elements().collect()))
        }
    };

    Ok((input, Some(attr)))
}

fn parse_stun_header<'a, I, E: ParseError<I>>(input: I) -> IResult<I, StunHeader, E>
where
    I: nom::Input<Item = u8> + nom::Compare<I> + nom::Compare<&'a [u8]> + AsBytes,
{
    let (input, msg_type) = parse_stun_message_type(input)?;
    let (input, msg_length) = be_u16(input)?;
    let (input, _) = tag(MAGIC_COOKIE.as_bytes())(input)?;
    let (input, tx_id) = take_txid(input)?;
    Ok((
        input,
        StunHeader {
            msg_type,
            msg_length,
            tx_id,
        },
    ))
}