import asyncio
from ipaddress import IPv4Address
import mysql.connector
import redis
import json
import requests
import time
import calendar
import os
import logging
import traceback
from datetime import datetime, timedelta

try:
    import geoip2.database
    from geoip2.errors import AddressNotFoundError
    geoip2_available = True
except ImportError:
    geoip2_available = False


class RedisClientHelper:
    def __init__(self, host='localhost', port=6379, db=0, decode_responses=True):
        """Initialize the Redis client."""
        self.redis_client = redis.StrictRedis(host=host, port=port, db=db, decode_responses=decode_responses)

    def scan_all_keys(self, pattern='*', count=1000):
        """Scan all keys in Redis."""
        cursor = 0
        self.keys = []
        while True:
            cursor, new_keys = self.redis_client.scan(cursor=cursor, match=pattern, count=count)
            self.keys.extend(new_keys)
            print('keys getting from redis',self.keys)
            if cursor == 0:
                break
        return self.keys
    

    def get_all_data(self, keys):
        """Get all data from Redis."""
        all_data = {}
        for key in keys:
            key_type = self.redis_client.type(key)
            if key_type == 'hash':
                all_data[key] = self.redis_client.hgetall(key)
            elif key_type == 'string':
                all_data[key] = self.redis_client.get(key)
            elif key_type == 'list':
                all_data[key] = self.redis_client.lrange(key, 0, -1)
            elif key_type == 'set':
                all_data[key] = list(self.redis_client.smembers(key))  # Convert set to list for JSON serialization
            elif key_type == 'zset':
                all_data[key] = self.redis_client.zrange(key, 0, -1, withscores=True)
        return all_data

class LoggerManager:
    def __init__(self, log_directory, delete_log_days=10, log_name="VPN"):
        self.log_directory = log_directory
        self.delete_log_days = delete_log_days
        self.log_name = log_name
        self.logger = None
        self.setup_logging()

    def setup_logging(self):
        """Sets up the logging configuration and creates the log file."""
        # Ensure the log directory exists
        if not os.path.exists(self.log_directory):
            os.makedirs(self.log_directory)
        
        log_file_path = os.path.join(self.log_directory, f'{self.log_name}_{datetime.now().strftime("%Y-%m-%d")}.log')
        
        # Create a log formatter and handler
        log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        file_handler = logging.FileHandler(log_file_path)
        file_handler.setFormatter(log_formatter)

        # Set up the logger
        self.logger = logging.getLogger(self.log_name)
        self.logger.addHandler(file_handler)
        self.logger.setLevel(logging.INFO)

    def remove_old_logs(self):
        """Removes log files older than the specified number of days."""
        cutoff_date = datetime.now() - timedelta(days=self.delete_log_days)

        for filename in os.listdir(self.log_directory):
            if filename.endswith('.log'):
                file_path = os.path.join(self.log_directory, filename)
                file_mod_time = datetime.fromtimestamp(os.path.getmtime(file_path))
                
                # If the file is older than the cutoff date, remove it
                if file_mod_time < cutoff_date:
                    os.remove(file_path)
                    self.logger.info(f"Removed old log file: {filename}")

        self.logger.info("Log cleanup complete.")

    def get_logger(self):
        """Returns the logger object for use in other classes."""
        return self.logger


class OpenVPNClient:
    def __init__(self, servers, db_conn,data_limit_mb,max_concurrent_connections=100):
        """
        Initialize the OpenVPN client with a list of servers, geoip_version, and max concurrent connections.
        """
        self.connected_servers_str = ""
        self.servers = servers
        self.gi = geoip2.database.Reader('/var/www/html/1375-staysolid-ovpn-backend/GeoLite2-City.mmdb')
        self.geoip_version = 2
        self.max_concurrent_connections = max_concurrent_connections
        self.semaphore = asyncio.Semaphore(self.max_concurrent_connections)
        self.session = {}
        self.servers_data = []
        self.connected_servers = []
        self.db_conn = db_conn
        self.data_limit_bytes = data_limit_mb * 1024 * 1024

        self._working_directory = os.path.dirname(os.path.abspath(__file__))
        self._log_directory = f'{self._working_directory}/logs'

        self.logger_manager = LoggerManager(self._log_directory,10,"OVPN_Monitor")
        self.logger = self.logger_manager.get_logger()

        # Initialize the Redis client helper
        self.redis_helper = RedisClientHelper()

        # Getting data from redis of active server
        start_time = time.time()
        all_keys = self.redis_helper.scan_all_keys()

        self.all_data_retrieved = self.redis_helper.get_all_data(all_keys)
        print('redis data',self.all_data_retrieved)
        end_time = time.time()

        print(f"Time taken to scan {len(all_keys)} keys and retrieve data: {end_time - start_time:.4f} seconds")

    def disconnect_user(self, username, remote_ip, total_data_usage):
        try:
            disconnect_command = f"kill {username}\n"
            self.s.sendall(disconnect_command.encode('utf-8'))
            print(f"Sent disconnect command for user: {username}")
        except Exception as e:
            print(f"Error sending disconnect command: {e}")

    def log_disconnect(self, username, remote_ip, total_data_usage, disconnect_time):
        log_entry = f"{disconnect_time},{username},{remote_ip},{total_data_usage}\n"
        
        with open(self.log_file, 'a') as log_file:
            log_file.write(log_entry)
            print(f"Logged disconnect: {log_entry.strip()}")

    def n_sessions(self):
        # self.logger.info('----------------2')
        """Update session data into the database."""
        try:
            with self.db_conn.cursor(dictionary=True) as db_cursor:

                print(self.servers_data)
                print('size',len(self.servers_data))
                self.logger.info(f'Total servers : {len(self.servers_data)}')
                self.logger.info(self.servers_data)
                for session_data in self.servers_data:
                    
                    username = session_data['username']
                    remote_ip = str(session_data['remote_ip'])
                    incoming_bytes_recv = session_data['bytes_recv']
                    incoming_bytes_sent = session_data['bytes_sent']
                    connected_since = str(session_data['connected_since'])
                    total_data_usage = session_data['total_data_usage']
                    last_seen = str(session_data['last_seen'])
                    location = session_data['location']
                    region = session_data['region']
                    latitude = session_data['latitude']
                    city = session_data['city']
                    country = session_data['country']
                    longitude = session_data['longitude']
                    local_ips = session_data['local_ip']
                    server_ip = session_data['host_name']
                    # print('username',username)
                    # print(f"SELECT o.username, o.current_bytes_recv, o.current_bytes_sent, o.current_total_bytes, sc.threshold_bytes FROM ovpn_server_logs AS o JOIN server_configurations AS sc ON sc.server_ip = o.server_ip WHERE o.username = {username}")
                    # db_cursor.execute("SELECT o.name,sc.threshold_bytes,o.user_device_id FROM ovpn_servers AS o JOIN server_configurations AS sc ON o.server_config_id = sc.id WHERE o.name = %s",(username,))

                    db_cursor.execute("SELECT o.username,os.user_device_id FROM ovpn_server_logs AS o JOIN ovpn_servers AS os ON o.username = os.name WHERE o.username = %s", (username,))
                    result = db_cursor.fetchone()
                    print('result',result)
                    # self.logger.info(f'result : {result}')
                    
                    if result:
                        print('update data into database')
                        file_name, user_device_id = result['username'],result['user_device_id']
                        # self.logger.info('insert2')
                        self.connected_servers.append(username)
                        self.connected_servers_str = "'" + "', '".join(self.connected_servers) + "'"

                        new_total_bytes = incoming_bytes_recv + incoming_bytes_sent

                        print('-----1------',new_total_bytes)
                        print(incoming_bytes_recv)
                        print(incoming_bytes_sent)
                        dt = time.gmtime()
                        ts = calendar.timegm(dt)

                        update_query = """ UPDATE ovpn_server_logs SET current_bytes_recv = %s, current_bytes_sent = %s, current_total_bytes = %s, is_latest = %s, connected_since = %s, last_seen = %s, region = %s, latitude = %s, city = %s, country = %s, longitude = %s, remote_ip = %s, local_ip = %s, server_ip = %s, updated_at = %s, user_device_id = %s WHERE username = %s"""
                        db_cursor.execute(update_query, (incoming_bytes_recv, incoming_bytes_sent, new_total_bytes, 1, connected_since, last_seen, region, latitude, city, country, longitude, remote_ip, local_ips, server_ip,ts, user_device_id, username))
                        print('update_query : ',update_query)
                        self.db_conn.commit()
                            

                    else:
                        try:
                            db_cursor.execute("SELECT os.user_device_id FROM ovpn_servers AS os WHERE os.name = %s", (username,))
                            d_id = db_cursor.fetchone()
                            print('New server coming',d_id)
                            # self.logger.info('insert')
                            self.connected_servers.append(username)
                            self.connected_servers_str = "'" + "', '".join(self.connected_servers) + "'"
                            dt = time.gmtime()
                            ts = calendar.timegm(dt)
                            insert_query = """
                            INSERT INTO ovpn_server_logs (username, current_bytes_recv, current_bytes_sent, current_total_bytes,is_latest ,connected_since , last_seen, region, latitude, city, country, longitude, remote_ip, local_ip, server_ip, created_at, updated_at, user_device_id)
                            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
                            %s, %s, %s, %s)
                            """
                            db_cursor.execute(insert_query, (username, incoming_bytes_recv, incoming_bytes_sent, incoming_bytes_recv + incoming_bytes_sent, 1, connected_since, last_seen, region, latitude, city, country, longitude, remote_ip, local_ips, server_ip, ts, ts, d_id['user_device_id']))
                            self.db_conn.commit()

                            print(f"Inserted new data usage record for {username}: recv = {incoming_bytes_recv}, sent = {incoming_bytes_sent}")
                        except Exception as e:
                            self.logger.error(e)
                            print(e)

                self.connected_servers_str_redis = "'" + "', '".join(self.redis_helper.keys) + "'"
                print('keys after add into string',self.connected_servers_str_redis)
                if not self.connected_servers_str:
                    self.connected_servers_str = "''"
                print('These are connected servers',self.connected_servers_str)
                print("server into redis",self.connected_servers_str_redis)
                db_cursor.execute(f"""SELECT username, current_bytes_recv, current_bytes_sent, current_total_bytes, total_bytes_recv, total_bytes_sent, total_bytes FROM ovpn_server_logs WHERE username NOT IN ({self.connected_servers_str}) AND username NOT IN ({self.connected_servers_str_redis}) AND is_latest = 1""")
                connected_servers_data = db_cursor.fetchall()
                print('Server for disconnect',connected_servers_data)

                if connected_servers_data:
                    update_data = []

                    for row in connected_servers_data:
                        username = row['username']
                        
                        redis_exists = self.redis_helper.redis_client.exists(username)

                        if redis_exists:
                            redis_data = self.all_data_retrieved.get(username, {})
                            
                            total_bytes_recv = row['total_bytes_recv'] + row['current_bytes_recv'] + int(redis_data.get('current_bytes_recv', 0))
                            total_bytes_sent = row['total_bytes_sent'] + row['current_bytes_sent'] + int(redis_data.get('current_bytes_sent', 0))
                            total_bytes = row['total_bytes'] + row['current_total_bytes'] + int(redis_data.get('current_total_bytes', 0))

                            update_data.append((total_bytes_recv, total_bytes_sent, total_bytes, username))

                            self.redis_helper.redis_client.delete(username)
                        else:
                            total_bytes_recv = row['total_bytes_recv'] + row['current_bytes_recv']
                            total_bytes_sent = row['total_bytes_sent'] + row['current_bytes_sent']
                            total_bytes = row['total_bytes'] + row['current_total_bytes']

                            # Add the updated data to the list
                            update_data.append((total_bytes_recv, total_bytes_sent, total_bytes, username))

                    if update_data:
                        update_query = """UPDATE ovpn_server_logs
                                        SET total_bytes_recv = %s, 
                                            total_bytes_sent = %s, 
                                            total_bytes = %s, 
                                            current_bytes_recv = 0, 
                                            current_bytes_sent = 0, 
                                            current_total_bytes = 0, 
                                            is_latest = 0 
                                        WHERE username = %s"""
                        db_cursor.executemany(update_query, update_data)
                        self.db_conn.commit()

        except mysql.connector.Error as err:
            self.logger.error(err)
            print(f"Error: {err}")
            self.db_conn.rollback()

    def get_geoip_data(self,session, geoip_version):
        try:
            if geoip_version == 1:
                # Legacy GeoIP lookup
                gir = self.gi.record_by_addr(str(session['remote_ip']))
                if gir is not None:
                    session['location'] = gir.get('country_code', '')
                    session['region'] = gir.get('region', '')
                    session['city'] = gir.get('city', '')
                    session['country'] = gir.get('country_name', '')
                    session['longitude'] = gir.get('longitude', '')
                    session['latitude'] = gir.get('latitude', '')
            
            elif geoip_version == 2:

                # GeoIP2 lookup
                gir = self.gi.city(str(session['remote_ip']))
                session['location'] = gir.country.iso_code if gir.country else ''
                session['region'] = gir.subdivisions.most_specific.iso_code if gir.subdivisions.most_specific else ''
                session['city'] = gir.city.name if gir.city else ''
                session['country'] = gir.country.name if gir.country else ''
                session['longitude'] = gir.location.longitude if gir.location else ''
                session['latitude'] = gir.location.latitude if gir.location else ''
        except geoip2.errors.AddressNotFoundError:
            pass
        except SystemError:
            pass


    def parse_status(self, data, server_ip, server_port, gi):
        """Parse the raw data and convert it to a structured dictionary."""
        if not data or data.strip() == '':
            print("No valid data to parse.")
            return {}

        lines = data.splitlines()

        parsing_clients = False
        parsing_routing_table = False

        for line in lines:
            # Start parsing the CLIENT_LIST section
            if line.startswith('Common Name'):
                parsing_clients = True
                continue
            elif line.startswith('ROUTING TABLE'):
                parsing_clients = False
                parsing_routing_table = True
                continue
            elif line.startswith('GLOBAL STATS'):
                break

            if parsing_clients and line:
                columns = line.split(',')
                username = columns[0].strip()  
                remote_ip = columns[1].split(':')[0].strip()  
                bytes_received = int(columns[2].strip()) 
                bytes_sent = int(columns[3].strip())  
                connected_since = datetime.strptime(columns[4].strip(), '%Y-%m-%d %H:%M:%S')

                total_data_usage = bytes_received + bytes_sent

                self.session = {
                    'host_name' : server_ip,
                    'session_key': f"{username}_{remote_ip}",
                    'username': username,
                    'remote_ip': IPv4Address(remote_ip),
                    'bytes_recv': bytes_received,
                    'bytes_sent': bytes_sent,
                    'connected_since': connected_since,
                    'total_data_usage': total_data_usage,
                    'last_seen': '',
                    'location': '',
                    'region': '',
                    'city': '',
                    'country': '',
                    'longitude': '',
                    'latitude': '',
                    'local_ip': ''
                }


                self.get_geoip_data(self.session, self.geoip_version)
                self.servers_data.append(self.session)

                session_key = f"{username}_{remote_ip}"

            elif parsing_routing_table and line:
                if line.startswith('Virtual Address'):
                    continue
                columns = line.split(',')
                if len(columns) >= 4:
                    local_ip = columns[0].strip()  
                    username = columns[1].strip()  
                    remote_ip = columns[2].split(':')[0].strip()  
                    last_seen = columns[3].strip()

                    session_key = f"{username}_{remote_ip}"
                    if session_key in self.session['session_key']:
                        self.session['local_ip'] = local_ip
                        self.session['last_seen'] = last_seen


    async def connect_to_openvpn(self, server_ip, server_port):
        """
        Connect to the OpenVPN management interface, send a command, and parse the response.
        """
        async with self.semaphore:
            try:
                reader, writer = await asyncio.open_connection(server_ip, server_port)
                print(f"Connected to OpenVPN management interface at {server_ip}:{server_port}")

                command = "status\n"
                writer.write(command.encode())
                await writer.drain()

                response = ""
                while True:
                    chunk = await reader.read(4096)
                    if not chunk:
                        break
                    response += chunk.decode('utf-8')
                    if "END" in chunk.decode('utf-8'):
                        break

                if not response:
                    print(f"No data received from {server_ip}:{server_port}")
                    return
                print('response =======',response)
                self.parse_status(response, server_ip, server_port, gi=None)
                # print(f"Parsed data from {server_ip}:{server_port}:\n{parsed_data}")
                # self.logger.info('--------------1')
                writer.close()
                await writer.wait_closed()

            except Exception as e:
                excp = traceback.format_exc()
                print("exp------", excp)
                self.logger.error(excp)
                print(f"Failed to connect to {server_ip}:{server_port} - Error: {e}")

    async def run(self):
        """
        Run the tasks for all servers asynchronously with controlled concurrency.
        """
        tasks = [self.connect_to_openvpn(server_ip, server_port) for server_ip, server_port in self.servers]
        await asyncio.gather(*tasks)
        self.n_sessions()


# Entry point
if __name__ == '__main__':

    db_conn = mysql.connector.connect(
        host="ovpn-centrealised-server.cl66yeq0k0b7.us-east-2.rds.amazonaws.com",
        user="xeg3ua7ooqu1",
        password="ahv3eigee2zo",
        database="xeg3ua7ooqu12",
        # port=5306
    )
    servers = [
        ('172.245.234.181', 5555),
        # ('13.200.64.163', 5555),
    ]
    
    openvpn_client = OpenVPNClient(servers,max_concurrent_connections=100, data_limit_mb=100000, db_conn=db_conn)
    asyncio.run(openvpn_client.run())
