use std::net::{Ipv4Addr, Ipv6Addr};

use rand::RngCore;
use serde::{Deserialize, Serialize};
use tracing::{trace, warn};

use super::{buffer::PacketBuffer, query::QueryType, Result};

#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub enum DnsRecord {
    UNKNOWN {
        domain: String,
        qtype: u16,
        data_len: u16,
        ttl: u32,
    }, // 0
    A {
        domain: String,
        addr: Ipv4Addr,
        ttl: u32,
    }, // 1
    NS {
        domain: String,
        host: String,
        ttl: u32,
    }, // 2
    CNAME {
        domain: String,
        host: String,
        ttl: u32,
    }, // 5
    SOA {
        domain: String,
        mname: String,
        nname: String,
        serial: u32,
        refresh: u32,
        retry: u32,
        expire: u32,
        minimum: u32,
        ttl: u32,
    }, // 6
    PTR {
        domain: String,
        pointer: String,
        ttl: u32,
    }, // 12
    MX {
        domain: String,
        priority: u16,
        host: String,
        ttl: u32,
    }, // 15
    TXT {
        domain: String,
        text: Vec<String>,
        ttl: u32,
    }, //16
    AAAA {
        domain: String,
        addr: Ipv6Addr,
        ttl: u32,
    }, // 28
    SRV {
        domain: String,
        priority: u16,
        weight: u16,
        port: u16,
        target: String,
        ttl: u32,
    }, // 33
    CAA {
        domain: String,
        flags: u8,
        length: u8,
        tag: String,
        value: String,
        ttl: u32,
    }, // 257
    AR {
        domain: String,
        ttl: u32,
    },
    AAAAR {
        domain: String,
        ttl: u32,
    },
}

impl DnsRecord {
    pub fn read(buffer: &mut PacketBuffer) -> Result<Self> {
        let mut domain = String::new();
        buffer.read_qname(&mut domain)?;

        let qtype_num = buffer.read_u16()?;
        let qtype = QueryType::from_num(qtype_num);
        let _ = buffer.read_u16()?;
        let ttl = buffer.read_u32()?;
        let data_len = buffer.read_u16()?;

        trace!("Reading DNS Record TYPE: {:?}", qtype);

        let header_pos = buffer.pos();

        match qtype {
            QueryType::A => {
                let raw_addr = buffer.read_u32()?;
                let addr = Ipv4Addr::new(
                    ((raw_addr >> 24) & 0xFF) as u8,
                    ((raw_addr >> 16) & 0xFF) as u8,
                    ((raw_addr >> 8) & 0xFF) as u8,
                    (raw_addr & 0xFF) as u8,
                );

                Ok(Self::A { domain, addr, ttl })
            }
            QueryType::AAAA => {
                let raw_addr1 = buffer.read_u32()?;
                let raw_addr2 = buffer.read_u32()?;
                let raw_addr3 = buffer.read_u32()?;
                let raw_addr4 = buffer.read_u32()?;
                let addr = Ipv6Addr::new(
                    ((raw_addr1 >> 16) & 0xFFFF) as u16,
                    (raw_addr1 & 0xFFFF) as u16,
                    ((raw_addr2 >> 16) & 0xFFFF) as u16,
                    (raw_addr2 & 0xFFFF) as u16,
                    ((raw_addr3 >> 16) & 0xFFFF) as u16,
                    (raw_addr3 & 0xFFFF) as u16,
                    ((raw_addr4 >> 16) & 0xFFFF) as u16,
                    (raw_addr4 & 0xFFFF) as u16,
                );

                Ok(Self::AAAA { domain, addr, ttl })
            }
            QueryType::NS => {
                let mut ns = String::new();
                buffer.read_qname(&mut ns)?;

                Ok(Self::NS {
                    domain,
                    host: ns,
                    ttl,
                })
            }
            QueryType::CNAME => {
                let mut cname = String::new();
                buffer.read_qname(&mut cname)?;

                Ok(Self::CNAME {
                    domain,
                    host: cname,
                    ttl,
                })
            }
            QueryType::SOA => {
                let mut mname = String::new();
                buffer.read_qname(&mut mname)?;

                let mut nname = String::new();
                buffer.read_qname(&mut nname)?;

                let serial = buffer.read_u32()?;
                let refresh = buffer.read_u32()?;
                let retry = buffer.read_u32()?;
                let expire = buffer.read_u32()?;
                let minimum = buffer.read_u32()?;

                Ok(Self::SOA {
                    domain,
                    mname,
                    nname,
                    serial,
                    refresh,
                    retry,
                    expire,
                    minimum,
                    ttl,
                })
            }
            QueryType::PTR => {
                let mut pointer = String::new();
                buffer.read_qname(&mut pointer)?;

                Ok(Self::PTR {
                    domain,
                    pointer,
                    ttl,
                })
            }
            QueryType::MX => {
                let priority = buffer.read_u16()?;
                let mut mx = String::new();
                buffer.read_qname(&mut mx)?;

                Ok(Self::MX {
                    domain,
                    priority,
                    host: mx,
                    ttl,
                })
            }
            QueryType::TXT => {
                let mut text = Vec::new();

                loop {
                    let mut s = String::new();
                    buffer.read_string(&mut s)?;

                    if s.len() == 0 {
                        break;
                    } else {
                        text.push(s);
                    }
                }

                Ok(Self::TXT { domain, text, ttl })
            }
            QueryType::SRV => {
                let priority = buffer.read_u16()?;
                let weight = buffer.read_u16()?;
                let port = buffer.read_u16()?;

                let mut target = String::new();
                buffer.read_qname(&mut target)?;

                Ok(Self::SRV {
                    domain,
                    priority,
                    weight,
                    port,
                    target,
                    ttl,
                })
            }
            QueryType::CAA => {
                let flags = buffer.read()?;
                let length = buffer.read()?;

                let mut tag = String::new();
                buffer.read_string_n(&mut tag, length)?;

                let value_len = (data_len as usize) + header_pos - buffer.pos;
                let mut value = String::new();
                buffer.read_string_n(&mut value, value_len as u8)?;

                Ok(Self::CAA {
                    domain,
                    flags,
                    length,
                    tag,
                    value,
                    ttl,
                })
            }
            QueryType::UNKNOWN(_) | _ => {
                buffer.step(data_len as usize)?;

                Ok(Self::UNKNOWN {
                    domain,
                    qtype: qtype_num,
                    data_len,
                    ttl,
                })
            }
        }
    }

    pub fn write(&self, buffer: &mut PacketBuffer) -> Result<usize> {
        let start_pos = buffer.pos();

        trace!("Writing DNS Record {:?}", self);

        match *self {
            Self::A {
                ref domain,
                ref addr,
                ttl,
            } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::A.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;
                buffer.write_u16(4)?;

                let octets = addr.octets();
                buffer.write_u8(octets[0])?;
                buffer.write_u8(octets[1])?;
                buffer.write_u8(octets[2])?;
                buffer.write_u8(octets[3])?;
            }
            Self::NS {
                ref domain,
                ref host,
                ttl,
            } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::NS.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;

                let pos = buffer.pos();
                buffer.write_u16(0)?;

                buffer.write_qname(host)?;

                let size = buffer.pos() - (pos + 2);
                buffer.set_u16(pos, size as u16)?;
            }
            Self::CNAME {
                ref domain,
                ref host,
                ttl,
            } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::CNAME.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;

                let pos = buffer.pos();
                buffer.write_u16(0)?;

                buffer.write_qname(host)?;

                let size = buffer.pos() - (pos + 2);
                buffer.set_u16(pos, size as u16)?;
            }
            Self::SOA {
                ref domain,
                ref mname,
                ref nname,
                serial,
                refresh,
                retry,
                expire,
                minimum,
                ttl,
            } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::SOA.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;

                let pos = buffer.pos();
                buffer.write_u16(0)?;

                buffer.write_qname(mname)?;
                buffer.write_qname(nname)?;
                buffer.write_u32(serial)?;
                buffer.write_u32(refresh)?;
                buffer.write_u32(retry)?;
                buffer.write_u32(expire)?;
                buffer.write_u32(minimum)?;

                let size = buffer.pos() - (pos + 2);
                buffer.set_u16(pos, size as u16)?;
            }
            Self::PTR {
                ref domain,
                ref pointer,
                ttl,
            } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::NS.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;

                let pos = buffer.pos();
                buffer.write_u16(0)?;

                buffer.write_qname(&pointer)?;

                let size = buffer.pos() - (pos + 2);
                buffer.set_u16(pos, size as u16)?;
            }
            Self::MX {
                ref domain,
                priority,
                ref host,
                ttl,
            } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::MX.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;

                let pos = buffer.pos();
                buffer.write_u16(0)?;

                buffer.write_u16(priority)?;
                buffer.write_qname(host)?;

                let size = buffer.pos() - (pos + 2);
                buffer.set_u16(pos, size as u16)?;
            }
            Self::TXT {
                ref domain,
                ref text,
                ttl,
            } => {
                buffer.write_qname(&domain)?;
                buffer.write_u16(QueryType::TXT.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;

                let pos = buffer.pos();
                buffer.write_u16(0)?;

                if text.is_empty() {
                    return Ok(buffer.pos() - start_pos);
                }

                for s in text {
                    buffer.write_u8(s.len() as u8)?;
                    buffer.write_string(&s)?;
                }
                let size = buffer.pos() - (pos + 2);
                buffer.set_u16(pos, size as u16)?;
            }
            Self::AAAA {
                ref domain,
                ref addr,
                ttl,
            } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::AAAA.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;
                buffer.write_u16(16)?;

                for octet in &addr.segments() {
                    buffer.write_u16(*octet)?;
                }
            }
            Self::SRV {
                ref domain,
                priority,
                weight,
                port,
                ref target,
                ttl,
            } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::SRV.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;

                let pos = buffer.pos();
                buffer.write_u16(0)?;

                buffer.write_u16(priority)?;
                buffer.write_u16(weight)?;
                buffer.write_u16(port)?;
                buffer.write_qname(target)?;

                let size = buffer.pos() - (pos + 2);
                buffer.set_u16(pos, size as u16)?;
            }
            Self::CAA {
                ref domain,
                flags,
                length,
                ref tag,
                ref value,
                ttl,
            } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::CAA.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;

                let pos = buffer.pos();
                buffer.write_u16(0)?;

                buffer.write_u8(flags)?;
                buffer.write_u8(length)?;
                buffer.write_string(tag)?;
                buffer.write_string(value)?;

                let size = buffer.pos() - (pos + 2);
                buffer.set_u16(pos, size as u16)?;
            }
            Self::AR { ref domain, ttl } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::A.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;
                buffer.write_u16(4)?;

                let mut rand = rand::thread_rng();
                buffer.write_u32(rand.next_u32())?;
            }
            Self::AAAAR { ref domain, ttl } => {
                buffer.write_qname(domain)?;
                buffer.write_u16(QueryType::A.to_num())?;
                buffer.write_u16(1)?;
                buffer.write_u32(ttl)?;
                buffer.write_u16(4)?;

                let mut rand = rand::thread_rng();
                buffer.write_u32(rand.next_u32())?;
                buffer.write_u32(rand.next_u32())?;
                buffer.write_u32(rand.next_u32())?;
                buffer.write_u32(rand.next_u32())?;
            }
            Self::UNKNOWN { .. } => {
                warn!("Skipping record: {self:?}");
            }
        }

        Ok(buffer.pos() - start_pos)
    }

    pub fn get_domain(&self) -> String {
        self.get_shared_domain().0
    }

    pub fn get_qtype(&self) -> QueryType {
        self.get_shared_domain().1
    }

    pub fn get_ttl(&self) -> u32 {
        self.get_shared_domain().2
    }

    fn get_shared_domain(&self) -> (String, QueryType, u32) {
        match self {
            DnsRecord::UNKNOWN {
                domain, ttl, qtype, ..
            } => (domain.clone(), QueryType::UNKNOWN(*qtype), *ttl),
            DnsRecord::AAAA { domain, ttl, .. } => (domain.clone(), QueryType::AAAA, *ttl),
            DnsRecord::A { domain, ttl, .. } => (domain.clone(), QueryType::A, *ttl),
            DnsRecord::NS { domain, ttl, .. } => (domain.clone(), QueryType::NS, *ttl),
            DnsRecord::CNAME { domain, ttl, .. } => (domain.clone(), QueryType::CNAME, *ttl),
            DnsRecord::SOA { domain, ttl, .. } => (domain.clone(), QueryType::SOA, *ttl),
            DnsRecord::PTR { domain, ttl, .. } => (domain.clone(), QueryType::PTR, *ttl),
            DnsRecord::MX { domain, ttl, .. } => (domain.clone(), QueryType::MX, *ttl),
            DnsRecord::TXT { domain, ttl, .. } => (domain.clone(), QueryType::TXT, *ttl),
            DnsRecord::SRV { domain, ttl, .. } => (domain.clone(), QueryType::SRV, *ttl),
            DnsRecord::CAA { domain, ttl, .. } => (domain.clone(), QueryType::CAA, *ttl),
            DnsRecord::AR { domain, ttl, .. } => (domain.clone(), QueryType::AR, *ttl),
            DnsRecord::AAAAR { domain, ttl, .. } => (domain.clone(), QueryType::AAAAR, *ttl),
        }
    }
}