use std::{
  fmt::Write,
  net::{IpAddr, Ipv4Addr, Ipv6Addr},
};

use ip_set::{IpRange, Ipv4Map, Ipv6Map};

use crate::{Answer, QType, Query};

const V_SPF1: &str = "v=spf1";
const REDIRECT: &str = "redirect=";
const EXP: &str = "exp=";
const IP4: &str = "ip4:";
const IP6: &str = "ip6:";
const INCLUDE: &str = "include:";
const EXISTS: &str = "exists:";
const ALL: &str = "all";
const A: &str = "a";
const MX: &str = "mx";

const PREFIX_V4: u8 = 32;
const PREFIX_V6: u8 = 128;
const MAX_LOOKUPS: u8 = 16;
const MAX_MX: usize = 10;

/// SPF check result (RFC 7208) / SPF 检查结果
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Status {
  Pass,
  Fail,
  SoftFail,
  Neutral,
  None,
  TempError,
  PermError,
}

#[cfg(feature = "cache")]
#[static_init::dynamic]
pub static CACHE: crate::Cache<Spf> = crate::Cache::new(600);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Act {
  #[default]
  Pass,
  Fail,
  SoftFail,
  Neutral,
}

#[derive(Debug, Clone)]
pub struct Host {
  pub act: Act,
  pub kind: HostKind,
}

#[derive(Debug, Clone)]
pub enum HostKind {
  A(Option<String>, u8),
  Mx(Option<String>, u8),
  Include(String),
  Exists(String),
  Redirect(String),
}

#[derive(Debug, Clone)]
pub struct Spf {
  pub ip4: Ipv4Map<Act>,
  pub ip6: Ipv6Map<Act>,
  pub all: Option<Act>,
  pub host: Vec<Host>,
  pub ttl: u64,
}

impl Spf {
  pub async fn verify<Q: Query>(q: &Q, domain: &str, ip: IpAddr) -> Status {
    let mut lookups = 0u8;
    let r = Self::check_host(q, domain, ip, &mut lookups).await;

    if !matches!(r, Status::Pass | Status::PermError | Status::TempError) {
      let prefix = if ip.is_ipv4() { PREFIX_V4 } else { PREFIX_V6 };
      if let DnsResult::Match = mx_match(q, domain, prefix, ip, &mut lookups).await {
        return Status::Pass;
      }
    }
    r
  }

  async fn check_host<Q: Query>(q: &Q, domain: &str, ip: IpAddr, lookups: &mut u8) -> Status {
    *lookups += 1;
    if *lookups > MAX_LOOKUPS {
      log::warn!("SPF {domain}: DNS lookup limit exceeded");
      return Status::PermError;
    }

    let li = match CACHE.query(q, domain).await {
      Ok(Some(li)) => li,
      Ok(None) => return Status::None,
      Err(e) => {
        log::warn!("SPF {domain}: {e}");
        return Status::TempError;
      }
    };

    if li.len() > 1 {
      log::warn!("SPF {domain}: multiple SPF records");
      return Status::PermError;
    }

    match li.into_iter().next() {
      Some(spf) => spf.eval(q, domain, ip, lookups).await,
      None => Status::None,
    }
  }

  fn check(&self, ip: IpAddr) -> Option<Act> {
    match ip {
      IpAddr::V4(addr) => self.check_v4(addr),
      IpAddr::V6(addr) => self.check_v6(addr),
    }
  }

  fn check_v4(&self, addr: Ipv4Addr) -> Option<Act> {
    self.ip4.get(addr).or(if self.host.is_empty() { self.all } else { None })
  }

  fn check_v6(&self, addr: Ipv6Addr) -> Option<Act> {
    self.ip6.get(addr).or(if self.host.is_empty() { self.all } else { None })
  }

  async fn eval<Q: Query>(&self, q: &Q, domain: &str, ip: IpAddr, lookups: &mut u8) -> Status {
    if let Some(act) = self.check(ip) {
      return act.into();
    }

    for h in &self.host {
      let result = match &h.kind {
        HostKind::Include(d) => {
          let r = Box::pin(Self::check_host(q, d, ip, lookups)).await;
          match r {
            Status::Pass => Some(h.act),
            Status::None | Status::PermError => return Status::PermError,
            Status::TempError => return Status::TempError,
            _ => None,
          }
        }
        HostKind::Redirect(_) => continue,
        HostKind::A(d, prefix) => {
          *lookups += 1;
          if *lookups > MAX_LOOKUPS {
            return Status::PermError;
          }
          let target = d.as_deref().unwrap_or(domain);
          match host_match(q, target, *prefix, ip).await {
            DnsResult::Match => Some(h.act),
            DnsResult::Error => return Status::TempError,
            _ => None,
          }
        }
        HostKind::Mx(d, prefix) => {
          *lookups += 1;
          if *lookups > MAX_LOOKUPS {
            return Status::PermError;
          }
          let target = d.as_deref().unwrap_or(domain);
          match mx_match(q, target, *prefix, ip, lookups).await {
            DnsResult::Match => Some(h.act),
            DnsResult::Error => return Status::PermError,
            _ => None,
          }
        }
        HostKind::Exists(d) => {
          *lookups += 1;
          if *lookups > MAX_LOOKUPS {
            return Status::PermError;
          }
          let target = expand_macro(d, ip, domain);
          host_exists(q, &target).await.then_some(h.act)
        }
      };

      if let Some(act) = result {
        return act.into();
      }
    }

    for h in &self.host {
      if let HostKind::Redirect(d) = &h.kind {
        let r = Box::pin(Self::check_host(q, d, ip, lookups)).await;
        return if r == Status::None {
          log::warn!("SPF {domain}: redirect {d} not found");
          Status::PermError
        } else {
          r
        };
      }
    }

    self.all.map_or(Status::Neutral, Into::into)
  }
}

enum DnsResult {
  Match,
  NoMatch,
  NotFound,
  Error,
}

async fn host_match<Q: Query>(q: &Q, host: &str, prefix: u8, ip: IpAddr) -> DnsResult {
  match ip {
    IpAddr::V4(target) => match super::a::CACHE.query(q, host).await {
      Ok(Some(li)) if li.iter().any(|a| Ipv4Addr::in_cidr(a.ip, prefix, target)) => {
        DnsResult::Match
      }
      Ok(Some(_)) => DnsResult::NoMatch,
      Ok(None) => DnsResult::NotFound,
      Err(_) => DnsResult::Error,
    },
    IpAddr::V6(target) => match super::aaaa::CACHE.query(q, host).await {
      Ok(Some(li)) if li.iter().any(|a| Ipv6Addr::in_cidr(a.ip, prefix, target)) => {
        DnsResult::Match
      }
      Ok(Some(_)) => DnsResult::NoMatch,
      Ok(None) => DnsResult::NotFound,
      Err(_) => DnsResult::Error,
    },
  }
}

async fn mx_match<Q: Query>(
  q: &Q,
  domain: &str,
  prefix: u8,
  ip: IpAddr,
  lookups: &mut u8,
) -> DnsResult {
  let li = match super::mx::CACHE.query(q, domain).await {
    Ok(Some(li)) => li,
    Ok(None) => return DnsResult::NotFound,
    Err(_) => return DnsResult::Error,
  };

  for (i, mx) in li.iter().enumerate() {
    if i >= MAX_MX {
      log::warn!("SPF {domain}: MX record limit exceeded");
      return DnsResult::Error;
    }
    *lookups += 1;
    if *lookups > MAX_LOOKUPS {
      log::warn!("SPF {domain}: DNS lookup limit exceeded");
      return DnsResult::Error;
    }
    match host_match(q, &mx.server, prefix, ip).await {
      DnsResult::Match => return DnsResult::Match,
      DnsResult::Error => return DnsResult::Error,
      _ => continue,
    }
  }
  DnsResult::NoMatch
}

async fn host_exists<Q: Query>(q: &Q, host: &str) -> bool {
  matches!(super::a::CACHE.query(q, host).await, Ok(Some(li)) if !li.is_empty())
}

fn expand_macro(s: &str, ip: IpAddr, domain: &str) -> String {
  if !s.contains('%') {
    return s.to_string();
  }

  let mut out = String::with_capacity(s.len() * 2);
  let mut chars = s.chars().peekable();

  while let Some(c) = chars.next() {
    if c != '%' {
      out.push(c);
      continue;
    }

    match chars.next() {
      Some('%') => out.push('%'),
      Some('_') => out.push(' '),
      Some('-') => {
        let _ = out.write_str("%20");
      }
      Some('{') => {
        let mut letter = None;
        let mut reverse = false;

        while let Some(&ch) = chars.peek() {
          if ch == '}' {
            chars.next();
            break;
          }
          chars.next();
          match ch {
            'i' | 'I' => letter = Some('i'),
            'd' | 'D' => letter = Some('d'),
            'r' | 'R' => reverse = true,
            _ => {}
          }
        }

        match letter {
          Some('i') if reverse => {
            let _ = out.write_str(&ip_reverse(ip));
          }
          Some('i') => {
            let _ = out.write_str(&ip.to_string());
          }
          Some('d') => {
            let _ = out.write_str(domain);
          }
          _ => {}
        }
      }
      Some(c) => {
        out.push('%');
        out.push(c);
      }
      None => out.push('%'),
    }
  }
  out
}

fn ip_reverse(ip: IpAddr) -> String {
  match ip {
    IpAddr::V4(v4) => {
      let o = v4.octets();
      format!("{}.{}.{}.{}", o[3], o[2], o[1], o[0])
    }
    IpAddr::V6(v6) => {
      let mut out = String::with_capacity(64);
      for b in v6.octets().iter().rev() {
        if !out.is_empty() {
          out.push('.');
        }
        let _ = write!(out, "{:x}.{:x}", b & 0xf, b >> 4);
      }
      out
    }
  }
}

impl From<Act> for Status {
  fn from(act: Act) -> Self {
    match act {
      Act::Pass => Status::Pass,
      Act::Fail => Status::Fail,
      Act::SoftFail => Status::SoftFail,
      Act::Neutral => Status::Neutral,
    }
  }
}

fn parse_act(s: &str) -> (Act, &str) {
  match s.as_bytes().first() {
    Some(b'+') => (Act::Pass, &s[1..]),
    Some(b'-') => (Act::Fail, &s[1..]),
    Some(b'~') => (Act::SoftFail, &s[1..]),
    Some(b'?') => (Act::Neutral, &s[1..]),
    _ => (Act::Pass, s),
  }
}

fn parse_prefix(s: &str, default: u8) -> u8 {
  s.parse().unwrap_or(default)
}

fn parse_ip<T: std::str::FromStr>(s: &str, default: u8) -> Option<(T, u8)> {
  let (addr, prefix) = s
    .split_once('/')
    .map_or((s, default), |(a, p)| (a, parse_prefix(p, default)));
  addr.parse().ok().map(|a| (a, prefix))
}

fn parse_domain_prefix(s: &str) -> (Option<String>, u8) {
  if s.is_empty() {
    return (None, PREFIX_V4);
  }
  if let Some(p) = s.strip_prefix('/') {
    return (None, parse_prefix(p, PREFIX_V4));
  }
  let rest = s.strip_prefix(':').map_or(s, |r| r);
  if let Some((domain, p)) = rest.split_once('/') {
    (Some(domain.into()), parse_prefix(p, PREFIX_V4))
  } else {
    (Some(rest.into()), PREFIX_V4)
  }
}

fn strip_mech<'a>(s: &'a str, mech: &str) -> Option<&'a str> {
  if s == mech {
    return Some("");
  }
  s.strip_prefix(mech)
    .filter(|r| r.starts_with(':') || r.starts_with('/'))
}

impl super::Parse for Spf {
  const QTYPE: QType = QType::TXT;

  fn li(answers: impl IntoIterator<Item = Answer>) -> Vec<Self> {
    answers
      .into_iter()
      .filter(|a| a.type_id == QType::TXT as u16)
      .filter_map(|a| parse_spf(&a.val, a.ttl.into()))
      .collect()
  }
}

struct SpfBuilder {
  ip4: Ipv4Map<Act>,
  ip6: Ipv6Map<Act>,
  all: Option<Act>,
  host: Vec<Host>,
  ttl: u64,
}

impl SpfBuilder {
  fn new(ttl: u64) -> Self {
    Self {
      ip4: Ipv4Map::new(),
      ip6: Ipv6Map::new(),
      all: None,
      host: Vec::new(),
      ttl,
    }
  }

  fn add_ip4(&mut self, act: Act, addr: Ipv4Addr, prefix: u8) {
    self.ip4.add_cidr(addr, prefix, act);
  }

  fn add_ip6(&mut self, act: Act, addr: Ipv6Addr, prefix: u8) {
    self.ip6.add_cidr(addr, prefix, act);
  }

  fn build(self) -> Spf {
    Spf {
      ip4: self.ip4,
      ip6: self.ip6,
      all: self.all,
      host: self.host,
      ttl: self.ttl,
    }
  }
}

fn parse_spf(raw: &str, ttl: u64) -> Option<Spf> {
  let val = raw.trim().trim_matches('"').to_lowercase();
  if !val.starts_with(V_SPF1) {
    return None;
  }

  let mut builder = SpfBuilder::new(ttl);

  for part in val.split_whitespace().skip(1) {
    if let Some(domain) = part.strip_prefix(REDIRECT) {
      builder.host.push(Host {
        act: Act::Pass,
        kind: HostKind::Redirect(domain.into()),
      });
      continue;
    }

    if part.starts_with(EXP) {
      continue;
    }

    let (act, rest) = parse_act(part);

    if rest == ALL {
      builder.all = Some(act);
      continue;
    }

    if let Some(s) = rest.strip_prefix(IP4) {
      if let Some((addr, prefix)) = parse_ip::<Ipv4Addr>(s, PREFIX_V4) {
        builder.add_ip4(act, addr, prefix);
      }
      continue;
    }

    if let Some(s) = rest.strip_prefix(IP6) {
      if let Some((addr, prefix)) = parse_ip::<Ipv6Addr>(s, PREFIX_V6) {
        builder.add_ip6(act, addr, prefix);
      }
      continue;
    }

    if let Some(domain) = rest.strip_prefix(INCLUDE) {
      builder.host.push(Host {
        act,
        kind: HostKind::Include(domain.into()),
      });
      continue;
    }

    if let Some(domain) = rest.strip_prefix(EXISTS) {
      builder.host.push(Host {
        act,
        kind: HostKind::Exists(domain.into()),
      });
      continue;
    }

    if let Some(s) = strip_mech(rest, A) {
      let (domain, prefix) = parse_domain_prefix(s);
      builder.host.push(Host {
        act,
        kind: HostKind::A(domain, prefix),
      });
      continue;
    }

    if let Some(s) = strip_mech(rest, MX) {
      let (domain, prefix) = parse_domain_prefix(s);
      builder.host.push(Host {
        act,
        kind: HostKind::Mx(domain, prefix),
      });
    }
  }

  Some(builder.build())
}
