from __future__ import annotations
import re
from typing import Dict, List

from dns_client import query_txt
from models import DmarcInfo, SpfInfo, DkimInfo


def parse_kv_tags(record: str) -> Dict[str, str]:
    tags: Dict[str, str] = {}
    for part in record.split(";"):
        part = part.strip()
        if not part or "=" not in part:
            continue
        k, v = part.split("=", 1)
        tags[k.strip().lower()] = v.strip()
    return tags


def analyze_dmarc(txt_records: List[str]) -> DmarcInfo:
    info = DmarcInfo(raw_records=txt_records[:])
    dmarc_recs = [r for r in txt_records if "v=dmarc1" in r.lower()]
    if not dmarc_recs:
        info.present = False
        return info

    info.present = True
    if len(dmarc_recs) > 1:
        info.ambiguous = True

    rec = dmarc_recs[0]
    try:
        tags = parse_kv_tags(rec)
        if tags.get("v", "").upper() != "DMARC1":
            info.parse_error = "Invalid DMARC version tag"
        p = tags.get("p")
        if p:
            info.policy = p.lower()
        pct = tags.get("pct")
        if pct and pct.isdigit():
            info.pct = int(pct)
        rua = tags.get("rua")
        info.rua_present = bool(rua and "mailto:" in rua.lower())
        info.adkim = tags.get("adkim")
        info.aspf = tags.get("aspf")
    except Exception as e:
        info.parse_error = f"Exception parsing DMARC: {e}"
    return info


def analyze_spf(txt_records: List[str]) -> SpfInfo:
    info = SpfInfo(raw_records=txt_records[:])
    spf_recs = [r for r in txt_records if "v=spf1" in r.lower()]
    info.spf_records_only = spf_recs[:]
    if not spf_recs:
        info.present = False
        return info

    info.present = True
    if len(spf_recs) > 1:
        info.multiple = True

    rec = spf_recs[0]
    try:
        m = re.findall(r"(^|\s)([+\-~?]?)(all)($|\s)", rec, flags=re.IGNORECASE)
        if m:
            qual = m[-1][1]
            info.all_qualifier = qual if qual else "+"
        else:
            info.all_qualifier = None

        if info.all_qualifier in ("~", "?", "+"):
            info.weak_all = True
    except Exception as e:
        info.parse_error = f"Exception parsing SPF: {e}"
    return info


COMMON_SELECTORS = ["default", "selector1", "selector2", "s1", "s2", "google", "smtp", "mail", "dkim"]


def probe_dkim_selectors(resolver, etld1: str, enabled: bool) -> DkimInfo:
    info = DkimInfo(probed=enabled)
    if not enabled:
        info.notes = "DKIM probing disabled (no email context)."
        return info

    found: List[str] = []
    for sel in COMMON_SELECTORS:
        name = f"{sel}._domainkey.{etld1}"
        txt = query_txt(resolver, name)
        if any("v=dkim1" in r.lower() for r in txt):
            found.append(sel)

    info.found_selectors = found
    if not found:
        info.notes = "No DKIM record found for common selectors (not proof of absence)."
    return info
