import os
import sys
import ssl
import json
import time
import logging
import logging.handlers
import datetime
import traceback
import signal
import subprocess
from pathlib import Path

import colorama
import gevent
from gevent import monkey
from gevent.pywsgi import WSGIServer
from gevent.pool import Pool
from gevent.threadpool import ThreadPool

monkey.patch_all()

os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'gvsds.settings')

import django
django.setup()

CPU_CORES = os.cpu_count() or 4
MAX_COROUTINES = CPU_CORES * 50
MAX_THREADS = CPU_CORES * 2

from gvsds.secrets import get as _secret_get
SQL_IP = _secret_get('database', 'host', fallback='127.0.0.1')

WHITE_HOSTS = [
    "www.gvsds.com", "gvsds.com", "api.gvsds.com", "127.0.0.1", "localhost", SQL_IP,
    "gvsds.dev", "yc8z.org", "ycbz.org", "gvsds.org", "gvsds.studio", "114.66.59.25",
    "librok.gvsds.com"
]
BAN_IPS = [
    "www.kanglaohui.com", "www.best-loving.cn", "11.4.5.14", "185.223.152.189"
]

PIRACY_DIR = Path("TempLate") / "Errors" / "Piracy"
PIRACY_LIST = os.listdir(PIRACY_DIR) if PIRACY_DIR.exists() else []

# 静态站点：host -> 站点根目录
STATIC_SITES = {
    "librok.gvsds.com": Path("TempLate") / "rok" / "lib",
}

LOG_CONFIG = {
    'disable_static_success': False,
    'disable_favicon': False,
    'disable_all_static': False,
    'color_whole_line': True,
    'static_extensions': ['.js', '.css', '.png', '.jpg', '.jpeg', '.gif', '.ico', '.svg', '.woff', '.woff2', '.ttf', '.eot'],
    'log_level': 'INFO',
    'show_domain': True,
    'show_full_host': True,
}

STATUS_COLOR_MAP = {
    200: colorama.Fore.GREEN,
    301: colorama.Fore.CYAN,
    302: colorama.Fore.CYAN,
    304: colorama.Fore.CYAN,
    400: colorama.Fore.YELLOW,
    403: colorama.Fore.RED,
    404: colorama.Fore.MAGENTA,
    500: colorama.Fore.RED + colorama.Style.BRIGHT,
    502: colorama.Fore.RED + colorama.Style.BRIGHT,
    503: colorama.Fore.RED + colorama.Style.BRIGHT,
}

LEVEL_COLOR_MAP = {
    'DEBUG': colorama.Fore.CYAN,
    'INFO': colorama.Fore.GREEN,
    'WARNING': colorama.Fore.YELLOW,
    'ERROR': colorama.Fore.RED,
    'CRITICAL': colorama.Fore.RED + colorama.Style.BRIGHT,
}

logger = None


class NoOpLog:
    @staticmethod
    def write(msg):
        pass
    @staticmethod
    def flush():
        pass


class CustomWSGIServer(WSGIServer):
    def handle_error(self, request, client_address):
        exc_type, exc_value, exc_traceback = self.get_exc_info()
        ignored_exceptions = (
            ssl.SSLEOFError, ConnectionResetError, BrokenPipeError, OSError, EOFError
        )
        if isinstance(exc_value, ignored_exceptions):
            logger.debug(f"客户端 {client_address} 连接异常断开: {exc_value}",
                         extra={'domain': 'system', 'status_code': 0})
            return
        super().handle_error(request, client_address)


class WholeLineColorFormatter(logging.Formatter):
    def format(self, record):
        timestamp = self.formatTime(record, datefmt='%Y-%m-%d %H:%M:%S')
        levelname = record.levelname
        message = record.getMessage()
        status_code = getattr(record, 'status_code', 0)
        domain = getattr(record, 'domain', '')
        host = getattr(record, 'host', '')

        if LOG_CONFIG['color_whole_line'] and status_code:
            color = STATUS_COLOR_MAP.get(status_code, LEVEL_COLOR_MAP.get(levelname, colorama.Fore.WHITE))
        else:
            color = LEVEL_COLOR_MAP.get(levelname, colorama.Fore.WHITE)

        log_parts = [f"[{timestamp}]", f"[{levelname:7s}]"]
        if LOG_CONFIG['show_domain']:
            if domain:
                log_parts.append(f"[域:{domain:20s}]")
            elif host:
                log_parts.append(f"[Host:{host:20s}]")

        log_parts.append(message)
        log_line = " ".join(log_parts)

        if LOG_CONFIG['color_whole_line']:
            log_line = f"{color}{log_line}{colorama.Style.RESET_ALL}"
        return log_line


class RequestFilter:
    @staticmethod
    def is_favicon(path):
        return path.lower() == '/favicon.ico'

    @staticmethod
    def is_static_file(path):
        ext = os.path.splitext(path)[1].lower()
        return ext in LOG_CONFIG['static_extensions']

    @staticmethod
    def should_skip_log(path, status_code):
        if LOG_CONFIG['disable_favicon'] and RequestFilter.is_favicon(path):
            return True
        if RequestFilter.is_static_file(path):
            if LOG_CONFIG['disable_all_static']:
                return True
            if LOG_CONFIG['disable_static_success'] and status_code < 400:
                return True
        return False

    @staticmethod
    def extract_domain(host):
        if not host:
            return ''
        return host.split(':')[0] if ':' in host else host

def GetIP(environ: dict = {}):
    x_forwarded_for_raw = environ.get('HTTP_X_FORWARDED_FOR', '')
    x_forwarded_for_list = [x.strip() for x in x_forwarded_for_raw.split(',') if x.strip()]
    x_forwarded_for = x_forwarded_for_list[0] if x_forwarded_for_list else ''
    return x_forwarded_for

class FilteredLoggingMiddleware:
    def __init__(self, app):
        self.app = app

    def __call__(self, environ, start_response):
        path = environ.get('PATH_INFO', '')
        query_string = environ.get('QUERY_STRING', '')
        full_path = path + (f"?{query_string}" if query_string else "")
        method = environ.get('REQUEST_METHOD', '')
        ip = environ.get('REMOTE_ADDR', '')
        host = environ.get('HTTP_HOST', '')
        domain = RequestFilter.extract_domain(host)
        start_time = datetime.datetime.now()

        real_ip = GetIP(environ)
        real_ip = real_ip if real_ip else ip


        if RequestFilter.should_skip_log(full_path, 200):
            def wrapped_start_response(status, headers, exc_info=None):
                return start_response(status, headers, exc_info)
            return self.app(environ, wrapped_start_response)

        def wrapped_start_response(status, headers, exc_info=None):
            try:
                status_code = int(status.split()[0]) if status else 0

                if RequestFilter.should_skip_log(full_path, status_code):
                    return start_response(status, headers, exc_info)

                duration = (datetime.datetime.now() - start_time).total_seconds()
                log_message = f"[{status_code:3d}] [{method.ljust(4)[:4]}] [{real_ip:15}] [{duration:.6f}s] {full_path}"

                extra = {
                    'status_code': status_code,
                    'domain': domain,
                    'host': host if LOG_CONFIG['show_full_host'] else ''
                }

                if status_code >= 500:
                    logger.error(log_message, extra=extra, exc_info=exc_info)
                elif status_code >= 400:
                    logger.warning(log_message, extra=extra)
                else:
                    logger.info(log_message, extra=extra)
            except Exception as e:
                logger.error(f"日志记录失败: {e}", exc_info=True)

            return start_response(status, headers, exc_info)

        return self.app(environ, wrapped_start_response)

import threading
import time
import re

class SecurityMiddleware:
    def __init__(self, app):
        self.app = app

        # 恶意参数正则
        self.bad_reg = re.compile(r'^(?=.{8,})(\d+[A-Za-z&]{1,}|[A-Za-z]\d+|&\d+)', re.I)

        # ==================== 文件配置 ====================
        self.ip_file = "./ips.txt"                # 黑名单IP
        self.ip_times_file = "./iptimes.txt"      # IP:次数
        self.total_times_file = "./times.txt"     # 总拦截次数

        # ==================== 内存存储 ====================
        self.seen_ips = self.load_file_to_set(self.ip_file)
        self.ip_counter = self.load_file_to_dict(self.ip_times_file)  # {ip: count}
        self.total_count = self.load_total_count()

        # ==================== 线程队列 & 异步写入 ====================
        self.queue = []
        self.queue_lock = threading.Lock()
        self.sync_interval = 5  # 每5秒批量落地一次
        self.start_async_saver()

    # ==================== 加载工具 ====================
    def load_file_to_set(self, path):
        s = set()
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    ip = line.strip()
                    if ip:
                        s.add(ip)
        return s

    def load_file_to_dict(self, path):
        d = {}
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line or ":" not in line:
                        continue
                    ip, cnt = line.split(":", 1)
                    try:
                        d[ip.strip()] = int(cnt.strip())
                    except:
                        continue
        return d

    def load_total_count(self):
        if os.path.exists(self.total_times_file):
            try:
                with open(self.total_times_file, "r", encoding="utf-8") as f:
                    return int(f.read().strip())
            except:
                return 0
        return 0

    # ==================== 异步落地方案（核心） ====================
    def start_async_saver(self):
        def saver():
            while True:
                time.sleep(self.sync_interval)
                self.sync_to_disk()

        t = threading.Thread(target=saver, daemon=True)
        t.start()

    def sync_to_disk(self):
        with self.queue_lock:
            if not self.queue:
                return
            queue_copy = self.queue.copy()
            self.queue.clear()

        # 批量更新
        for ip in queue_copy:
            self.ip_counter[ip] = self.ip_counter.get(ip, 0) + 1
            self.total_count += 1
            self.seen_ips.add(ip)

        # 写入磁盘
        with open(self.ip_times_file, "w", encoding="utf-8") as f:
            for ip, cnt in self.ip_counter.items():
                f.write(f"{ip}:{cnt}\n")

        with open(self.ip_file, "w", encoding="utf-8") as f:
            for ip in self.seen_ips:
                f.write(f"{ip}\n")

        with open(self.total_times_file, "w", encoding="utf-8") as f:
            f.write(str(self.total_count))

    # ==================== 记录IP（入队列，不阻塞） ====================
    def add_ip_to_queue(self, ip):
        with self.queue_lock:
            self.queue.append(ip)

    # ==================== 获取客户端IP ====================
    def get_ip(self, environ):
        forward = environ.get('HTTP_X_FORWARDED_FOR', '')
        if forward:
            ip = forward.split(',')[0].strip()
            if ip:
                return ip
        return environ.get('REMOTE_ADDR', '127.0.0.1')

    # ==================== 请求入口 ====================
    def __call__(self, environ, start_response):
        query = environ.get('QUERY_STRING', '')

        # ==================== 拦截恶意参数 ====================
        if query and self.bad_reg.search(query):
            ip = self.get_ip(environ)
            self.add_ip_to_queue(ip)

            # 直接关闭连接，不产生日志
            start_response('444 No Response', [])
            return []

        # ==================== 原有逻辑保持不变 ====================
        host = environ.get('HTTP_HOST', '').lower()
        path = environ.get('PATH_INFO', '')

        # 跳转到 www.gvsds.com
        if host == "gvsds.com":
            location = f'https://www.gvsds.com{path}'
            if query:
                location += f'?{query}'
            start_response('301 Moved Permanently', [('Location', location), ('Content-Type', 'text/plain')])
            return [b'Redirecting to www.gvsds.com...']

        # robots.txt
        if path == "/robots.txt":
            start_response('200 OK', [('Content-Type', 'text/plain')])
            return [b'User-agent: *\nDisallow: /assets/']

        # 盗版域名拦截
        domain_suffix = ".".join(host.split(".")[-2:]) if "." in host else host
        if domain_suffix in PIRACY_LIST:
            try:
                content = open(PIRACY_DIR / domain_suffix, encoding="utf-8").read()
            except Exception:
                content = "Access Denied"
            start_response('403 Forbidden', [('Content-Type', 'text/html')])
            return [content.encode('utf-8')]

        # 非法HOST + 黑名单IP
        x_forwarded_for = self.get_ip(environ)
        host_no_port = host.split(':')[0] if ':' in host else host
        if (host not in WHITE_HOSTS and host_no_port not in WHITE_HOSTS
                and domain_suffix not in WHITE_HOSTS) or x_forwarded_for in BAN_IPS:
            try:
                content = open(Path("TempLate") / "Errors" / "_403_HOST.html", encoding="utf-8").read()
            except Exception:
                content = "Forbidden"
            start_response('403 Forbidden', [('Content-Type', 'text/html')])
            return [content.encode('utf-8')]

        # 拦截脚本后缀
        lower_path = path.lower().split("?")[0]
        if "." in lower_path and lower_path.split(".")[-1] in ["php", "asp", "aspx", "cgi"]:
            start_response('403 Forbidden', [('Content-Type', 'text/plain')])
            return [b'Server does not support this type of page.']

        # 静态站点路由：将指定 host 的请求直接映射到本地目录
        if host in STATIC_SITES:
            return self._serve_static_site(environ, start_response, host, path)

        return self.app(environ, start_response)

    def _serve_static_site(self, environ, start_response, host, path):
        """为静态站点（host -> 目录）提供文件服务
        - 默认入口：目录请求 / 返回 index.html
        - 仅允许访问站点根目录内的文件（防止 ../ 越权）
        """
        import mimetypes
        mimetypes.add_type('application/javascript', '.js')
        mimetypes.add_type('text/css', '.css')
        mimetypes.add_type('image/svg+xml', '.svg')

        site_root = STATIC_SITES[host].resolve()
        # 清理路径：去掉前导 /，URL 反解码，过滤 ..
        rel_path = path.lstrip('/')
        if not rel_path:
            rel_path = 'index.html'
        try:
            from urllib.parse import unquote
            rel_path = unquote(rel_path)
        except Exception:
            pass
        # 防止目录越权
        if '..' in rel_path.split('/'):
            start_response('403 Forbidden', [('Content-Type', 'text/plain')])
            return [b'Forbidden']

        target = (site_root / rel_path).resolve()
        try:
            target.relative_to(site_root)
        except ValueError:
            start_response('403 Forbidden', [('Content-Type', 'text/plain')])
            return [b'Forbidden']

        # 目录则尝试 index.html
        if target.is_dir():
            target = target / 'index.html'
        if not target.is_file():
            start_response('404 Not Found', [('Content-Type', 'text/html')])
            try:
                content = open(Path("TempLate") / "Errors" / "404.html", encoding="utf-8").read().encode('utf-8')
            except Exception:
                content = b'404 Not Found'
            return [content]

        try:
            data = target.read_bytes()
        except Exception as e:
            start_response('500 Internal Server Error', [('Content-Type', 'text/plain')])
            return [f'Read error: {e}'.encode('utf-8')]

        ctype, _ = mimetypes.guess_type(str(target))
        if not ctype:
            ctype = 'application/octet-stream'
        start_response('200 OK', [
            ('Content-Type', ctype),
            ('Content-Length', str(len(data))),
            ('Cache-Control', 'no-cache'),
        ])
        return [data]


def SetupLogging():
    global logger
    log_dir = "./logs"
    os.makedirs(log_dir, exist_ok=True)
    log_filename = os.path.join(log_dir, f"server_{datetime.datetime.now().strftime('%Y-%m-%d')}.log")

    log_level = getattr(logging, LOG_CONFIG['log_level'].upper(), logging.INFO)
    formatter = WholeLineColorFormatter()

    file_handler = logging.handlers.RotatingFileHandler(
        log_filename, maxBytes=10*1024*1024, backupCount=5, encoding='utf-8'
    )
    file_handler.setFormatter(formatter)

    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)

    logging.basicConfig(level=log_level, handlers=[file_handler, console_handler], force=True)

    for lib in ['django', 'django.request', 'django.server', 'django.db.backends']:
        logging.getLogger(lib).setLevel(logging.WARNING)

    logger = logging.getLogger('server')
    logger.setLevel(log_level)
    logger.propagate = False

    if not logger.handlers:
        logger.addHandler(file_handler)
        logger.addHandler(console_handler)

    return logger


class ConnectionPoolManager:
    _instance = None
    _lock = gevent.lock.BoundedSemaphore(1)

    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._init_pools()
        return cls._instance

    def _init_pools(self):
        self.coroutine_pool = Pool(MAX_COROUTINES)
        self.thread_pool = ThreadPool(MAX_THREADS)


def CreateSslContext(cert_path, key_path):
    context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
    context.load_cert_chain(certfile=cert_path, keyfile=key_path)
    context.minimum_version = ssl.TLSVersion.TLSv1_2
    context.maximum_version = ssl.TLSVersion.TLSv1_3
    context.set_ciphers('ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305')
    context.options |= ssl.OP_CIPHER_SERVER_PREFERENCE
    return context


def RedirectApp(environ, start_response):
    query = environ.get('QUERY_STRING', '')
    path = environ.get('PATH_INFO', '')
    if query:
        path += '?' + query
    host = environ.get('HTTP_HOST', '')
    location = f'https://{host}{path}'

    domain = RequestFilter.extract_domain(host)
    logger.info(f"[301] [GET] {path} - {environ.get('REMOTE_ADDR', '')} - 0.000000s",
                extra={'status_code': 301, 'domain': domain, 'host': host})

    start_response('301 Moved Permanently', [
        ('Location', location),
        ('Content-Type', 'text/plain'),
        ('Access-Control-Allow-Origin', f'https://{host}'),
        ('Access-Control-Allow-Credentials', 'true')
    ])
    return [b'Redirecting to HTTPS...']


def get_django_wsgi_application():
    from django.core.wsgi import get_wsgi_application
    app = get_wsgi_application()
    app = FilteredLoggingMiddleware(app)
    app = SecurityMiddleware(app)
    from gvsdsdk.msync.sio import create_sio_app
    app = create_sio_app(app)
    return app


class HotReloader:
    def __init__(self, watch_dirs=None, extensions=None, debounce=1.0):
        self.watch_dirs = watch_dirs or ['./gvsds']
        self.extensions = extensions or {'.py'}
        self.debounce = debounce
        self._observer = None
        self._changed = gevent.event.Event()
        self._last_change_time = 0
        self._lock = gevent.lock.BoundedSemaphore(1)

    def _should_watch(self, path):
        ext = os.path.splitext(path)[1].lower()
        if ext not in self.extensions:
            return False
        parts = os.path.normpath(path).split(os.sep)
        if '__pycache__' in parts or '.git' in parts:
            return False
        return True

    def _on_change(self, event):
        if event.is_directory:
            return
        if not self._should_watch(event.src_path):
            return
        now = time.time()
        with self._lock:
            if now - self._last_change_time < self.debounce:
                return
            self._last_change_time = now
        rel_path = os.path.relpath(event.src_path, '.')
        logger.warning(f"[热重启] 检测到文件变更: {rel_path}", extra={'domain': 'system', 'status_code': 0})
        self._changed.set()

    def start(self):
        from watchdog.observers import Observer
        from watchdog.events import FileSystemEventHandler

        class Handler(FileSystemEventHandler):
            def __init__(self, callback):
                super().__init__()
                self.callback = callback

            def on_modified(self, event):
                self.callback(event)

            def on_created(self, event):
                self.callback(event)

        self._observer = Observer()
        handler = Handler(self._on_change)
        for d in self.watch_dirs:
            abs_dir = os.path.abspath(d)
            if os.path.isdir(abs_dir):
                self._observer.schedule(handler, abs_dir, recursive=True)
        self._observer.daemon = True
        self._observer.start()
        logger.info(f"[热重启] 文件监听已启动: {self.watch_dirs}", extra={'domain': 'system', 'status_code': 0})

    def stop(self):
        if self._observer:
            self._observer.stop()
            self._observer.join(timeout=3)

    def wait_for_change(self, timeout=None):
        self._changed.clear()
        return self._changed.wait(timeout)

    @property
    def is_changed(self):
        return self._changed.is_set()


class ServerManager:
    def __init__(self):
        global logger
        if logger is None:
            SetupLogging()

        self.pool_manager = ConnectionPoolManager()
        self.servers = []
        self.ssl_context = CreateSslContext(
            os.environ.get('SSL_CERT_PATH', './SSL/gvsds.com_public.crt'),
            os.environ.get('SSL_KEY_PATH', './SSL/gvsds.com.key')
        )

        self.django_app = get_django_wsgi_application()

        logger.info("Django 应用初始化完成", extra={'domain': 'system', 'status_code': 0})

    def add_http_redirect_server(self, port=80):
        server = CustomWSGIServer(
            ('0.0.0.0', port), RedirectApp,
            spawn=self.pool_manager.coroutine_pool,
            log=NoOpLog(), error_log=NoOpLog()
        )
        self.servers.append(('HTTP Redirect', port, server, None))
        return self

    def add_django_server(self, port, is_ssl=True, app_name="Django App"):
        kwargs = {
            'listener': ('0.0.0.0', port),
            'application': self.django_app,
            'spawn': self.pool_manager.coroutine_pool,
            'log': NoOpLog(),
            'error_log': NoOpLog(),
            'suppress_ragged_eofs': True,
            'do_handshake_on_connect': False,
        }
        if is_ssl:
            kwargs['ssl_context'] = self.ssl_context

        server = CustomWSGIServer(**kwargs)
        self.servers.append((app_name, port, server, None))
        return self

    def start_all(self, hot_reload=False):
        logger.info("=== 启动 Django 服务器集群 ===", extra={'domain': 'system', 'status_code': 0})
        logger.info(f"协程池: {MAX_COROUTINES} | 线程池: {MAX_THREADS}", extra={'domain': 'system', 'status_code': 0})

        from django.urls import get_resolver
        resolver = get_resolver()
        logger.info("=== Django URL 路由已加载 ===", extra={'domain': 'system', 'status_code': 0})

        running_servers = []
        for name, port, server, _ in self.servers:
            logger.info(f"启动 {name} 服务器 on 0.0.0.0:{port}", extra={'domain': 'system', 'status_code': 0})
            gevent.sleep(0.1)
            g = gevent.spawn(server.serve_forever)
            running_servers.append((name, port, server, g))

        self.servers = running_servers

        reloader = None
        if hot_reload:
            reloader = HotReloader(watch_dirs=['./gvsds', './Main.py'])
            reloader.start()

        try:
            while True:
                if hot_reload and reloader:
                    reloader.wait_for_change(timeout=1)
                    if reloader.is_changed:
                        logger.warning("[热重启] 正在重启服务器...", extra={'domain': 'system', 'status_code': 0})
                        self.stop_all()
                        reloader.stop()
                        os.execv(sys.executable, [sys.executable] + sys.argv)
                else:
                    gevent.sleep(3600)
        except KeyboardInterrupt:
            if reloader:
                reloader.stop()
            self.stop_all()

    def stop_all(self):
        logger.info("=== 停止所有服务器 ===", extra={'domain': 'system', 'status_code': 0})
        for item in self.servers:
            if len(item) == 4:
                _, _, server, g = item
                try:
                    server.stop()
                except:
                    pass
        try:
            self.pool_manager.coroutine_pool.kill()
            self.pool_manager.thread_pool.kill()
        except:
            pass
        logger.info("所有服务器已停止", extra={'domain': 'system', 'status_code': 0})


if __name__ == "__main__":
    colorama.init(autoreset=True)
    os.environ['GEVENT_THREADPOOL_SIZE'] = str(MAX_THREADS)

    hot_reload = '--reload' in sys.argv

    server_manager = None
    try:
        logger = SetupLogging()
        if hot_reload:
            logger.info("[热重启] 模式已启用 (--reload)", extra={'domain': 'system', 'status_code': 0})
        server_manager = ServerManager()
        server_manager.add_http_redirect_server(80)
        server_manager.add_django_server(443, is_ssl=True, app_name="Main Django App")
        server_manager.start_all(hot_reload=hot_reload)
    except Exception as e:
        traceback.print_exc()
        if logger:
            logger.critical(f"服务器启动失败: {e}", exc_info=True, extra={'domain': 'system', 'status_code': 500})
        if server_manager:
            server_manager.stop_all()
    finally:
        if server_manager:
            server_manager.stop_all()
        gevent.killall([], timeout=5)