from functools import cache
import json
from flask import (
    Flask,
    make_response,
    redirect,
    request,
    jsonify,
    render_template,
    send_from_directory,
    send_file,
)
import os
import json
import requests
from datetime import datetime, timedelta
import dotenv

dotenv.load_dotenv()

app = Flask(__name__)

# Proxy PAC configuration
PROXY_ADDRESS = os.getenv('PROXY_ADDRESS', '127.0.0.1:9590')
# URL for IANA TLD list
IANA_TLD_URL = "https://data.iana.org/TLD/tlds-alpha-by-domain.txt"
# Default TLDs to always skip (in addition to IANA list)
DEFAULT_SKIP_TLDS = [
    'localhost', 'test', 'invalid', 'example', 'local',
    'arpa', 'onion', 'internal', 'private', 'home'
]
# Additional TLDs to skip specified in environment
ADDITIONAL_SKIP_TLDS = os.getenv('ADDITIONAL_SKIP_TLDS', '').split(',')
if ADDITIONAL_SKIP_TLDS == ['']:
    ADDITIONAL_SKIP_TLDS = []

# IANA TLD list cache
tld_list_cache = {
    'tlds': [],
    'last_updated': None,
    'initialized': False
}

def get_iana_tlds():
    """Download and parse the IANA TLD list."""
    now = datetime.now()
    
    # Check if we need to refresh the cache (daily)
    if (tld_list_cache['last_updated'] is None or 
        now - tld_list_cache['last_updated'] > timedelta(days=1)):
        try:
            response = requests.get(IANA_TLD_URL)
            if response.status_code == 200:
                # Parse the TLD list (skip the header line, convert to lowercase)
                tlds = [line.strip().lower() for line in response.text.splitlines() 
                       if line.strip() and not line.startswith('#')]
                tld_list_cache['tlds'] = tlds
                tld_list_cache['last_updated'] = now
                tld_list_cache['initialized'] = True
                print(f"Downloaded {len(tlds)} TLDs from IANA")
            else:
                print(f"Failed to download IANA TLD list: {response.status_code}")
                # If we failed but have a previous cache, keep using it
                if not tld_list_cache['tlds']:
                    # Otherwise use an empty list
                    tld_list_cache['tlds'] = []
        except Exception as e:
            print(f"Error downloading IANA TLD list: {e}")
            # If exception occurs but we have a previous cache, keep using it
            if not tld_list_cache['tlds']:
                # Otherwise use an empty list
                tld_list_cache['tlds'] = []
    
    # Combine IANA TLDs with our default skip TLDs and additional skip TLDs
    all_tlds = list(set(tld_list_cache['tlds'] + DEFAULT_SKIP_TLDS + ADDITIONAL_SKIP_TLDS))
    return all_tlds

def generate_pac_script(proxy_addr, skip_tlds):
    """Generate a Proxy Auto-Configuration script."""
    skipped_tlds = "', '".join(skip_tlds)
    return f"""
function FindProxyForURL(url, host) {{
    var skipped = ['{skipped_tlds}'];

    // skip any TLD in the list 
    var tld = host;
    var lastDot = tld.lastIndexOf('.');
    if (lastDot != -1) {{
      tld = tld.substr(lastDot+1);
    }}
    tld = tld.toLowerCase();

    if (skipped.includes(tld)) {{
      return 'DIRECT';
    }}

    // skip IP addresses
    var isIpV4Addr = /^(\\d+.)(\\d+.)(\\d+.)(\\d+)$/;
    if (isIpV4Addr.test(host)) {{
       return "DIRECT";
    }}

    // loosely check if IPv6
    if (lastDot == -1 && host.split(':').length > 2) {{
      return "DIRECT";
    }}

    return "PROXY {proxy_addr}";
}}
"""

def find(name, path):
    for root, dirs, files in os.walk(path):
        if name in files:
            return os.path.join(root, name)

# Assets routes
@app.route("/assets/<path:path>")
def send_assets(path):
    if path.endswith(".json"):
        return send_from_directory(
            "templates/assets", path, mimetype="application/json"
        )

    if os.path.isfile("templates/assets/" + path):
        return send_from_directory("templates/assets", path)

    # Try looking in one of the directories
    filename: str = path.split("/")[-1]
    if (
        filename.endswith(".png")
        or filename.endswith(".jpg")
        or filename.endswith(".jpeg")
        or filename.endswith(".svg")
    ):
        if os.path.isfile("templates/assets/img/" + filename):
            return send_from_directory("templates/assets/img", filename)
        if os.path.isfile("templates/assets/img/favicon/" + filename):
            return send_from_directory("templates/assets/img/favicon", filename)

    return render_template("404.html"), 404


# region Special routes
@app.route("/favicon.png")
def faviconPNG():
    return send_from_directory("templates/assets/img", "favicon.png")


@app.route("/proxy.pac")
def proxy_pac():
    """Serve the Proxy Auto-Configuration file."""
    proxy_addr = PROXY_ADDRESS
    skip_tlds = get_iana_tlds()
    
    response = make_response(generate_pac_script(proxy_addr, skip_tlds))
    response.headers["Content-Type"] = "application/x-ns-proxy-autoconfig"
    return response

@app.route("/.well-known/<path:path>")
def wellknown(path):
    # Try to proxy to https://nathan.woodburn.au/.well-known/
    req = requests.get(f"https://nathan.woodburn.au/.well-known/{path}")
    return make_response(
        req.content, 200, {"Content-Type": req.headers["Content-Type"]}
    )


# endregion


# region Main routes
@app.route("/")
def index():
    return render_template("index.html")


@app.route("/<path:path>")
def catch_all(path: str):
    if os.path.isfile("templates/" + path):
        return render_template(path)

    # Try with .html
    if os.path.isfile("templates/" + path + ".html"):
        return render_template(path + ".html")

    if os.path.isfile("templates/" + path.strip("/") + ".html"):
        return render_template(path.strip("/") + ".html")

    # Try to find a file matching
    if path.count("/") < 1:
        # Try to find a file matching
        filename = find(path, "templates")
        if filename:
            return send_file(filename)

    return render_template("404.html"), 404


# endregion


# region Error Catching
# 404 catch all
@app.errorhandler(404)
def not_found(e):
    return render_template("404.html"), 404


# endregion

# Replace before_first_request with a startup initialization
# This runs once at server startup
with app.app_context():
    get_iana_tlds()

if __name__ == "__main__":
    # Initialize TLD list at startup if running directly
    if not tld_list_cache['initialized']:
        get_iana_tlds()
    app.run(debug=True, port=5000, host="0.0.0.0")