import socket
import threading
import RNS
import RNS.vendor.umsgpack as umsgpack
from .common import APP_NAME, load_or_create_identity

class ProxynetServer:
    def __init__(self, identity_name="proxynet_server"):
        self.identity = load_or_create_identity(identity_name)
        self.destination = RNS.Destination(
            self.identity,
            RNS.Destination.IN,
            RNS.Destination.SINGLE,
            APP_NAME,
            "rproxy"
        )
        self.destination.set_link_established_callback(self.link_established)
        self.active_links = {}

    def announce(self):
        self.destination.announce()
        RNS.log(f"Server announced: {RNS.prettyhexrep(self.destination.hash)}", RNS.LOG_NOTICE)

    def link_established(self, link):
        RNS.log(f"Link established with {RNS.prettyhexrep(link.get_remote_identity().hash if link.get_remote_identity() else b'unknown')}", RNS.LOG_NOTICE)
        link.set_packet_callback(self.packet_received)
        link.set_link_closed_callback(self.link_closed)

    def link_closed(self, link):
        RNS.log("Link closed", RNS.LOG_NOTICE)
        # Clean up any associated sockets
        if link in self.active_links:
            target_sock = self.active_links[link].get('target_sock')
            if target_sock:
                try:
                    target_sock.close()
                except:
                    pass
            del self.active_links[link]

    def packet_received(self, message, packet):
        try:
            data = umsgpack.unpackb(message)
            msg_type = data.get('type')

            if msg_type == 'connect':
                self.handle_connect(data, packet.link)
            elif msg_type == 'data':
                self.handle_data(data, packet.link)
            elif msg_type == 'udp_packet':
                self.handle_udp_packet(data, packet.link)
            elif msg_type == 'close':
                self.handle_close(packet.link)

        except Exception as e:
            RNS.log(f"Error processing packet: {e}", RNS.LOG_ERROR)

    def handle_udp_packet(self, data, link):
        host = data.get('host')
        port = data.get('port')
        payload = data.get('payload')
        
        try:

            udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            udp_sock.settimeout(2.0)
            udp_sock.sendto(payload, (host, port))
            
            # Try to receive a response
            try:
                response, addr = udp_sock.recvfrom(65535)
                # Send back to client
                resp_msg = umsgpack.packb({
                    'type': 'udp_response',
                    'host': addr[0],
                    'port': addr[1],
                    'payload': response
                })
                RNS.Packet(link, resp_msg).send()
            except socket.timeout:
                pass
            finally:
                udp_sock.close()
                
        except Exception as e:
            RNS.log(f"UDP handling error: {e}", RNS.LOG_DEBUG)

    def handle_connect(self, data, link):
        host = data.get('host')
        port = data.get('port')
        RNS.log(f"Connection request to {host}:{port}", RNS.LOG_INFO)

        try:
            target_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            target_sock.settimeout(10)
            target_sock.connect((host, port))
            target_sock.settimeout(None)

            self.active_links[link] = {
                'target_sock': target_sock,
                'active': True
            }

            # Start a thread to read from the target socket and send back over Reticulum
            threading.Thread(target=self.target_reader, args=(link, target_sock), daemon=True).start()

            # Send success response
            response = umsgpack.packb({'type': 'response', 'success': True})
            RNS.Packet(link, response).send()

        except Exception as e:
            RNS.log(f"Failed to connect to {host}:{port}: {e}", RNS.LOG_ERROR)
            response = umsgpack.packb({'type': 'response', 'success': False, 'error': str(e)})
            RNS.Packet(link, response).send()
            link.teardown()

    def handle_data(self, data, link):
        if link in self.active_links:
            target_sock = self.active_links[link]['target_sock']
            try:
                target_sock.sendall(data.get('payload'))
            except Exception as e:
                RNS.log(f"Error sending to target: {e}", RNS.LOG_ERROR)
                self.handle_close(link)

    def handle_close(self, link):
        if link in self.active_links:
            target_sock = self.active_links[link]['target_sock']
            try:
                target_sock.close()
            except:
                pass
            self.active_links[link]['active'] = False
            del self.active_links[link]
        link.teardown()

    def target_reader(self, link, target_sock):
        try:
            while True:
                data = target_sock.recv(16384)
                if not data:
                    break
                
                # Split data into chunks if necessary, though Reticulum Link handles some of this
                # For now, just send it. RNS.Link handles mtu.
                msg = umsgpack.packb({'type': 'data', 'payload': data})
                RNS.Packet(link, msg).send()
        except Exception as e:
            RNS.log(f"Target reader error: {e}", RNS.LOG_DEBUG)
        finally:
            self.handle_close(link)

def run_server(identity_name="proxynet_server"):
    server = ProxynetServer(identity_name)
    server.announce()
    
    import time
    while True:
        time.sleep(600)
        server.announce()