import os

HAPROXY_CFG = '/etc/haproxy/haproxy.cfg'

def sanitize_name(name):
    """Convert hostname/name to valid ACL name"""
    return name.replace('.', '_').replace('-', '_').replace('/', '_').replace(':', '_')

def frontend_exists_at_port(frontend_ip, frontend_port):
    """Check if frontend already exists at specific port"""
    if not os.path.exists(HAPROXY_CFG):
        return None
    
    try:
        with open(HAPROXY_CFG, 'r') as f:
            content = f.read()
            lines = content.split('\n')
            
            for i, line in enumerate(lines):
                if line.strip().startswith('frontend'):
                    for j in range(i+1, min(i+10, len(lines))):
                        if lines[j].strip().startswith('bind'):
                            bind_info = lines[j].strip().split(' ', 1)[1]
                            bind_part = bind_info.split(' ssl ')[0].strip()
                            if f"{frontend_ip}:{frontend_port}" in bind_part:
                                return line.strip().split(' ', 1)[1]  # Zwróć nazwę frontendu
                        elif lines[j].strip().startswith('frontend') or lines[j].strip().startswith('backend'):
                            break
    except Exception as e:
        print(f"[HAPROXY_CONFIG] Error: {e}", flush=True)
    
    return None

def add_acl_to_frontend(frontend_name, acl_name, hostname, backend_name):
    if not os.path.exists(HAPROXY_CFG):
        return False
    
    try:
        with open(HAPROXY_CFG, 'r') as f:
            lines = f.readlines()
        
        frontend_idx = -1
        for i, line in enumerate(lines):
            if 'frontend' in line and frontend_name in line:
                frontend_idx = i
                break
        
        if frontend_idx == -1:
            print(f"[HAPROXY_CONFIG] Frontend '{frontend_name}' not found", flush=True)
            return False
        
        for line in lines[frontend_idx:]:
            if acl_name in line and 'acl' in line:
                print(f"[HAPROXY_CONFIG] ACL '{acl_name}' already exists", flush=True)
                return True 
            if line.strip().startswith('backend'):
                break
        
        insert_idx = frontend_idx + 1
        for i in range(frontend_idx + 1, len(lines)):
            if lines[i].strip().startswith('backend') or lines[i].strip().startswith('frontend'):
                insert_idx = i
                break
            if 'use_backend' in lines[i] or 'default_backend' in lines[i]:
                insert_idx = i + 1
        
        # Wstaw ACL i use_backend
        acl_line = f"    acl {acl_name} hdr(host) -i {hostname}\n"
        use_backend_line = f"    use_backend {backend_name} if {acl_name}\n"
        
        lines.insert(insert_idx, use_backend_line)
        lines.insert(insert_idx, acl_line)
        
        with open(HAPROXY_CFG, 'w') as f:
            f.writelines(lines)
        
        print(f"[HAPROXY_CONFIG] ACL '{acl_name}' added to frontend '{frontend_name}'", flush=True)
        return True
    except Exception as e:
        print(f"[HAPROXY_CONFIG] Error adding ACL: {e}", flush=True)
        return False

def is_backend_exist(backend_name):
    if not os.path.exists(HAPROXY_CFG):
        return False
    
    try:
        with open(HAPROXY_CFG, 'r') as haproxy_cfg:
            for line in haproxy_cfg:
                line = line.strip()
                if line.startswith('backend') and not line.startswith('#'):
                    parts = line.split()
                    if len(parts) >= 2 and parts[1] == backend_name:
                        return True
    except Exception as e:
        print(f"[HAPROXY_CONFIG] Error checking backend: {e}", flush=True)
    
    return False

def count_frontends_and_backends():
    if not os.path.exists(HAPROXY_CFG):
        return 0, 0, 0, 0, 0
    
    frontend_count = 0
    backend_count = 0
    acl_count = 0
    layer7_count = 0
    layer4_count = 0
    
    try:
        with open(HAPROXY_CFG, 'r') as haproxy_cfg:
            content = haproxy_cfg.read()
            lines = content.split('\n')
            
            for line in lines:
                line_stripped = line.strip()
                if line_stripped.startswith('frontend'):
                    frontend_count += 1
                    if 'mode http' in content:
                        layer7_count += 1
                    elif 'mode tcp' in content:
                        layer4_count += 1
                elif line_stripped.startswith('backend'):
                    backend_count += 1
                elif line_stripped.startswith('acl'):
                    acl_count += 1
    except Exception as e:
        print(f"[HAPROXY_CONFIG] Error counting: {e}", flush=True)
    
    return frontend_count, backend_count, acl_count, layer7_count, layer4_count

def update_haproxy_config(frontend_name, frontend_ip, frontend_port, lb_method, protocol, backend_name,
                          backend_servers, health_check, health_check_tcp, health_check_link, sticky_session,
                          add_header, header_name, header_value, sticky_session_type, is_acl, acl_name,
                          acl_action, acl_backend_name, use_ssl, ssl_cert_path, https_redirect, is_dos,
                          ban_duration, limit_requests, forward_for, is_forbidden_path, forbidden_name,
                          allowed_ip, forbidden_path, sql_injection_check, is_xss, is_remote_upload,
                          add_path_based, redirect_domain_name, root_redirect, redirect_to, is_webshells,
                          del_server_header=False, backend_ssl_redirect=False, ssl_redirect_backend_name='',
                          ssl_redirect_port='80', frontend_hostname='', add_custom_acl=False,
                          custom_acl_name='', custom_acl_type='path_beg', custom_acl_value='',
                          custom_acl_action='route', custom_acl_backend='', custom_acl_redirect_url=''):
    
    os.makedirs(os.path.dirname(HAPROXY_CFG), exist_ok=True)
    
    unique_backend_name = f"{backend_name}_{sanitize_name(frontend_hostname)}" if frontend_hostname else backend_name
    
    if is_backend_exist(unique_backend_name):
        return f"Backend {unique_backend_name} already exists. Cannot add duplicate."
    
    is_no_lb = lb_method == 'no-lb'
    if is_no_lb and len(backend_servers) > 1:
        backend_servers = backend_servers[:1]
    
    try:
        # ===== CHECK IF FRONTEND EXISTS AT PORT =====
        existing_frontend = frontend_exists_at_port(frontend_ip, frontend_port)
        
        if existing_frontend:
            print(f"[HAPROXY] Found existing frontend '{existing_frontend}' at {frontend_ip}:{frontend_port}", flush=True)
            
            with open(HAPROXY_CFG, 'a') as haproxy_cfg:
                # ===== BACKEND =====
                haproxy_cfg.write(f"\nbackend {unique_backend_name}\n")
                
                if not is_no_lb:
                    haproxy_cfg.write(f"    balance {lb_method}\n")
                
                if sticky_session and not is_no_lb:
                    if sticky_session_type == "cookie":
                        haproxy_cfg.write(f"    cookie SERVERID insert indirect nocache\n")
                    elif sticky_session_type == "source":
                        haproxy_cfg.write(f"    stick-table type ip size 200k expire 30m\n")
                        haproxy_cfg.write(f"    stick on src\n")
                
                if health_check and protocol == 'http':
                    haproxy_cfg.write(f"    option httpchk GET {health_check_link}\n")
                elif health_check_tcp and protocol == 'tcp':
                    haproxy_cfg.write(f"    option tcp-check\n")
                
                if add_header:
                    haproxy_cfg.write(f"    http-response add-header {header_name} {header_value}\n")
                
                if del_server_header:
                    haproxy_cfg.write(f"    http-response del-header Server\n")
                
                if forward_for:
                    haproxy_cfg.write(f"    option forwardfor\n")
                
                # Add servers
                for server_name, server_ip, server_port, maxconn in backend_servers:
                    maxconn_str = f" maxconn {maxconn}" if maxconn else ""
                    
                    if health_check and protocol == 'http':
                        haproxy_cfg.write(f"    server {server_name} {server_ip}:{server_port}{maxconn_str} check\n")
                    else:
                        haproxy_cfg.write(f"    server {server_name} {server_ip}:{server_port}{maxconn_str}\n")
            
            acl_name_sanitized = f"is_{sanitize_name(frontend_hostname)}" if frontend_hostname else f"is_{unique_backend_name}"
            add_acl_to_frontend(existing_frontend, acl_name_sanitized, frontend_hostname or 'localhost', unique_backend_name)
            
            # ===== REDIRECT HTTP→HTTPS (jeśli zaznaczony) =====
            if backend_ssl_redirect and ssl_redirect_backend_name:
                unique_redirect_backend_name = f"{ssl_redirect_backend_name}_redirect_{sanitize_name(frontend_hostname)}" if frontend_hostname else f"{ssl_redirect_backend_name}_redirect"
                
                existing_http_frontend = frontend_exists_at_port(frontend_ip, ssl_redirect_port)
                
                if existing_http_frontend:
                    print(f"[HAPROXY] Adding redirect ACL to existing HTTP frontend '{existing_http_frontend}'", flush=True)
                    
                    with open(HAPROXY_CFG, 'a') as haproxy_cfg:
                        haproxy_cfg.write(f"\nbackend {unique_redirect_backend_name}\n")
                        haproxy_cfg.write(f"    mode http\n")
                        haproxy_cfg.write(f"    redirect scheme https code 301 if !{{ ssl_fc }}\n")
                    
                    if frontend_hostname:
                        acl_name_redirect = f"is_{sanitize_name(frontend_hostname)}_redirect"
                        add_acl_to_frontend(existing_http_frontend, acl_name_redirect, frontend_hostname, unique_redirect_backend_name)
                else:
                    print(f"[HAPROXY] Creating new HTTP redirect frontend at {frontend_ip}:{ssl_redirect_port}", flush=True)
                    
                    with open(HAPROXY_CFG, 'a') as haproxy_cfg:
                        generic_http_redirect_name = f"http_redirect_frontend"
                        
                        haproxy_cfg.write(f"\nfrontend {generic_http_redirect_name}\n")
                        haproxy_cfg.write(f"    bind {frontend_ip}:{ssl_redirect_port}\n")
                        haproxy_cfg.write(f"    mode http\n")
                        
                        if frontend_hostname:
                            acl_name_redirect = f"is_{sanitize_name(frontend_hostname)}_redirect"
                            haproxy_cfg.write(f"    acl {acl_name_redirect} hdr(host) -i {frontend_hostname}\n")
                            haproxy_cfg.write(f"    use_backend {unique_redirect_backend_name} if {acl_name_redirect}\n")
                        else:
                            haproxy_cfg.write(f"    default_backend {unique_redirect_backend_name}\n")
                        
                        # Redirect backend
                        haproxy_cfg.write(f"\nbackend {unique_redirect_backend_name}\n")
                        haproxy_cfg.write(f"    mode http\n")
                        haproxy_cfg.write(f"    redirect scheme https code 301 if !{{ ssl_fc }}\n")
            
            return f"Backend added to existing frontend"
        
        # ===== TWORZENIE NOWEGO FRONTENDU (GENERYCZNE NAZWY) =====
        # Generuj generyczną nazwę frontendu
        generic_frontend_name = f"https_frontend" if use_ssl else f"http_frontend"
        
        print(f"[HAPROXY] Creating new frontend '{generic_frontend_name}' at {frontend_ip}:{frontend_port}", flush=True)
        
        with open(HAPROXY_CFG, 'a') as haproxy_cfg:
            # ===== PRIMARY FRONTEND (GENERIC NAME) =====
            haproxy_cfg.write(f"\nfrontend {generic_frontend_name}\n")
            haproxy_cfg.write(f"    bind {frontend_ip}:{frontend_port}")
            
            if use_ssl:
                haproxy_cfg.write(f" ssl crt {ssl_cert_path}")
            
            haproxy_cfg.write("\n")
            
            # Headers zaraz po BIND/CERT
            haproxy_cfg.write(f"    http-request set-header X-Forwarded-For %[src]\n")
            if use_ssl:
                haproxy_cfg.write(f"    http-request set-header X-Forwarded-Proto https\n")
            else:
                haproxy_cfg.write(f"    http-request set-header X-Forwarded-Proto http\n")
            
            haproxy_cfg.write(f"    mode {protocol}\n")
            
            # ACL dla pierwszego vhost
            acl_name_sanitized = None
            if frontend_hostname:
                acl_name_sanitized = f"is_{sanitize_name(frontend_hostname)}"
                haproxy_cfg.write(f"    acl {acl_name_sanitized} hdr(host) -i {frontend_hostname}\n")
            
            if not is_no_lb:
                haproxy_cfg.write(f"    balance {lb_method}\n")
                if forward_for:
                    haproxy_cfg.write(f"    option forwardfor\n")
            
            # Protections
            if is_dos:
                haproxy_cfg.write(f"    stick-table type ip size 1m expire {ban_duration} store http_req_rate(1m)\n")
                haproxy_cfg.write(f"    http-request track-sc0 src\n")
                haproxy_cfg.write(f"    acl abuse sc_http_req_rate(0) gt {limit_requests}\n")
                haproxy_cfg.write(f"    http-request silent-drop if abuse\n")
            
            if sql_injection_check:
                haproxy_cfg.write("    acl is_sql_injection urlp_reg -i (union|select|insert|update|delete|drop|@@|1=1|`1)\n")
                haproxy_cfg.write("    acl is_long_uri path_len gt 400\n")
                haproxy_cfg.write("    acl semicolon_path path_reg -i ^.*;.*\n")
                haproxy_cfg.write("    acl is_sql_injection2 urlp_reg -i (;|substring|extract|union\\s+all|order\\s+by)\\s+(\\d+|--\\+)\n")
                haproxy_cfg.write(f"    http-request deny if is_sql_injection or is_long_uri or semicolon_path or is_sql_injection2\n")
            
            if is_xss:
                haproxy_cfg.write("    acl is_xss_attack urlp_reg -i (<|>|script|alert|onerror|onload|javascript)\n")
                haproxy_cfg.write("    acl is_xss_attack_2 urlp_reg -i (<\\s*script\\s*|javascript:|<\\s*img\\s*src\\s*=|<\\s*a\\s*href\\s*=|<\\s*iframe\\s*src\\s*=|\\bon\\w+\\s*=|<\\s*input\\s*[^>]*\\s*value\\s*=|<\\s*form\\s*action\\s*=|<\\s*svg\\s*on\\w+\\s*=)\n")
                haproxy_cfg.write("    acl is_xss_attack_hdr hdr_reg(Cookie|Referer|User-Agent) -i (<|>|script|alert|onerror|onload|javascript)\n")
                haproxy_cfg.write(f"    http-request deny if is_xss_attack or is_xss_attack_2 or is_xss_attack_hdr\n")
            
            if is_webshells:
                haproxy_cfg.write("    acl blocked_webshell path_reg -i /(cmd|shell|backdoor|webshell|phpspy|c99|kacak|b374k|log4j|log4shell|wsos|madspot|malicious|evil).*\\.php.*\n")
                haproxy_cfg.write(f"    http-request deny if blocked_webshell\n")
            
            if https_redirect:
                haproxy_cfg.write(f"    redirect scheme https code 301 if !{{ ssl_fc }}\n")
            
            if del_server_header:
                haproxy_cfg.write(f"    http-response del-header Server\n")
            
            # Backend routing
            if acl_name_sanitized:
                haproxy_cfg.write(f"    use_backend {unique_backend_name} if {acl_name_sanitized}\n")
            else:
                haproxy_cfg.write(f"    default_backend {unique_backend_name}\n")
            
            # ===== BACKEND =====
            haproxy_cfg.write(f"\nbackend {unique_backend_name}\n")
            
            if not is_no_lb:
                haproxy_cfg.write(f"    balance {lb_method}\n")
            
            if sticky_session and not is_no_lb:
                if sticky_session_type == "cookie":
                    haproxy_cfg.write(f"    cookie SERVERID insert indirect nocache\n")
                elif sticky_session_type == "source":
                    haproxy_cfg.write(f"    stick-table type ip size 200k expire 30m\n")
                    haproxy_cfg.write(f"    stick on src\n")
            
            if health_check and protocol == 'http':
                haproxy_cfg.write(f"    option httpchk GET {health_check_link}\n")
            elif health_check_tcp and protocol == 'tcp':
                haproxy_cfg.write(f"    option tcp-check\n")
            
            if add_header:
                haproxy_cfg.write(f"    http-response add-header {header_name} {header_value}\n")
            
            if del_server_header:
                haproxy_cfg.write(f"    http-response del-header Server\n")
            
            if forward_for:
                haproxy_cfg.write(f"    option forwardfor\n")
            
            for server_name, server_ip, server_port, maxconn in backend_servers:
                maxconn_str = f" maxconn {maxconn}" if maxconn else ""
                
                if health_check and protocol == 'http':
                    haproxy_cfg.write(f"    server {server_name} {server_ip}:{server_port}{maxconn_str} check\n")
                else:
                    haproxy_cfg.write(f"    server {server_name} {server_ip}:{server_port}{maxconn_str}\n")
            
            # ===== REDIRECT HTTP -> HTTPS (GENERIC NAME) =====
            if backend_ssl_redirect and ssl_redirect_backend_name:
                unique_redirect_backend_name = f"{ssl_redirect_backend_name}_redirect_{sanitize_name(frontend_hostname)}" if frontend_hostname else f"{ssl_redirect_backend_name}_redirect"
                
                # Check if HTTP frontend exists
                existing_http_frontend = frontend_exists_at_port(frontend_ip, ssl_redirect_port)
                
                if not existing_http_frontend:
                    generic_http_redirect_name = f"http_redirect_frontend"
                    
                    haproxy_cfg.write(f"\nfrontend {generic_http_redirect_name}\n")
                    haproxy_cfg.write(f"    bind {frontend_ip}:{ssl_redirect_port}\n")
                    haproxy_cfg.write(f"    mode http\n")
                    
                    if frontend_hostname:
                        acl_name_redirect = f"is_{sanitize_name(frontend_hostname)}_redirect"
                        haproxy_cfg.write(f"    acl {acl_name_redirect} hdr(host) -i {frontend_hostname}\n")
                        haproxy_cfg.write(f"    use_backend {unique_redirect_backend_name} if {acl_name_redirect}\n")
                    else:
                        haproxy_cfg.write(f"    default_backend {unique_redirect_backend_name}\n")
                else:
                    if frontend_hostname:
                        acl_name_redirect = f"is_{sanitize_name(frontend_hostname)}_redirect"
                        add_acl_to_frontend(existing_http_frontend, acl_name_redirect, frontend_hostname, unique_redirect_backend_name)
                
                # Redirect backend
                haproxy_cfg.write(f"\nbackend {unique_redirect_backend_name}\n")
                haproxy_cfg.write(f"    mode http\n")
                haproxy_cfg.write(f"    redirect scheme https code 301 if !{{ ssl_fc }}\n")
            
            return "Configuration updated successfully!"
        
    except Exception as e:
        print(f"[HAPROXY_CONFIG] Error updating config: {e}", flush=True)
        return f"Error: {e}"