import base64
import http.client
import io
import json
import secrets
import ssl
import time
import urllib.parse
from dataclasses import dataclass

from .transport import GetTransport, TransportError
from .verify import AttestationError, verify_document


@dataclass(frozen=True)
class Response:
    status_code: int
    headers: tuple
    content: bytes

    @property
    def text(self):
        return self.content.decode("utf-8", errors="replace")

    def json(self):
        return json.loads(self.content)


class _TLSFile(io.RawIOBase):
    def __init__(self, stream):
        self.stream = stream

    def readable(self):
        return True

    def readinto(self, buffer):
        data = self.stream.read(len(buffer))
        buffer[:len(data)] = data
        return len(data)


class InnerTLS:
    def __init__(self, transport, server_name):
        self.transport = transport
        self.deadline = transport.deadline
        self.incoming, self.outgoing = ssl.MemoryBIO(), ssl.MemoryBIO()
        context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
        context.check_hostname = False
        # Authentication is deferred to Nitro, before any application request.
        context.verify_mode = ssl.CERT_NONE
        context.minimum_version = ssl.TLSVersion.TLSv1_3
        self.tls = context.wrap_bio(self.incoming, self.outgoing, server_side=False, server_hostname=server_name)
        while True:
            try:
                self.tls.do_handshake()
                self._flush()
                break
            except (ssl.SSLWantReadError, ssl.SSLWantWriteError):
                self._pump()
        self.peer_der = self.tls.getpeercert(binary_form=True)

    def _check_time(self):
        if time.monotonic() > self.deadline:
            raise TimeoutError("inner TLS operation exceeded deadline")

    def _flush(self):
        self._check_time()
        if self.outgoing.pending:
            data = self.transport.send(self.outgoing.read())
            if data:
                self.incoming.write(data)

    def _pump(self):
        self._check_time()
        pending = self.outgoing.pending
        self._flush()
        if not pending:
            data = self.transport.exchange()
            if data:
                self.incoming.write(data)
            elif self.transport.eof:
                self.incoming.write_eof()

    def sendall(self, data):
        offset = 0
        while offset < len(data):
            try:
                # Flush bounded TLS batches so a large POST starts reaching the
                # enclave before its HTTP header deadline expires.
                offset += self.tls.write(data[offset:offset + 16384])
                self._flush()
            except (ssl.SSLWantReadError, ssl.SSLWantWriteError):
                self._pump()

    def read(self, length):
        while True:
            self._check_time()
            try:
                return self.tls.read(length)
            except ssl.SSLZeroReturnError:
                return b""
            except (ssl.SSLWantReadError, ssl.SSLWantWriteError):
                self._pump()

    def makefile(self, mode):
        if mode != "rb":
            raise ValueError("read-only HTTP stream")
        return io.BufferedReader(_TLSFile(self))

    def request(self, path, host, *, limit, method="GET", body=b""):
        if not path.startswith("/") or any(ord(c) < 33 or ord(c) > 126 for c in path + host):
            raise ValueError("HTTP request path must contain escaped ASCII")
        self.deadline = time.monotonic() + self.transport.timeout
        self.transport.deadline = self.deadline
        if method not in ("GET", "POST") or not isinstance(body, bytes) or len(body) > 180*1024:
            raise ValueError("invalid inner request method or body")
        if method == "GET" and body:
            raise ValueError("GET body is not supported")
        extra = f"Content-Type: application/json\r\nContent-Length: {len(body)}\r\n" if method == "POST" else ""
        self.sendall((f"{method} {path} HTTP/1.1\r\nHost: {host}\r\nAccept-Encoding: identity\r\n"
                      f"{extra}Connection: keep-alive\r\n\r\n").encode("ascii") + body)

        response = http.client.HTTPResponse(self)
        response.begin()
        if response.length is not None and response.length > limit:
            raise TransportError("inner HTTP response exceeds limit")
        body = response.read(limit+1)
        if len(body) > limit:
            raise TransportError("inner HTTP response exceeds limit")
        result = Response(response.status, tuple(response.getheaders()), body)
        response.close()
        return result


class Relay:
    """A verified inner TLS connection; no application data is sent before verify()."""
    def __init__(self, endpoint, *, expected_pcr0, require_graviton5=True, timeout=60,
                 insecure_local=False, allow_dev=False):
        if allow_dev and not (insecure_local and urllib.parse.urlsplit(endpoint).hostname in ("127.0.0.1", "localhost", "::1")):
            raise ValueError("development mode is restricted to explicit loopback endpoints")
        if not expected_pcr0 and not (allow_dev and insecure_local):
            raise ValueError("an independently pinned PCR0 is required")
        self.endpoint = endpoint
        self.expected_pcr0 = expected_pcr0
        self.require_graviton5 = require_graviton5
        self.timeout = timeout
        self.insecure_local = insecure_local
        self.allow_dev = allow_dev
        self._stream = None
        self.attestation = None

    def verify(self):
        # Explicit verification establishes a new session with a new nonce.
        self._stream = None
        self.attestation = None
        transport = GetTransport(self.endpoint, timeout=self.timeout, insecure_local=self.insecure_local)
        host = urllib.parse.urlsplit(self.endpoint).hostname
        stream = InnerTLS(transport, host)
        nonce = secrets.token_bytes(32)
        response = stream.request("/v1/attestation?nonce=" + base64.urlsafe_b64encode(nonce).rstrip(b"=").decode(), host, limit=65536)
        if response.status_code != 200:
            raise AttestationError(f"attestation returned HTTP {response.status_code}")
        data = response.json()
        if data.get("mode") == "dev":
            if not (self.allow_dev and self.insecure_local):
                raise AttestationError("development enclave provides no attestation")
            self.attestation = {"verified": False, "mode": "dev", "policy": data.get("policy", {})}
        else:
            raw = base64.b64decode(data["attestation_document_b64"], validate=True)
            self.attestation = verify_document(raw, nonce=nonce, peer_der=stream.peer_der,
                                               expected_pcr0=self.expected_pcr0,
                                               require_graviton5=self.require_graviton5)
        self._stream = stream
        return self.attestation

    def get(self, url):
        target = urllib.parse.urlsplit(url)
        if target.scheme != "https" or not target.hostname or target.username or target.password or target.fragment:
            raise ValueError("target must be an HTTPS URL without credentials or fragment")
        if self._stream is None:
            self.verify()
        path = "/f/https/" + target.netloc + (target.path or "/")
        if target.query:
            path += "?" + target.query
        try:
            return self._stream.request(path, urllib.parse.urlsplit(self.endpoint).hostname, limit=10*1024*1024)
        except Exception:
            self._stream = None
            raise


    def send_request(self, url, *, method="GET", headers=None, body=b""):
        """Forward one HTTPS request. No automatic redirects or application retries.

        The outer transport still uses only GET. A failure after dispatch can
        mean the destination acted but its response was lost; do not blindly
        retry a non-idempotent request.
        """
        if method not in ("GET", "HEAD", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"):
            raise ValueError("unsupported method")
        if isinstance(body, str):
            body = body.encode("utf-8")
        if not isinstance(body, bytes) or len(body) > 102400:
            raise ValueError("request body must be at most 100 KiB")
        if method in ("GET", "HEAD") and body:
            raise ValueError("GET and HEAD bodies are unsupported")
        target = urllib.parse.urlsplit(url)
        if target.scheme != "https" or not target.hostname or target.username or target.password or target.fragment:
            raise ValueError("target must be HTTPS without credentials or fragment")
        pairs = list((headers or {}).items()) if isinstance(headers, dict) else list(headers or [])
        raw = json.dumps({"url":url,"method":method,"headers":pairs,
                          "body_b64":base64.b64encode(body).decode("ascii")},
                         ensure_ascii=False,separators=(",", ":")).encode("utf-8")
        if len(raw) > 180*1024:
            raise ValueError("request command too large")
        if self._stream is None:
            self.verify()
        try:
            result = self._stream.request("/v1/commands/send_request", urllib.parse.urlsplit(self.endpoint).hostname,
                                          limit=15*1024*1024,method="POST",body=raw)
            if result.status_code != 200:
                raise TransportError(f"request command returned HTTP {result.status_code}")
            value = result.json()
            content = base64.b64decode(value["body_b64"],validate=True)
            if len(content) > 10*1024*1024 or type(value["status"]) is not int or not 100 <= value["status"] <= 599:
                raise TransportError("invalid request response")
            response_headers = tuple((k,base64.b64decode(v,validate=True).decode("latin-1")) for k,v in value["headers"])
            response_headers += tuple(("x-attested-relay-"+key,value[key]) for key in ("record","puzzle","evidence"))
            return Response(value["status"],response_headers,content)
        except Exception:
            self._stream = None
            raise

    def write_pastebin(self, tag, content):
        """Append a paste under a shared tag password."""
        return self.write_paste(tag, content)

    def read_pastebin_tag(self, tag, *, after=None, limit=10):
        """Read a page of pastes under a shared tag password."""
        return self.read_pastes(tag, after=after, limit=limit)

    @staticmethod
    def new_paste_tag():
        """Create a random shared read/write password; exchange it privately."""
        return secrets.token_urlsafe(32)

    def _paste_request(self, path, tag, fields):
        if not isinstance(tag, str) or not 1 <= len(tag.encode("utf-8")) <= 256:
            raise ValueError("tag must be 1 through 256 UTF-8 bytes")
        if self._stream is None:
            self.verify()
        policy = self.attestation.get("policy", {})
        if policy.get("paste_protocol_version") != 1 or policy.get("paste_max_bytes") != 102400:
            raise AttestationError("measured enclave does not support this paste protocol")
        raw = json.dumps({"tag": tag, **fields}, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
        try:
            result = self._stream.request(path, urllib.parse.urlsplit(self.endpoint).hostname,
                                          limit=1500000, method="POST", body=raw)
            if result.status_code not in (200, 201):
                raise TransportError(f"paste operation returned HTTP {result.status_code}")
            return result.json()
        except Exception:
            self._stream = None
            raise

    def write_paste(self, tag, content):
        """Append up to 100 KiB; knowing the tag grants immediate read/write access."""
        if isinstance(content, str):
            content = content.encode("utf-8")
        if not isinstance(content, bytes) or len(content) > 102400:
            raise ValueError("paste content must be bytes or text of at most 100 KiB")
        return self._paste_request("/v1/commands/write_pastebin", tag, {"content_b64": base64.b64encode(content).decode("ascii")})

    def read_pastes(self, tag, *, after=None, limit=10):
        """Read one authenticated page; the untrusted host can omit entries.

        Cursors are hash order, not a live subscription: begin again at the first
        page when polling to discover new entries that sort before an old cursor.
        """
        import re
        if type(limit) is not int or not 1 <= limit <= 10:
            raise ValueError("paste page limit must be 1 through 10")
        if after is not None and (not isinstance(after, str) or not re.fullmatch(r"[0-9a-f]{64}\.paste\.json", after)):
            raise ValueError("invalid paste cursor")
        result = self._paste_request("/v1/commands/read_pastebin_tag", tag, {"after": after, "limit": limit})
        if not isinstance(result, dict) or set(result) != {"pastes", "next_cursor"} or not isinstance(result["pastes"], list) or len(result["pastes"]) > limit:
            raise TransportError("invalid paste page")
        for item in result["pastes"]:
            item["content"] = base64.b64decode(item.pop("content_b64"), validate=True)
            if len(item["content"]) > 102400:
                raise TransportError("oversized paste")
        return result