import os
import sys
import yaml
import requests
import subprocess
from datetime import datetime
from loguru import logger
from dotenv import load_dotenv
from pathlib import Path

# Constants
SCRIPT_DIR = Path(__file__).resolve().parent
CADDYFILE_PATH = "/etc/caddy/Caddyfile"
CF_DOMAINS_FILE = SCRIPT_DIR / 'cf_domains.yaml'
ENV_FILE = SCRIPT_DIR / '.env'
LOG_FILE = SCRIPT_DIR / 'cf_script.log'

# Load environment variables
load_dotenv(ENV_FILE)

# Configure logger
logger.remove()  # Remove default handler
logger.add(sys.stderr, format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> <level>{level}</level> {message}", level="INFO", colorize=True)
logger.add(LOG_FILE, format="{time:YYYY-MM-DD HH:mm:ss} {level} {message}", rotation="10 MB", retention="1 week", level="DEBUG")

def get_current_ip():
    return requests.get('https://am.i.mullvad.net/ip').text.strip()

def update_env_ip(ip):
    env_vars = {}
    if ENV_FILE.exists():
        with open(ENV_FILE, 'r') as f:
            for line in f:
                if '=' in line:
                    key, value = line.strip().split('=', 1)
                    env_vars[key] = value

    env_vars['CURRENT_IP'] = ip

    with open(ENV_FILE, 'w') as f:
        for key, value in env_vars.items():
            f.write(f"{key}={value}\n")

def sort_yaml():
    with open(CF_DOMAINS_FILE, 'r') as file:
        data = yaml.safe_load(file)

    sorted_data = dict(sorted(data.items()))

    for domain, subdomains in sorted_data.items():
        if isinstance(subdomains, dict):
            sorted_data[domain] = dict(sorted(subdomains.items()))

    with open(CF_DOMAINS_FILE, 'w') as file:
        yaml.dump(sorted_data, file, sort_keys=False)

def usage():
    logger.error("Incorrect usage")
    print("Usage: cf <full-domain> [--ip <ip address>] --port <port>")
    print("       cf ddns [--force]")
    print("       cf all [--force]")
    sys.exit(1)

def update_caddyfile(full_domain, caddy_ip, port):
    with open(CADDYFILE_PATH, 'a') as file:
        file.write(f"""
{full_domain} {{
    reverse_proxy {caddy_ip}:{port}
    tls {{
        dns cloudflare {{env.CLOUDFLARE_API_TOKEN}}
    }}
}}
""")
    logger.info(f"Configuration appended to {CADDYFILE_PATH} with correct formatting.")

def update_dns_record(zone_id, dns_id, name, ip):
    url = f'https://api.cloudflare.com/client/v4/zones/{zone_id}/dns_records/{dns_id}'
    headers = {
        'Authorization': f'Bearer {os.getenv("CLOUDFLARE_API_TOKEN")}',
        'Content-Type': 'application/json'
    }
    data = {
        'type': 'A',
        'name': name,
        'content': ip,
        'ttl': 120,
        'proxied': True
    }
    response = requests.put(url, headers=headers, json=data)
    return response.json()

def ddns(force=False):
    current_ip = get_current_ip()
    last_ip = os.getenv('CURRENT_IP')

    if current_ip != last_ip or force:
        update_env_ip(current_ip)

        with open(CF_DOMAINS_FILE, 'r') as f:
            domains = yaml.safe_load(f)

        for domain, data in domains.items():
            zone_id = data['_id']
            for subdomain, dns_id in data.items():
                if subdomain == '_id':
                    continue
                full_domain = f"{subdomain}.{domain}" if subdomain != '@' else domain
                logger.info(f"Updating {full_domain}")
                result = update_dns_record(zone_id, dns_id, full_domain, current_ip)
                if result.get('success'):
                    logger.info(f"Successfully updated {full_domain}")
                else:
                    logger.error(f"Failed to update {full_domain}: {result.get('errors', 'Unknown error')}")

        logger.info(f"DDNS update completed at {datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')} UTC")
    else:
        logger.info("IP address hasn't changed. No updates needed.")

def get_existing_record(zone_id, name):
    url = f"https://api.cloudflare.com/client/v4/zones/{zone_id}/dns_records?type=A&name={name}"
    headers = {
        "Authorization": f"Bearer {os.getenv('CLOUDFLARE_API_TOKEN')}",
        "Content-Type": "application/json"
    }
    response = requests.get(url, headers=headers)
    result = response.json()
    if result.get('success') and result['result']:
        return result['result'][0]
    return None

def update_or_create_record(zone_id, subdomain, domain, cloudflare_ip):
    # For root domain, we need to use the full domain name instead of "@"
    full_name = domain if subdomain == "@" else f"{subdomain}.{domain}"

    existing_record = get_existing_record(zone_id, full_name)

    headers = {
        "Authorization": f"Bearer {os.getenv('CLOUDFLARE_API_TOKEN')}",
        "Content-Type": "application/json"
    }
    data = {
        "type": "A",
        "name": full_name,
        "content": cloudflare_ip,
        "ttl": 120,
        "proxied": True
    }

    if existing_record:
        # Check if the existing record has the same content
        if existing_record['content'] == cloudflare_ip:
            logger.info(f"Record for {full_name} already exists with the correct IP. No update needed.")
            return {"success": True, "result": existing_record}

        url = f"https://api.cloudflare.com/client/v4/zones/{zone_id}/dns_records/{existing_record['id']}"
        response = requests.put(url, headers=headers, json=data)
    else:
        url = f"https://api.cloudflare.com/client/v4/zones/{zone_id}/dns_records"
        response = requests.post(url, headers=headers, json=data)

    result = response.json()

    # Handle the specific error code
    if not result.get('success') and result.get('errors'):
        for error in result['errors']:
            if error.get('code') == 81058:  # "A record with the same settings already exists"
                logger.info(f"Record for {full_name} already exists with the correct settings. No update needed.")
                return {"success": True, "result": existing_record or data}

    return result

def main():
    if len(sys.argv) < 2:
        usage()

    if sys.argv[1] in ['ddns', 'all']:
        force = '--force' in sys.argv
        ddns(force)
        if sys.argv[1] == 'ddns':
            return

    if os.geteuid() != 0 and sys.argv[1] not in ['ddns', 'all']:
        logger.error("This script must be run as root for Caddyfile modifications. Try using 'sudo'.")
        sys.exit(1)

    full_domain = sys.argv[1]
    caddy_ip = "localhost"
    port = None

    i = 2
    while i < len(sys.argv):
        if sys.argv[i] in ['--ip', '-i']:
            caddy_ip = sys.argv[i+1]
            i += 2
        elif sys.argv[i] in ['--port', '-p']:
            port = sys.argv[i+1]
            i += 2
        else:
            usage()

    if not port:
        usage()

    parts = full_domain.split('.')
    if len(parts) == 2:
        domain = full_domain
        subdomain = "@"
    else:
        subdomain = parts[0]
        domain = '.'.join(parts[1:])

    with open(CF_DOMAINS_FILE, 'r') as file:
        cf_domains = yaml.safe_load(file)

    if domain not in cf_domains:
        logger.error(f"Domain {domain} not found in cf_domains.yaml")
        sys.exit(1)

    zone_id = cf_domains[domain].get('_id')
    if not zone_id:
        logger.error(f"Zone ID for {domain} could not be found.")
        sys.exit(1)

    cloudflare_ip = os.getenv('CURRENT_IP')

    result = update_or_create_record(zone_id, subdomain, domain, cloudflare_ip)

    if result.get('success'):
        record_id = result['result'].get('id')
        if record_id:
            if subdomain == "@":
                cf_domains[domain]["@"] = record_id
            else:
                cf_domains[domain][subdomain] = record_id
            with open(CF_DOMAINS_FILE, 'w') as file:
                yaml.dump(cf_domains, file)
            logger.info(f"A record for {full_domain} created/updated and cf_domains.yaml updated successfully.")
        else:
            logger.info(f"A record for {full_domain} already exists with correct settings. No update needed.")
        sort_yaml()
        logger.info("YAML file sorted after update.")
        update_caddyfile(full_domain, caddy_ip, port)
    else:
        error_message = result.get('errors', [{'message': 'Unknown error', 'code': 'N/A'}])[0]
        logger.error(f"Failed to create/update A record for {full_domain}. Error: {error_message.get('message')} (Code: {error_message.get('code')})")

    sort_yaml()
    logger.info("YAML file sorted at end.")

    logger.info("Restarting caddy!")
    subprocess.run(['sudo', 'systemctl', 'restart', 'caddy'], check=True)

if __name__ == "__main__":
    main()