"""Gemeinsame Bausteine: Verlauf, aktive Verbindungen, Relay, Ausgangsverbindungen."""
import asyncio
import ipaddress
import itertools
import logging
import socket
import time
from dataclasses import dataclass, field

from . import config
from .db import db
from .policy import Decision, policy

log = logging.getLogger("proxy")


# --------------------------------------------------------------------------- Verlauf
class HistoryLogger:
    COLUMNS = ("ts", "username", "client_ip", "proto", "method", "host", "port", "url",
               "status", "detail", "bytes_up", "bytes_down", "duration")

    def __init__(self):
        self.buffer: list[tuple] = []

    def add(self, **e) -> None:
        if policy.settings.get("log_enabled", "1") != "1":
            return
        e.setdefault("ts", time.time())
        self.buffer.append(tuple(e.get(c, 0 if c in ("bytes_up", "bytes_down", "duration") else None)
                                 for c in self.COLUMNS))
        if len(self.buffer) > 50_000:  # Notbremse, falls die DB hängt
            del self.buffer[:10_000]

    def flush(self) -> None:
        if not self.buffer:
            return
        rows, self.buffer = self.buffer, []
        db.executemany(f"INSERT INTO history({','.join(self.COLUMNS)}) VALUES({','.join('?' * len(self.COLUMNS))})",
                       rows)


history = HistoryLogger()


# --------------------------------------------------------------------------- aktive Verbindungen
_ids = itertools.count(1)


@dataclass
class Conn:
    username: str | None
    client_ip: str
    proto: str
    host: str
    port: int
    method: str = ""
    url: str = ""
    id: int = field(default_factory=lambda: next(_ids))
    started: float = field(default_factory=time.time)
    bytes_up: int = 0
    bytes_down: int = 0
    tasks: list = field(default_factory=list)
    closers: list = field(default_factory=list)

    def to_dict(self) -> dict:
        return {"id": self.id, "username": self.username, "client_ip": self.client_ip, "proto": self.proto,
                "host": self.host, "port": self.port, "method": self.method, "url": self.url,
                "started": self.started, "bytes_up": self.bytes_up, "bytes_down": self.bytes_down}

    def kill(self) -> None:
        for t in self.tasks:
            t.cancel()
        for c in self.closers:
            try:
                c()
            except Exception:
                pass

    def finish(self, status: str = "ok", detail: str = "") -> None:
        registry.pop(self.id, None)
        history.add(ts=self.started, username=self.username, client_ip=self.client_ip, proto=self.proto,
                    method=self.method, host=self.host, port=self.port, url=self.url, status=status,
                    detail=detail, bytes_up=self.bytes_up, bytes_down=self.bytes_down,
                    duration=round(time.time() - self.started, 3))


registry: dict[int, Conn] = {}


def register(conn: Conn) -> Conn:
    registry[conn.id] = conn
    return conn


# --------------------------------------------------------------------------- Relay
async def _pipe(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, conn: Conn, up: bool) -> None:
    try:
        while True:
            data = await reader.read(65536)
            if not data:
                break
            writer.write(data)
            if up:
                conn.bytes_up += len(data)
            else:
                conn.bytes_down += len(data)
            await writer.drain()
    except (ConnectionError, OSError):
        pass
    finally:
        try:
            if writer.can_write_eof() and not writer.is_closing():
                writer.write_eof()
        except (OSError, RuntimeError):
            pass


async def relay(conn: Conn, c_reader, c_writer, r_reader, r_writer) -> None:
    """Leitet Daten in beide Richtungen weiter, bis eine Seite schließt."""
    up = asyncio.create_task(_pipe(c_reader, r_writer, conn, True))
    down = asyncio.create_task(_pipe(r_reader, c_writer, conn, False))
    conn.tasks += [up, down]
    conn.closers += [c_writer.close, r_writer.close]
    try:
        done, pending = await asyncio.wait({up, down}, return_when=asyncio.FIRST_COMPLETED)
        if pending:
            # Halb geschlossen: der anderen Richtung noch etwas Zeit geben
            await asyncio.wait(pending, timeout=60)
            for t in pending:
                t.cancel()
    finally:
        for w in (c_writer, r_writer):
            try:
                w.close()
            except Exception:
                pass


# --------------------------------------------------------------------------- Ausgang
class Blocked(Exception):
    def __init__(self, decision: Decision):
        super().__init__(decision.reason)
        self.decision = decision


async def resolve(host: str, port: int, kind=socket.SOCK_STREAM) -> list[tuple]:
    loop = asyncio.get_running_loop()
    infos = await loop.getaddrinfo(host.strip("[]"), port, type=kind)
    seen, out = set(), []
    for family, _, _, _, sockaddr in infos:
        if sockaddr[0] not in seen:
            seen.add(sockaddr[0])
            out.append((family, sockaddr))
    return out


async def check_and_resolve(user: dict | None, host: str, port: int, kind=socket.SOCK_STREAM) -> list[tuple]:
    """Prüft Regeln für Hostname + jede aufgelöste IP. Gibt erlaubte Adressen zurück."""
    first = policy.decide(user, host, port)
    if not first.allowed:
        raise Blocked(first)
    try:
        addrs = await resolve(host, port, kind)
    except socket.gaierror as e:
        raise OSError(f"DNS-Auflösung fehlgeschlagen: {e.strerror or e}")
    allowed, last = [], first
    for family, sockaddr in addrs:
        d = policy.decide(user, host, port, ipaddress.ip_address(sockaddr[0].split("%")[0]))
        if d.allowed:
            allowed.append((family, sockaddr))
        else:
            last = d
    if not allowed:
        raise Blocked(last)
    return allowed


async def open_upstream(user: dict | None, host: str, port: int):
    addrs = await check_and_resolve(user, host, port)
    err: Exception | None = None
    for family, sockaddr in addrs:
        try:
            return await asyncio.wait_for(asyncio.open_connection(sockaddr[0], port, family=family),
                                          timeout=config.CONNECT_TIMEOUT)
        except (OSError, asyncio.TimeoutError) as e:
            err = e
    if isinstance(err, asyncio.TimeoutError):
        raise OSError("Zeitüberschreitung beim Verbindungsaufbau")
    raise err or OSError("Verbindung fehlgeschlagen")


def peer_ip(writer: asyncio.StreamWriter) -> str:
    peer = writer.get_extra_info("peername")
    ip = peer[0] if peer else "?"
    if ip.startswith("::ffff:"):
        ip = ip[7:]
    return ip