import requests
import json
import argparse
import mysql.connector
import re
import os
from dotenv import load_dotenv

load_dotenv()

class NginxProxyManager:
    def __init__(self, host, email, password):
        self.host = host
        self.email = email
        self.password = password
        self.token = None

    def login(self):
        url = f"{self.host}/api/tokens"
        payload = {
            "identity": self.email,
            "secret": self.password
        }
        headers = {
            "Content-Type": "application/json"
        }
        response = requests.post(url, headers=headers, data=json.dumps(payload))
        if response.status_code == 200:
            self.token = response.json()["token"]
            print("Login successful.")
        else:
            print(f"Failed to login: {response.text}")
            exit(1)

    def get_proxy_hosts(self):
        if not self.token:
            self.login()
        
        url = f"{self.host}/api/nginx/proxy-hosts"
        headers = {
            "Authorization": f"Bearer {self.token}"
        }
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            return response.json()
        else:
            print(f"Failed to get proxy hosts: {response.text}")
            return []

    def get_proxy_host(self, host_id):
        if not self.token:
            self.login()

        url = f"{self.host}/api/nginx/proxy-hosts/{host_id}"
        headers = {
            "Authorization": f"Bearer {self.token}"
        }
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            return response.json()
        else:
            print(f"Failed to get proxy host {host_id}: {response.text}")
            return None

    def update_proxy_host_config(self, host_id, config):
        if not self.token:
            self.login()

        url = f"{self.host}/api/nginx/proxy-hosts/{host_id}"
        payload = {
            "advanced_config": config
        }
        headers = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json"
        }
        response = requests.put(url, headers=headers, data=json.dumps(payload))
        if response.status_code == 200:
            print(f"Successfully updated proxy host {host_id}")
        else:
            print(f"Failed to update proxy host {host_id}: {response.text}")

def get_streams_from_db(db_host, db_user, db_pass, db_name, db_port):
    try:
        conn = mysql.connector.connect(
            host=db_host,
            user=db_user,
            password=db_pass,
            database=db_name,
            port=db_port
        )
        cursor = conn.cursor(dictionary=True)
        cursor.execute("SELECT streamName, streamURL FROM streams")
        streams = cursor.fetchall()
        cursor.close()
        conn.close()
        return streams
    except mysql.connector.Error as err:
        print(f"Error connecting to database: {err}")
        return []

def update_config_with_streams(config, streams):
    for stream in streams:
        stream_name = stream['streamName']
        stream_url = stream['streamURL']
        # Use a more specific regex to avoid replacing parts of other URLs
        pattern = re.compile(f'(location ~ \^/{re.escape(stream_name)}\(\.\*\)\$ {{\s*return 302 )([^;]+)(;\\s*}})')
        config = pattern.sub(f'\\1{stream_url}/$1$is_args$args\\3', config)
    return config

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Modify Nginx Proxy Manager custom configuration.")
    parser.add_argument("--list-hosts", action="store_true", help="List all proxy hosts")
    parser.add_argument("--host-id", type=int, help="The ID of the proxy host to modify")
    parser.add_argument("--config-file", type=str, help="Path to the file containing the new advanced configuration")
    parser.add_argument("--download-config", type=str, help="Path to save the current advanced configuration")
    parser.add_argument("--update-from-db", action="store_true", help="Update the configuration from the database")

    args = parser.parse_args()

    npm_host = os.getenv("NPM_HOST")
    npm_email = os.getenv("NPM_EMAIL")
    npm_password = os.getenv("NPM_PASSWORD")
    db_host = os.getenv("DBHOST")
    db_user = os.getenv("DBUSER")
    db_pass = os.getenv("DBPASS")
    db_name = os.getenv("DATABASE")
    db_port = os.getenv("DBPORT")

    npm = NginxProxyManager(npm_host, npm_email, npm_password)
    npm.login()

    if args.list_hosts:
        hosts = npm.get_proxy_hosts()
        for host in hosts:
            print(f"ID: {host['id']}, Domains: {', '.join(host['domain_names'])}")
    
    if args.host_id and args.download_config:
        host = npm.get_proxy_host(args.host_id)
        if host:
            with open(args.download_config, 'w') as f:
                f.write(host.get('advanced_config', ''))
            print(f"Configuration for host {args.host_id} downloaded to {args.download_config}")

    if args.host_id and args.config_file:
        with open(args.config_file, 'r') as f:
            config = f.read()
        npm.update_proxy_host_config(args.host_id, config)

    if args.host_id and args.update_from_db:
        host = npm.get_proxy_host(args.host_id)
        if host:
            current_config = host.get('advanced_config', '')
            streams = get_streams_from_db(db_host, db_user, db_pass, db_name, db_port)
            if streams:
                new_config = update_config_with_streams(current_config, streams)
                npm.update_proxy_host_config(args.host_id, new_config)