from __future__ import annotations
import asyncio
import hashlib
import json
import logging
from datetime import datetime, timezone
from dataclasses import dataclass
from typing import Optional
from kademlia.network import Server
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat

logger = logging.getLogger(__name__)

DEFAULT_DHT_PORT = 63012
@dataclass
class PeerInfo:
    peer_id: str
    username: str = ""
    last_seen: str = ""
    noise_pubkey_hex: str = ""
    home_relay: str = ""
    ed25519_pubkey_hex: str = ""
    signature: str = ""
    
    def _get_signable_data(self) -> bytes:
        d = f"{self.peer_id}:{self.username}:{self.last_seen}:{self.noise_pubkey_hex}:{self.home_relay}:{self.ed25519_pubkey_hex}"
        return d.encode()

    def to_json(self) -> str:
        return json.dumps({
            "peer_id": self.peer_id,
            "username": self.username,
            "last_seen": self.last_seen,
            "noise_pubkey_hex": self.noise_pubkey_hex,
            "home_relay": self.home_relay,
            "ed25519_pubkey_hex": self.ed25519_pubkey_hex,
            "signature": self.signature
        })
    @classmethod
    def from_json(cls, data:str) -> "PeerInfo":
        d = json.loads(data)
        return cls(
            peer_id=d["peer_id"],
            username=d.get("username", ""),
            last_seen=d.get("last_seen", ""),
            noise_pubkey_hex=d.get("noise_pubkey_hex", ""),
            home_relay=d.get("home_relay", ""),
            ed25519_pubkey_hex=d.get("ed25519_pubkey_hex", ""),
            signature=d.get("signature", "")
        )
    
    @property
    def fingerprint(self) -> str:
        if not self.noise_pubkey_hex:
            return ""
        import hashlib
        return hashlib.sha256(bytes.fromhex(self.noise_pubkey_hex)).hexdigest()
    
    def is_valid(self) -> bool:
        if not self.ed25519_pubkey_hex or not self.signature:
            return False
        try:
            pub_bytes = bytes.fromhex(self.ed25519_pubkey_hex)
            if hashlib.sha256(pub_bytes).hexdigest() != self.peer_id:
                return False
            pub_key = Ed25519PublicKey.from_public_bytes(pub_bytes)
            pub_key.verify(bytes.fromhex(self.signature), self._get_signable_data())
            return True
        except Exception:
            return False

class DHTNode:
    def __init__(
        self,
        peer_id: str,
        host: str = "0.0.0.0",
        port: int = DEFAULT_DHT_PORT,
        bootstrap_nodes: list[tuple[str, int]] | None = None,
        ed25519_private_key=None
    ):
        self.peer_id = peer_id
        self.host = host
        self.port = port
        self.bootstrap_nodes = bootstrap_nodes or []
        self.ed25519_private_key = ed25519_private_key
        self._server = Server(ksize=8, alpha=5)
        self._started = False
    
    async def start(self) -> None:
        await self._server.listen(self.port)
        self._server.protocol.wait_timeout = 2.0
        self._started = True
        logger.info(f"TNT?? nah DHT node listening on port {self.port}")
        
        if self.bootstrap_nodes:
            asyncio.create_task(self._server.bootstrap(self.bootstrap_nodes))
            logger.info(f"bootsrapped with {self.bootstrap_nodes}")
        else:
            logger.warning("No nodes configured | in local mode")
    
    async def reannounce(self, username: str = "", noise_pubkey_hex: str = "", home_relay: str = "") -> None:
        while self._started:
            await self.announce(username=username, noise_pubkey_hex=noise_pubkey_hex, home_relay=home_relay)
            await asyncio.sleep(60)
    
    async def stop(self) -> None:
        if self._started:
            self._server.stop()
            self._started = False
            logger.info("DHT node stopped")
            
    async def announce(self, username: str = "", noise_pubkey_hex: str = "", home_relay: str = "") -> None:
        info = PeerInfo(
            peer_id=self.peer_id,
            username=username,
            last_seen=datetime.now(timezone.utc).isoformat(),
            noise_pubkey_hex=noise_pubkey_hex,
            home_relay=home_relay
        )
        if self.ed25519_private_key:
            pub_bytes = self.ed25519_private_key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
            info.ed25519_pubkey_hex = pub_bytes.hex()
            info.signature = self.ed25519_private_key.sign(info._get_signable_data()).hex()
        await self._server.set(self.peer_id, info.to_json())
        logger.info(f"Announced self to DHT {self.peer_id[:12]} on relay {home_relay}")
        
    async def lookup(self, peer_id: str) -> Optional[PeerInfo]:
        result = await self._server.get(peer_id)
        if result is None:
            logger.debug(f"peer {peer_id[:12]} not found")
            return None
        try:
            info = PeerInfo.from_json(result)
            if not info.is_valid():
                logger.error(f"DHT signature verification failed for {peer_id[:12]} it is unsigned or invalid signature")
                return None
            return info
        except Exception as e:
            logger.error(f"parsing failed for peer info {peer_id[:12]} {e}")
            return None

    async def is_online(self, peer_id: str, threshold_seconds: int = 90) -> bool:
        info = await self.lookup(peer_id)
        if not info or not info.last_seen:
            return False
        last_seen = datetime.fromisoformat(info.last_seen)
        age = (datetime.now(timezone.utc) - last_seen).total_seconds()
        return age < threshold_seconds
    
    async def get_value(self, key: str) -> str | None:
        return await self._server.get(key)
    
    async def set_value(self, key: str, value: str) -> None:
        await self._server.set(key, value)
        
    async def __aenter__(self):
        await self.start()
        return self
    
    async def __aexit__(self, *_):
        await self.stop()
