import dns
import httpx
from dataclasses import dataclass, field
from dns import rcode as dnsrcode
from dns import reversename
from dns.message import Message, QueryMessage
from dns.rrset import RRset
from functools import cached_property
from typing import Self, Optional, List

from fmtr.tools.string_tools import join

TTL_CODE_DEFAULTS = {
    dnsrcode.NOERROR: 300,  # Successful query
    dnsrcode.FORMERR: 60,  # Format error
    dnsrcode.SERVFAIL: 10,  # Server failure
    dnsrcode.NXDOMAIN: 60 * 60,  # Non-existent domain
    dnsrcode.NOTIMP: 60,  # Not implemented
    dnsrcode.REFUSED: 60,  # Refused
    dnsrcode.YXDOMAIN: 600,  # Name exists when it should not
    dnsrcode.YXRRSET: 600,  # RR Set exists when it should not
    dnsrcode.NXRRSET: 300,  # RR Set that should exist does not
    dnsrcode.NOTAUTH: 60,  # Not authorized
    dnsrcode.NOTZONE: 60  # Name not contained in zone
}

@dataclass
class BaseDNSData:
    """

    DNS response object.

    """
    wire: bytes

    @cached_property
    def message(self) -> Message:
        return dns.message.from_wire(self.wire)

    @classmethod
    def from_message(cls, message: Message) -> Self:
        return cls(message.to_wire())

@dataclass
class Response(BaseDNSData):
    """

    DNS response object.

    """

    http: Optional[httpx.Response] = None
    blocked_by: Optional[str] = None

    @classmethod
    def from_http(cls, response: httpx.Response) -> Self:
        """

        Initialise from an HTTP response.

        """
        self = cls(response.content, http=response)
        return self

    @property
    def answer(self) -> Optional[RRset]:
        """

        Get the latest answer, if one exists.

        """
        if not self.message.answer:
            return None
        return self.message.answer[-1]

    @property
    def rcode(self) -> dnsrcode.Rcode:
        return self.message.rcode()

    @property
    def rcode_text(self) -> str:
        return dnsrcode.to_text(self.rcode)

    @property
    def ttl(self) -> int:
        """

        Get median TTL from answers, falling back to authority, then to error-code defaults.

        """
        answers = self.message.answer or self.message.authority
        if answers:
            ttls = [answer.ttl for answer in answers]
            ttl = min(ttls)
            return ttl

        ttl = TTL_CODE_DEFAULTS.get(self.rcode, dnsrcode.NXDOMAIN)
        return ttl



    def __str__(self):
        """

        Put answer and ID text in string representation.

        """
        answer = self.answer

        if answer:
            answer = join(answer.to_text().splitlines(), sep=', ')

        string = join([answer, self.message.flags], sep=', ')
        string = f'{self.__class__.__name__}({string})'
        return string



@dataclass
class Request(BaseDNSData):
    """

    DNS request object.

    """
    wire: bytes

    @cached_property
    def question(self) -> RRset:
        return self.message.question[0]

    @cached_property
    def is_valid(self):
        return len(self.message.question) != 0

    @cached_property
    def type(self):
        return self.question.rdtype

    @cached_property
    def type_text(self):
        return dns.rdatatype.to_text(self.type)

    @cached_property
    def name(self):
        return self.question.name

    @cached_property
    def name_text(self):
        return self.name.to_text()

    def get_response_template(self):
        message = dns.message.make_response(self.message)
        message.flags |= dns.flags.RA
        return message

    @cached_property
    def blackhole(self) -> Response:
        blackhole = self.get_response_template()
        blackhole.set_rcode(dns.rcode.NXDOMAIN)
        response = Response.from_message(blackhole)
        return response


@dataclass
class Exchange:
    """

    Entire DNS exchange for a DNS Proxy: request -> upstream response -> response

    """
    ip: str
    port: int

    request: Request
    response: Optional[Response] = None
    answers_pre: List[RRset] = field(default_factory=list)
    is_internal: bool = False
    client_name: Optional[str] = None
    is_complete: bool = False

    @property
    def addr(self):
        return self.ip, self.port

    @classmethod
    def from_wire(cls, wire: bytes, **kwargs) -> Self:
        request = Request(wire)
        response = Response.from_message(request.get_response_template())
        return cls(request=request, response=response, **kwargs)

    @cached_property
    def client(self):
        return f'{self.ip}:{self.port}'

    @property
    def question_last(self) -> RRset:
        """

        Create an RRset surrogate representing the latest/current question.
        This can be the original question - or a hybrid one if we've injected our own answers into the Exchange.
        If there's a response, use its answers, else fall back to answers_pre, else to the original question.

        """
        answers = self.answers_pre
        if self.response:
            answers = self.response.message.answer or answers

        if answers:
            rrset = answers[-1]
            rdtype = self.request.type
            ttl = self.request.question.ttl
            rdclass = self.request.question.rdclass
            name = next(iter(rrset.items.keys())).to_text()
            rrset_surrogate = dns.rrset.from_text(
                name=name,
                ttl=ttl,
                rdtype=rdtype,
                rdclass=rdclass,
            )

            return rrset_surrogate
        else:
            return self.request.question

    @property
    def query_last(self) -> QueryMessage:
        """

        Create a query (e.g. for use by upstream) based on the last question.

        """

        question_last = self.question_last
        query = dns.message.make_query(qname=question_last.name, rdclass=question_last.rdclass, rdtype=question_last.rdtype, id=self.request.message.id)
        return query

    @property
    def key(self):
        """

        Hashable key for caching

        """
        data = tuple(self.request.question.to_text().split())
        return data

    @cached_property
    def reverse(self) -> Self:
        """

        Create an Exchange for a reverse lookup of this Exchange's client IP.

        """
        name = reversename.from_address(self.ip)
        query = dns.message.make_query(name, dns.rdatatype.PTR)
        exchange = self.__class__.from_wire(query.to_wire(), ip=self.ip, port=self.port, is_internal=True)
        return exchange
