from core.errors.exceptions import *
from core.errors.logger import logger

from collections.abc import Callable
from core.Constants import Constants
from core.controllers.ConfigurationController import ConfigurationController
from core.controllers.ProfileController import ProfileController
from core.controllers.SessionStateController import SessionStateController
from core.controllers.SystemStateController import SystemStateController
from core.models.session.SessionProfile import SessionProfile
from core.models.system.SystemProfile import SystemProfile
from core.models.system.SystemState import SystemState
from core.observers.ConnectionObserver import ConnectionObserver
from core.services.WebServiceApiService import WebServiceApiService
from essentials.modules.TorModule import TorModule
from essentials.services.ConnectionService import ConnectionService
from pathlib import Path
from core.Errors import InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ConnectionTerminationError, CommandNotFoundError, PaymentRequiredError
from subprocess import CalledProcessError
from typing import Union, Optional, Any
import os
import random
import re
import shutil
import signal
import subprocess
import sys
import tempfile
import time
from core.models.OperatorProxySession import OperatorProxySession


class ConnectionController:

    @staticmethod
    def with_preferred_connection(*args, task: Callable[..., Any], connection_observer: Optional[ConnectionObserver] = None, **kwargs):
        """
        This function does a task with a preferred connection type.

        It is currently being CALLED UPON with keyword based kwargs only, because any args would execute immediately.

        However, the *args keeps it open to future uses.
        """

        connection = ConfigurationController.get_connection()

        if connection == 'system':
            return task(*args, **kwargs)

        elif connection == 'tor':
            return ConnectionController.__with_tor_connection(*args, task=task, connection_observer=connection_observer, **kwargs)

        else:
            return None

    @staticmethod
    def establish_connection(profile: Union[SessionProfile, SystemProfile], ignore: tuple[type[Exception]] = (), connection_observer: Optional[ConnectionObserver] = None):
        from core.controllers.ConnectionController import ConnectionController
        
        ConnectionController.verify_singbox_installation()
        
        connection = profile.connection

        # needs_proxy_configuration comes from the child connection models
        # for the system it returns false blindly. for the session is checks if the connection type is masked.
        if connection.needs_proxy_configuration() and not profile.has_proxy_configuration():

            if profile.has_subscription():

                if not profile.subscription.has_been_activated():
                    ProfileController.activate_subscription(profile, connection_observer=connection_observer)

                proxy_configuration = ConnectionController.with_preferred_connection(profile.subscription.billing_code, task=WebServiceApiService.get_proxy_configuration, connection_observer=connection_observer)

                if proxy_configuration is None:
                    raise InvalidSubscriptionError()

                profile.attach_proxy_configuration(proxy_configuration)

            else:
                raise MissingSubscriptionError()

        # The needs_wireguard_configuration comes from the connection model, and can be transitioned to get it directly from the object's data
        # The has_wireguard_configuration comes from each polymorph object doing an os check on if the wg config exists
        if connection.needs_wireguard_configuration() and not profile.has_wireguard_configuration():

            if profile.has_subscription():

                if not profile.subscription.has_been_activated():
                    ProfileController.activate_subscription(profile, connection_observer=connection_observer)

                ProfileController.register_wireguard_session(profile, connection_observer=connection_observer)

            else:

                if profile.is_system_profile():

                    if ConnectionController.system_uses_wireguard_interface() and SystemStateController.exists():

                        try:
                            ConnectionController.terminate_system_connection()
                        except ConnectionTerminationError:
                            pass

                raise MissingSubscriptionError()

        if connection.needs_operator_proxy() and not profile.has_operator_proxy_session():

            if profile.has_subscription():

                if not profile.subscription.has_been_activated():
                    ProfileController.activate_subscription(profile, connection_observer=connection_observer)

                operator_proxy_session = ConnectionController.with_preferred_connection(
                    profile.subscription.billing_code,
                    profile.subscription.operator_id,
                    connection.get_protocol(),
                    task=WebServiceApiService.post_operator_proxy,
                    connection_observer=connection_observer
                )

                if operator_proxy_session is None:
                    raise PaymentRequiredError()

                profile.attach_operator_proxy_session(operator_proxy_session)

            else:
                raise MissingSubscriptionError()

        if profile.is_session_profile():

            try:
                return ConnectionController.establish_session_connection(profile, ignore=ignore, connection_observer=connection_observer)

            except ConnectionError:

                if ConnectionController.__should_renegotiate(profile):

                    ProfileController.register_wireguard_session(profile, connection_observer=connection_observer)
                    return ConnectionController.establish_session_connection(profile, ignore=ignore, connection_observer=connection_observer)

                else:
                    raise ConnectionError('The connection could not be established.')

        if profile.is_system_profile():

            try:
                return ConnectionController.establish_system_connection(profile, ignore=ignore, connection_observer=connection_observer)

            except ConnectionError:

                if ConnectionController.__should_renegotiate(profile):

                    ProfileController.register_wireguard_session(profile, connection_observer=connection_observer)
                    return ConnectionController.establish_system_connection(profile, ignore=ignore, connection_observer=connection_observer)

                else:
                    raise ConnectionError('The connection could not be established.')

        return None

    @staticmethod
    def establish_session_connection(profile: SessionProfile, ignore: tuple[type[Exception]] = (), connection_observer: Optional[ConnectionObserver] = None):

        session_directory = tempfile.mkdtemp(prefix='hv-')
        session_state = SessionStateController.get_or_new(profile.id)

        port_number = None
        proxy_port_number = None

        # this is a check from SessionConnection of if there's a systemwide with mask
        if profile.connection.is_unprotected():

            if not ConnectionController.system_uses_wireguard_interface():

                if not ConnectionUnprotectedError in ignore:
                    raise ConnectionUnprotectedError('Connection unprotected while the system is not using a WireGuard interface.')
                else:
                    ProfileController.disable(profile)

        if profile.connection.code == 'tor':

            port_number = ConnectionService.get_random_available_port_number()
            ConnectionController.establish_tor_session_connection(port_number, connection_observer=connection_observer)

            session_state.network_port_numbers.tor.append(port_number)

        elif profile.connection.code == 'wireguard':

            if ConfigurationController.get_endpoint_verification_enabled():
                ProfileController.verify_wireguard_endpoint(profile, ignore=ignore)

            port_number = ConnectionService.get_random_available_port_number()
            ConnectionController.establish_wireguard_session_connection(profile, session_directory, port_number)
            session_state.network_port_numbers.wireguard.append(port_number)

        elif profile.connection.code in ('vless', 'hysteria2'):

            if not profile.has_operator_proxy_session():
                raise MissingSubscriptionError()

            operator_proxy_session = profile.get_operator_proxy_session()
            port_number = ConnectionService.get_random_available_port_number()

            if profile.connection.code == 'vless':
                from core.controllers.encrypted_proxy.VlessController import VlessController
                VlessController.enable(operator_proxy_session, port_number)
                session_state.network_port_numbers.vless.append(port_number)

            elif profile.connection.code == 'hysteria2':
                from core.controllers.encrypted_proxy.HysteriaController import HysteriaController
                HysteriaController.enable(operator_proxy_session, port_number)
                session_state.network_port_numbers.hysteria2.append(port_number)

        if profile.connection.masked:

            while proxy_port_number is None or proxy_port_number == port_number:
                proxy_port_number = ConnectionService.get_random_available_port_number()

            ConnectionController.establish_proxy_session_connection(profile, session_directory, port_number, proxy_port_number)
            session_state.network_port_numbers.proxy.append(proxy_port_number)

        if not profile.connection.is_unprotected():
            ConnectionController.await_connection(proxy_port_number or port_number, connection_observer=connection_observer)

        SessionStateController.update_or_create(session_state)

        return proxy_port_number or port_number

    @staticmethod
    def _signal_previous_process() -> None:

        system_state = SystemStateController.get()
        if system_state is None:
            return
        pid = system_state.pid
        if pid is None or pid == os.getpid():
            return
        try:
            os.kill(pid, signal.SIGUSR1)
            time.sleep(0.5)
        except (ProcessLookupError, PermissionError):
            pass

    @staticmethod
    def establish_system_connection(profile: SystemProfile, ignore: tuple[type[Exception]] = (), connection_observer: Optional[ConnectionObserver] = None):
        ConnectionController._signal_previous_process()

        system_state = SystemStateController.get()
        if system_state is not None:
            try:
                ConnectionController.terminate_system_connection(connection_observer=connection_observer)
            except ConnectionTerminationError:
                pass

        if profile.connection.needs_operator_proxy():
            ok = ConnectionController.establish_encrypted_proxy_connection(
                profile, socks5_port=1080, observer=connection_observer
            )
            if not ok:
                raise ConnectionError('The connection could not be established.')
            token = SystemStateController.create(profile.id)
            if connection_observer is not None:
                connection_observer.notify('connected_token', {'session_token': token})
            return

        if ConfigurationController.get_endpoint_verification_enabled():
            ProfileController.verify_wireguard_endpoint(profile, ignore=ignore)

        try:
            ConnectionController.__establish_system_connection(profile, connection_observer)

        except ConnectionError:
            try:
                ConnectionController.terminate_system_connection()
            except ConnectionTerminationError:
                pass
            raise ConnectionError('The connection could not be established.')

        except CalledProcessError:
            try:
                ConnectionController.terminate_system_connection()
            except ConnectionTerminationError:
                pass
            try:
                ConnectionController.__establish_system_connection(profile, connection_observer)
            except (ConnectionError, CalledProcessError):
                try:
                    ConnectionController.terminate_system_connection()
                except ConnectionTerminationError:
                    pass
                raise ConnectionError('The connection could not be established.')

        ConnectionController.terminate_tor_connection()
        time.sleep(1.0)

    @staticmethod
    def establish_encrypted_proxy_connection(profile, socks5_port: int = 1080, observer=None) -> bool:
        """Shared encrypted proxy connection logic for CLI and GUI."""
        import socket
        operator_proxy_session = profile.get_operator_proxy_session()
        protocol = profile.connection.code

        if protocol == 'vless':
            from core.controllers.encrypted_proxy.VlessController import VlessController
            from core.services.encrypted_proxy.vless_service import parse_vless_link
            server_ip = operator_proxy_session.server_ip
            if server_ip is None:
                vless = parse_vless_link(operator_proxy_session.links[0])
                server_ip = socket.gethostbyname(vless['host'])
            return VlessController(socks5_port).enable(
                operator_proxy_session.links[0],
                operator_proxy_session.username,
                server_ip,
                observer
            )
        elif protocol == 'hysteria2':
            from core.controllers.encrypted_proxy.HysteriaController import HysteriaController
            server_ip = operator_proxy_session.server_ip
            if server_ip is None:
                server_ip = socket.gethostbyname(operator_proxy_session.operator_hysteria2_host)
            return HysteriaController(socks5_port).enable(
                operator_proxy_session.username,
                operator_proxy_session.password,
                operator_proxy_session.operator_hysteria2_host,
                server_ip,
                observer
            )
        return False

    @staticmethod
    def establish_tor_connection(connection_observer: Optional[ConnectionObserver] = None):

        tor_module = TorModule(Constants.HV_TOR_STATE_HOME)
        tor_module.start_service(connection_observer)

        for session_state in SessionStateController.all():

            for port_number in session_state.network_port_numbers.tor:
                tor_module.create_session(port_number)

    @staticmethod
    def terminate_tor_connection():

        tor_module = TorModule(Constants.HV_TOR_STATE_HOME)
        tor_module.stop_service()

    @staticmethod
    def establish_tor_session_connection(port_number: int, connection_observer: Optional[ConnectionObserver] = None):
        try:
            tor_module = TorModule(Constants.HV_TOR_STATE_HOME)
            tor_module.create_session(port_number, connection_observer)
        except TorServiceInitializationError as e:
            logger.error(f"TorServiceInitializationError. Tor Can't Start: {e}")
        except Exception as e:
            logger.error(f"Tor Can't Start: {e}")

    @staticmethod
    def terminate_tor_session_connection(port_number: int):

        tor_module = TorModule(Constants.HV_TOR_STATE_HOME)
        tor_module.destroy_session(port_number)

    @staticmethod
    def establish_wireguard_session_connection(profile: SessionProfile, session_directory: str, port_number: int):

        if not profile.has_wireguard_configuration():
            raise FileNotFoundError('No valid WireGuard configuration file detected.')

        wireguard_session_directory = f'{session_directory}/wireguard'
        Path(wireguard_session_directory).mkdir(exist_ok=True, mode=0o700)

        wireproxy_configuration_file_path = f'{wireguard_session_directory}/wireproxy.conf'
        Path(wireproxy_configuration_file_path).touch(exist_ok=True, mode=0o600)

        with open(wireproxy_configuration_file_path, 'w') as wireproxy_configuration_file:
            wireproxy_configuration_file.write(f'WGConfig = {profile.get_wireguard_configuration_path()}\n\n[Socks5]\nBindAddress = 127.0.0.1:{str(port_number)}\n')

        return subprocess.Popen((f'{Constants.HV_RUNTIME_DATA_HOME}/wireproxy/wireproxy', '-c', wireproxy_configuration_file_path), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)

    @staticmethod
    def establish_proxy_session_connection(profile: SessionProfile, session_directory: str, port_number: int, proxy_port_number: int):

        if shutil.which('proxychains4') is None:
            raise CommandNotFoundError('proxychains4')

        if shutil.which('microsocks') is None:
            raise CommandNotFoundError('microsocks')

        if profile.has_proxy_configuration():
            proxy_configuration = profile.get_proxy_configuration()
        else:
            raise FileNotFoundError('No valid proxy configuration file detected.')

        proxy_session_directory = f'{session_directory}/proxy'
        Path(proxy_session_directory).mkdir(parents=True, exist_ok=True, mode=0o700)

        proxychains_proxy_list = ''

        if port_number is not None:
            proxychains_proxy_list = f'socks5 127.0.0.1 {port_number}\n'

        proxychains_proxy_list = proxychains_proxy_list + f'socks5 {proxy_configuration.ip_address} {proxy_configuration.port_number} {proxy_configuration.username} {proxy_configuration.password}'
        proxychains_template_file_path = f'{Constants.HV_RUNTIME_DATA_HOME}/proxychains.ptpl'

        with open(proxychains_template_file_path, 'r') as proxychains_template_file:

            proxychains_configuration_file_path = f'{proxy_session_directory}/proxychains.conf'
            Path(proxychains_configuration_file_path).touch(exist_ok=True, mode=0o600)

            proxychains_configuration_file_contents = proxychains_template_file.read().format(proxy_list=proxychains_proxy_list)

            with open(proxychains_configuration_file_path, 'w') as proxychains_configuration_file:
                proxychains_configuration_file.write(proxychains_configuration_file_contents)

        return subprocess.Popen(('proxychains4', '-f', proxychains_configuration_file_path, 'microsocks', '-p', str(proxy_port_number)), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)

    @staticmethod
    def terminate_system_connection(connection_observer: Optional[ConnectionObserver] = None):
        system_state = SystemStateController.get()
        if system_state is not None:
            profile = ProfileController.get(system_state.profile_id)
            if profile is not None and profile.connection.needs_operator_proxy():
                protocol = profile.connection.get_protocol()
                if protocol == 'vless':
                    from core.controllers.encrypted_proxy.VlessController import VlessController
                    VlessController(1080).disable(connection_observer)
                elif protocol == 'hysteria2':
                    from core.controllers.encrypted_proxy.HysteriaController import HysteriaController
                    HysteriaController(1080).disable(connection_observer)
                SystemState.dissolve()
                return

        if shutil.which('nmcli') is None:
            raise CommandNotFoundError('nmcli')

        if SystemStateController.exists():
            process = subprocess.Popen(('nmcli', 'connection', 'delete', 'wg'), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
            completed_successfully = not bool(os.waitpid(process.pid, 0)[1] >> 8)

            if completed_successfully or not ConnectionController.system_uses_wireguard_interface():
                subprocess.run(('nmcli', 'connection', 'delete', 'hv-ipv6-sink'), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                ConnectionController.terminate_tor_connection()
                try:
                    from core.utils.encrypted_proxy import killswitch
                    killswitch.disarm()
                except Exception:
                    pass
                SystemState.dissolve()
                if connection_observer is not None:
                    connection_observer.notify('disconnected', {})
            else:
                raise ConnectionTerminationError('The connection could not be terminated.')

    @staticmethod
    def get_proxies(port_number: int):

        return dict(
            http=f'socks5h://127.0.0.1:{port_number}',
            https=f'socks5h://127.0.0.1:{port_number}'
        )

    @staticmethod
    def await_connection(port_number: Optional[int] = None, connection_observer: Optional[ConnectionObserver] = None):

        if port_number is None:
            ConnectionController.await_network_interface()

        for retry_count in range(Constants.MAX_CONNECTION_ATTEMPTS):

            if connection_observer is not None:

                connection_observer.notify('connecting', dict(
                    retry_interval=Constants.CONNECTION_RETRY_INTERVAL,
                    maximum_number_of_attempts=Constants.MAX_CONNECTION_ATTEMPTS,
                    attempt_count=retry_count + 1
                ))

            try:

                ConnectionController.__test_connection(port_number)
                return

            except ConnectionError:

                time.sleep(Constants.CONNECTION_RETRY_INTERVAL)
                retry_count += 1

        raise ConnectionError('The connection could not be established.')

    @staticmethod
    def await_network_interface():

        network_interface_is_activated = False

        retry_interval = .5
        maximum_number_of_attempts = 10
        attempt = 0

        while not network_interface_is_activated and attempt < maximum_number_of_attempts:

            time.sleep(retry_interval)

            network_interface_is_activated = ConnectionController.system_uses_wireguard_interface()
            attempt += 1

        if not network_interface_is_activated:
            raise ConnectionError('The network interface could not be activated.')

    @staticmethod
    def system_uses_wireguard_interface():

        if shutil.which('ip') is None:
            raise CommandNotFoundError('ip')

        process = subprocess.Popen(('ip', 'route', 'get', '192.0.2.1'), stdout=subprocess.PIPE)
        process_output = str(process.stdout.read())

        return bool(re.search('dev wg', str(process_output)))

    @staticmethod
    def __establish_system_connection(profile: SystemProfile, connection_observer: Optional[ConnectionObserver] = None):

        if shutil.which('dbus-send') is None:
            raise CommandNotFoundError('dbus-send')

        if shutil.which('nmcli') is None:
            raise CommandNotFoundError('nmcli')

        ConnectionController.terminate_system_connection()

        try:
            process_output = subprocess.check_output(('nmcli', 'connection', 'import', '--temporary', 'type', 'wireguard', 'file', profile.get_wireguard_configuration_path()), text=True)
        except CalledProcessError:
            raise ConnectionError('The connection could not be established.')

        try:

            connection_id = (m := re.search(r'(?<=\()([a-f0-9-]+?)(?=\))', process_output)) and m.group(1)
            ipv6_method = subprocess.check_output(('nmcli', '-g', 'ipv6.method', 'connection', 'show', connection_id), text=True).strip()

        except CalledProcessError:
            raise ConnectionError('The connection could not be established.')

        if ipv6_method in ('disabled', 'ignore'):

            try:
                subprocess.run(('dbus-send', '--system', '--print-reply', '--dest=org.freedesktop.NetworkManager', '/org/freedesktop/NetworkManager', 'org.freedesktop.DBus.Properties.Set', 'string:org.freedesktop.NetworkManager', 'string:ConnectivityCheckEnabled', 'variant:boolean:false'), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
            except CalledProcessError:
                raise ConnectionError('The connection could not be established.')

            try:
                subprocess.run(('nmcli', 'connection', 'add', 'type', 'dummy', 'save', 'no', 'con-name', 'hv-ipv6-sink', 'ifname', 'hvipv6sink0', 'ipv6.method', 'manual', 'ipv6.addresses', 'fd7a:fd4b:54e3:077c::/64', 'ipv6.gateway', 'fd7a:fd4b:54e3:077c::1', 'ipv6.dns', '::1', 'ipv6.route-metric', '72'), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
            except CalledProcessError:
                raise ConnectionError('The connection could not be established.')

        # Arm killswitch for WireGuard
        try:
            wg_server_ip = ConnectionController.__extract_wireguard_endpoint(profile)
            if wg_server_ip:
                from core.utils.encrypted_proxy import killswitch
                killswitch.arm(wg_server_ip, 'wg')
        except Exception:
            pass

        token = SystemStateController.create(profile.id)
        if connection_observer is not None:
            connection_observer.notify('connected_token', {'session_token': token})

        try:
            ConnectionController.await_connection(connection_observer=connection_observer)

        except ConnectionError:
            raise ConnectionError('The connection could not be established.')

    @staticmethod
    def __extract_wireguard_endpoint(profile):
        import re, socket
        try:
            config = profile.get_wireguard_configuration()
            if config is None:
                return None
            match = re.search(r'Endpoint\s*=\s*([^\s:]+):\d+', config)
            if not match:
                return None
            host = match.group(1)
            if re.match(r'^(\d{1,3}\.){3}\d{1,3}$', host):
                return host
            return socket.gethostbyname(host)
        except Exception:
            return None

    @staticmethod
    def __with_tor_connection(*args, task: Callable[..., Any], connection_observer: Optional[ConnectionObserver] = None, **kwargs):

        port_number = ConnectionService.get_random_available_port_number()
        ConnectionController.establish_tor_session_connection(port_number, connection_observer=connection_observer)

        ConnectionController.await_connection(port_number, connection_observer=connection_observer)
        task_output = task(*args, proxies=ConnectionController.get_proxies(port_number), **kwargs)

        ConnectionController.terminate_tor_session_connection(port_number)

        return task_output

    @staticmethod
    def __test_connection(port_number: Optional[int] = None, timeout: float = 4.0):

        request_urls = [Constants.PING_URL]
        proxies = None

        if os.environ.get('PING_URL') is None:

            request_urls.extend([
                'https://hc1.simplifiedprivacy.net',
                'https://hc2.simplifiedprivacy.org',
                'https://hc3.hydraveil.net'
            ])

            random.shuffle(request_urls)

        if port_number is not None:
            proxies = ConnectionController.get_proxies(port_number)

        for request_url in request_urls:

            command = [
                sys.executable, '-u', '-c', 'import requests, sys\n'
                'try:\n'
                f'    response = requests.get(\'{request_url}\', proxies={proxies}, timeout={timeout})\n'
                '    response.raise_for_status(); print(response.text)\n'
                'except requests.exceptions.RequestException:\n'
                '    sys.exit(1)'
            ]

            try:

                _response = subprocess.check_output(command, text=True, timeout=timeout)
                return None

            except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
                pass

        raise ConnectionError('The connection could not be established.')

    @staticmethod
    def __should_renegotiate(profile: Union[SessionProfile, SystemProfile]):

        if not profile.has_subscription():
            raise MissingSubscriptionError()

        if profile.connection.needs_wireguard_configuration() and profile.has_wireguard_configuration():

            if profile.subscription.has_been_activated():
                return True

        return False

    @staticmethod
    def verify_singbox_installation():
        import subprocess
        import os
        from core.Errors import SingboxNotInstalledException

        singbox_binary = '/usr/bin/sing-box'
        singbox_wrapper = '/usr/local/bin/hydraveil-singbox'

        if not os.path.exists(singbox_binary):
            raise SingboxNotInstalledException(f'sing-box binary not found at {singbox_binary}')

        if not os.path.exists(singbox_wrapper):
            raise SingboxNotInstalledException(f'sing-box wrapper not found at {singbox_wrapper}')

        try:
            subprocess.run([singbox_binary, 'version'], capture_output=True, timeout=5, check=True)
        except subprocess.CalledProcessError:
            raise SingboxNotInstalledException('sing-box binary check failed')
        except subprocess.TimeoutExpired:
            raise SingboxNotInstalledException('sing-box verification timeout')