"""
NAT Engine Module for VPN
Implements Network Address Translation
"""

import asyncio
import time
from typing import Dict, Optional, Tuple
from dataclasses import dataclass
import logging
import subprocess

logger = logging.getLogger(__name__)

@dataclass
class NATSession:
    virtual_ip: str
    virtual_port: int
    real_ip: str
    real_port: int
    host_port: int
    created_time: float
    last_activity: float
    bytes_in: int = 0
    bytes_out: int = 0

class NATEngine:
    def __init__(self, logger_instance=None):
        self.sessions: Dict[str, NATSession] = {}
        self.port_mappings: Dict[int, Tuple[str, int]] = {}
        self.next_port = 10000
        self.cleanup_interval = 300  # 5 minutes
        self._cleanup_task = None
        self.logger = logger_instance if logger_instance else logging.getLogger(__name__)

    async def start(self):
        """Start the NAT engine"""
        self._cleanup_task = asyncio.create_task(self._cleanup_loop())
        self.logger.info("NAT engine started")

    async def stop(self):
        """Stop the NAT engine"""
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
        self.sessions.clear()
        self.port_mappings.clear()
        self.logger.info("NAT engine stopped")

    def create_session(self, virtual_ip: str, virtual_port: int,
                      real_ip: str, real_port: int) -> NATSession:
        """Create a new NAT session"""
        host_port = self._allocate_port()
        session = NATSession(
            virtual_ip=virtual_ip,
            virtual_port=virtual_port,
            real_ip=real_ip,
            real_port=real_port,
            host_port=host_port,
            created_time=time.time(),
            last_activity=time.time()
        )
        
        session_key = f"{virtual_ip}:{virtual_port}"
        self.sessions[session_key] = session
        self.port_mappings[host_port] = (virtual_ip, virtual_port)
        
        self.logger.debug(f"Created NAT session: {session_key} -> {real_ip}:{real_port}")
        return session

    def lookup_session(self, ip: str, port: int) -> Optional[NATSession]:
        """Look up a NAT session by real IP and port"""
        for session in self.sessions.values():
            if session.real_ip == ip and session.real_port == port:
                return session
        return None

    def get_session_by_virtual(self, ip: str, port: int) -> Optional[NATSession]:
        """Get session by virtual IP and port"""
        return self.sessions.get(f"{ip}:{port}")

    def remove_session(self, session: NATSession):
        """Remove a NAT session"""
        session_key = f"{session.virtual_ip}:{session.virtual_port}"
        if session_key in self.sessions:
            del self.sessions[session_key]
        if session.host_port in self.port_mappings:
            del self.port_mappings[session.host_port]

    def _allocate_port(self) -> int:
        """Allocate a new port for NAT"""
        while self.next_port in self.port_mappings:
            self.next_port += 1
            if self.next_port > 65535:
                self.next_port = 10000
        return self.next_port

    async def _cleanup_loop(self):
        """Periodically clean up expired sessions"""
        while True:
            try:
                await asyncio.sleep(self.cleanup_interval)
                current_time = time.time()
                
                for session_key, session in list(self.sessions.items()):
                    if current_time - session.last_activity > self.cleanup_interval:
                        self.remove_session(session)
                        self.logger.debug(f"Removed expired session: {session_key}")
                        
            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"Error in cleanup loop: {e}")

    def get_stats(self) -> Dict:
        """Get NAT engine statistics"""
        return {
            "active_sessions": len(self.sessions),
            "allocated_ports": len(self.port_mappings),
            "total_bytes_in": sum(s.bytes_in for s in self.sessions.values()),
            "total_bytes_out": sum(s.bytes_out for s in self.sessions.values())
        }




    async def setup_nat(self, interface: str):
        """Setup NAT configuration"""
        # In this Windows implementation, NAT is handled at the application level
        # through the traffic forwarding engine and doesn't require system-level configuration
        logger.info(f"NAT configuration ready for interface {interface}")
        pass