# -*- coding: utf-8 -*-
from bcoding import bdecode
import configparser
from copy import deepcopy
from datetime import datetime
from io import StringIO
import email
import logging
import numpy as np
import os
import pathlib
import re
import requests
import socket
import ssl
import struct
import threading
from functools import lru_cache, wraps
from random import random

from pymrf4.entities import SiteConfig
# from pymrf4.namespaced_lock import NamespaceLock
from pymrf4.socks5_proxy import Socks5Handler, ThreadingTCPServer
from pymrf4.utils import get_random_ratio, convert_size, format_info_hash, format_info_hash_compact
from pymrf4.db import db
from pymrf4.models import Torrent, Announce

logger = logging.getLogger(__name__)

BUFSIZE = 2048
# MAX_UPLOAD_SPEED = 50 * 1024 * 1024   # 官方给出的数据是125MBps，这里求稳先给50MBps，[TODO] 而且这个数值是应该放在config中的
# MAX_SHARE_RATIO = 2
# Minimal numbers of leechers that should MRF stop modifying
# MIN_LEECHERS = 8
# Maximum overall safe upload speed
# MAX_OVERALL_UPLOAD_SPEED = 6.18 * 1024 * 1024

# THREAD_LOCAL = threading.local()

SERVER = None

'''
A class represents the request of a tracker, all information are from the
URI of a tracker request
'''
class TrackerRequest:
    RE_KEY = r'(?<=key=)[A-Fa-f0-9]+'
    RE_EVENT = r'(?<=event=)[A-Za-z]+'
    RE_INFO_HASH = r'(?<=info_hash=).[^&]+'
    RE_UPLOAD = r'(?<=uploaded=)\d+'
    RE_DOWNLOAD = r'(?<=downloaded=)\d+'
    RE_LEFT = r'(?<=left=)\d+'
    RE_PEER_ID = r'(?<=peer_id=).[^&]+'


    def __init__(self, key, info_hash, event, downloaded, uploaded, left, peer_id):
        self.key = key # session handle (libtorrent implementation)
        self.info_hash = info_hash
        self.event = event
        self.downloaded = downloaded
        self.uploaded = uploaded
        self.left = left
        self.peer_id = peer_id


def parse_announce_uri(request_uri: str) -> TrackerRequest:
    key = re.search(TrackerRequest.RE_KEY, request_uri)
    event = re.search(TrackerRequest.RE_EVENT, request_uri)
    info_hash = re.search(TrackerRequest.RE_INFO_HASH, request_uri)
    uploaded = re.search(TrackerRequest.RE_UPLOAD, request_uri)
    downloaded = re.search(TrackerRequest.RE_DOWNLOAD, request_uri)
    left = re.search(TrackerRequest.RE_LEFT, request_uri)
    peer_id = re.search(TrackerRequest.RE_PEER_ID, request_uri)

    key = key.group(0)
    if event:
        event = event.group(0)
    info_hash = info_hash.group(0)
    uploaded = uploaded.group(0)
    downloaded = downloaded.group(0)
    left = left.group(0)

    uploaded =  int(uploaded)
    downloaded = int(downloaded)
    left = int(left)

    return TrackerRequest(key, info_hash, event, downloaded, uploaded, left, peer_id)

def parse_http_message(msg):
    """
    return a tuple like [METHOD, URI, HEADERS]
    """
    startline, headers = msg.split("\r\n", 1)
    startline_elements = [ x.strip(' \t') for x in startline.split(' ')]
    method = startline_elements[0]
    uri = startline_elements[1]
    message = email.message_from_file(StringIO(headers))
    headers = dict(message.items())
    return (method, uri, headers)


def wrap_ssl_socket(plain_socket:socket):
    certfile = os.path.join(pathlib.Path(__file__).parent, "certs/cert.pem")
    keyfile = os.path.join(pathlib.Path(__file__).parent, "certs/key.pem")

    ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS)
    ssl_context.load_cert_chain(certfile, keyfile)

    # ssl_context.verify_mode = ssl.CERT_REQUIRED
    # ssl_context.load_verify_locations(cafile=os.path.relpath(certifi.where()),
    #    capath=None, cadata=None);
    
    wrapped_ssl_socket = ssl_context.wrap_socket(plain_socket, server_side=True, do_handshake_on_connect=True)
    return wrapped_ssl_socket


class MRFServer(ThreadingTCPServer):
    def set_config(self, config):
        self.config = config


    def find_config_for_tracker(self, hostname) -> SiteConfig:
        for k, v in self.config.items():
            if hostname in v.trackers:
                return v
        return None

def _parse_dns_query_name(data):
    """Extract the queried domain name from a raw DNS query packet."""
    offset = 12  # skip 12-byte header
    labels = []
    while offset < len(data):
        length = data[offset]
        if length == 0:
            offset += 1
            break
        labels.append(data[offset + 1:offset + 1 + length].decode())
        offset += 1 + length
    return ".".join(labels), offset


def _build_dns_response(query, ip_addr):
    """Build a minimal DNS A-record response for *query* returning *ip_addr*."""
    # Copy transaction ID from query
    txn_id = query[:2]
    flags = struct.pack("!H", 0x8180)  # QR=1, AA=1, RD=1, RA=1
    counts = struct.pack("!HHHH", 1, 1, 0, 0)  # QDCOUNT=1, ANCOUNT=1

    # Question section – copy from the original query
    _, qend = _parse_dns_query_name(query)
    qend += 4  # QTYPE(2) + QCLASS(2)
    question = query[12:qend]

    # Answer section: pointer to name at offset 12, TYPE=A, CLASS=IN, TTL=60
    answer = b"\xc0\x0c"
    answer += struct.pack("!HHIH", 1, 1, 60, 4)
    answer += socket.inet_aton(ip_addr)

    return txn_id + flags + counts + question + answer


DOH_SERVER = "https://223.5.5.5/resolve"  # DNS over HTTPS server URL


class MRFHandler(Socks5Handler):
    """
    Override
    """
    def resolve_domain_doh(self, domain):
        """Resolve a domain name to an IP address using DNS over HTTPS."""
        params = {"name": domain, "type": "A"}
        response = requests.get(DOH_SERVER, params=params)

        if response.status_code == 200:
            answer = response.json()["Answer"]
            if answer:
                logger.debug("Resolved domain %s to %s", domain, answer[0]["data"])
                return answer[0]["data"]

        raise Exception(f"Failed to resolve domain: {domain}")

    def relay_udp_payload(self, dst_addr, dst_port, payload):
        """For DNS requests (port 53), resolve via DoH; otherwise forward normally."""
        if dst_port == 53:
            return self._relay_dns_via_doh(payload)
        return super().relay_udp_payload(dst_addr, dst_port, payload)

    def _relay_dns_via_doh(self, dns_query):
        """Parse the DNS query, resolve the domain via DoH, and build a DNS response."""
        domain, _ = _parse_dns_query_name(dns_query)
        logger.debug("DNS via DoH: %s", domain)
        ip_addr = self.resolve_domain_doh(domain)
        return _build_dns_response(dns_query, ip_addr)

    def duplex_transport(self, remote_addr, remote_port):
        config = self.server.find_config_for_tracker(remote_addr)
        if config is not None: # or remote_port == 443):
            if remote_port == 443:
                self.https = True
                self.connection = wrap_ssl_socket(self.connection)
            else:
                self.https = False

            self.handle_request_over_http(remote_addr, config)
        else:
            super().duplex_transport(remote_addr, remote_port)


    def handle_request_over_http(self, hostname, config):
        req_data = self.connection.recv(BUFSIZE)
        if not req_data:
            return

        # Assumption 1: that BUFSIZE is everything that include whole message
        data_str = req_data.decode("ascii")

        method, uri, headers = parse_http_message(data_str)
        logger.debug(f"Request: {method} {uri} {headers}")

        is_tracker_announce = False
        is_tracker_request = False

        # Remove "Host" for "follow_redirects" to work
        if "Host" in headers: del headers["Host"]

        if re.match("^/announce\.php", uri, re.IGNORECASE) or re.match("^/announce", uri, re.IGNORECASE):
            logger.debug("ORIGINAL: \r\n" + data_str)
            is_tracker_request = True
            uri_mod, headers_mod, torrent_id = do_modify_3(uri, headers, config=config)

        elif re.match("^/scrape\.php", uri, re.IGNORECASE) or re.match("^/scrape", uri, re.IGNORECASE):
            is_tracker_request = True
            uri_mod = uri
            headers_mod = deepcopy(headers)
        else:
            # as is
            uri_mod = uri
            headers_mod = deepcopy(headers)
            headers_mod.pop("Referer", "")
            headers_mod["User-Agent"] = "Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/100.0.4896.127 Safari/537.36"

        if is_tracker_announce:
            headers_mod_str = "\r\n".join([f"{k}: {v}" for k, v in headers_mod.items()])
            http_request_str = f'{method} {uri_mod} HTTP/1.1\r\n{headers_mod_str}\r\n\r\n'
            logger.debug("Modified: \r\n" + http_request_str)

        try:
            scheme = "https" if self.https else "http"
            full_uri = f"{scheme}://{hostname}{uri_mod}"
            res = requests.get(full_uri, headers=headers_mod, allow_redirects=True)
        except:
            pass
        else:
            res_headers = deepcopy(res.headers)

            res_headers["Content-Length"] = len(res.content)
            if "Content-Encoding" in res_headers: del res_headers["Content-Encoding"] # plain
            if "Transfer-Encoding" in res_headers: del res_headers["Transfer-Encoding"]

            res_header_msg = "\r\n".join([f"{k}: {v}" for k, v in res_headers.items()])
            res_message = f"HTTP/1.1 {res.status_code} {res.reason}\r\n{res_header_msg}\r\n\r\n"
            res_all_bytes = res_message.encode("ascii") + res.content
            if is_tracker_announce:
                logger.debug(res_message + str(res.content[:200]) + "...[" + str(len(res.content)) + " bytes of content]")

            if is_tracker_request:
                try:
                    decoded_resp = bdecode(res.content)
                    logger.debug("Response decoded: ", decoded_resp)
                    if not "failure reason" in decoded_resp:
                        peers = len(decoded_resp.get('peers', [])) / 6
                        leechers = decoded_resp.get('incomplete', 0)
                        seeders = decoded_resp.get('complete', 0)
                        # current_torrent_id = THREAD_LOCAL.CURRENT_TORRENT_ID

                        if peers == 0 and leechers == 0 and seeders == 0:
                            logger.debug(f"Announce response all 0, ignore.")
                        elif torrent_id is not None:
                            current_torrent: Torrent = Torrent.get_by_id(torrent_id)

                            current_torrent.peers = peers
                            current_torrent.leechers = leechers
                            current_torrent.seeders = seeders
                            proper_info_hash = current_torrent.info_hash
                            logger.info(f"{proper_info_hash} leechers: {leechers} seeders: {seeders} peers: {peers}")

                            current_torrent.save()

                            # THREAD_LOCAL.CURRENT_TORRENT_ID = None
                    else:
                        logger.error(f"Error: {decoded_resp}")
                except Exception as e:
                    logger.error(e)

            self.connection.setblocking(True)
            try:
                self.connection.sendall(res_all_bytes)
            except ConnectionResetError:
                pass

            # sleep(1)
            self.connection.shutdown(socket.SHUT_RDWR)
        # finally:
        #     fcm.save_all()


def generate_exponential_speed(min_speed, max_speed, \
                               leechers, max_anticipate_leechers = 100, \
                               min_leecher_factor=0.125, scale_factor=4):
    # Generate a speed from an exponential distribution
    # The speed is affected by the number of leechers, more leechers higher speed
    # Several constants need to take care

    # Here 4 is a magic number, if need lower speed, adjust the number higher a bit to 4.8.
    # Here use 100 as a magic number, use leechers / 100 to calculate a factor, work along with SCALE_FACTOR

    scale = (max_speed - min_speed) / scale_factor * np.clip(leechers / max_anticipate_leechers, min_leecher_factor, 1)

    speed = np.random.exponential(scale=scale)

    return np.clip(speed, min_speed, max_speed)


# def lru_cache_by_argument(arg_pos, maxsize=128):
#     def decorator(func):
#         cache = {}

#         @wraps(func)
#         def wrapper(*args, **kwargs):
#             key = args[arg_pos]
#             if key in cache:
#                 logger.info("Cached result.")
#                 return cache[key]
#             result = func(*args, **kwargs)
#             if len(cache) >= maxsize:
#                 cache.pop(next(iter(cache)))  # Simple LRU eviction
#             cache[key] = result
#             return result

#         return wrapper

#     return decorator


@db.connection_context()
def do_modify_3(request_uri:str, headers:dict, config:SiteConfig):
    req = parse_announce_uri(request_uri)
    info_hash = format_info_hash(req.info_hash).decode()

    torrent, new_torrent = Torrent.get_or_create(info_hash=info_hash, site=config.name)

    update_stats_pre_mod(torrent, req)

    if new_torrent:
        if req.event == "started":
            start_new_session(torrent)
            upload_actual, upload_speed_actual = zero_upload(torrent, req)
        else:
            # report original
            upload_actual, upload_speed_actual = original_upload(torrent, req)
    else:
        if req.event == "started":
            start_new_session(torrent)
            upload_actual, upload_speed_actual = zero_upload(torrent, req)
        elif req.event == "completed":
            # zero-upload announce
            upload_actual, upload_speed_actual = original_upload(torrent, req)
        else:
            if torrent.is_seeding:
                if config.force_mod_upload: # 无论leechers数量, 强制修改上传
                    upload_actual, upload_speed_actual = mod_upload(torrent, req, config)

                elif torrent.leechers <= 1:
                    upload_actual, upload_speed_actual = zero_delta_upload(torrent, req)
                elif 1 < torrent.leechers < config.min_leachers_to_mod:
                    # reasonable speed
                    upload_actual, upload_speed_actual = min_speed_upload(torrent, req, int(config.uspeed_min))
                else:
                    upload_actual, upload_speed_actual = mod_upload(torrent, req, config)

            else:
                upload_actual, upload_speed_actual = zero_delta_upload(torrent, req)

    if req.event == "stopped":
        stop_session(torrent)

    new_announce = create_announce(torrent, req, upload_actual, upload_speed_actual)

    modified_uri = request_uri
    modified_uri = modify_uri_with_upload_actual(modified_uri, upload_actual)
    modified_uri = modify_uri_with_supported_peer_id(modified_uri)

    modified_headers = modify_headers(headers)

    update_stats_post_mod(torrent, req, new_announce)

    return modified_uri, modified_headers, torrent.id


def zero_upload(torrent: Torrent, req: TrackerRequest) -> int:
    return 0, 0


def zero_delta_upload(torrent: Torrent, req: TrackerRequest) -> int:
    return torrent.session_upload_actual, 0


def original_upload(torrent: Torrent, req: TrackerRequest) -> int:
    upload_actual = req.uploaded
    elapsed_since_last_announce = torrent.elapsed_since_last_announce()
    upload_speed = upload_actual / elapsed_since_last_announce
    return req.uploaded, upload_speed


def mod_upload(torrent: Torrent, req: TrackerRequest, config: SiteConfig) -> int:
    upload_speed = generate_exponential_speed(config.uspeed_min, config.uspeed_max, \
                                              torrent.leechers, config.max_anticipate_leechers, \
                                              config.min_leechers_factor, config.scale_factor)
    upload_speed = upload_speed * (1 - random() * config.upload_reduction)  # 速度减免
    upload_actual = calc_upload_amount_based_on_speed(upload_speed, torrent)

    return upload_actual, upload_speed

def min_speed_upload(torrent: Torrent, req: TrackerRequest, min_up_speed) -> int:
    upload_speed = get_random_ratio(0, min_up_speed)
    upload_actual = calc_upload_amount_based_on_speed(upload_speed, torrent)

    return upload_actual, upload_speed


def calc_upload_amount_based_on_speed(speed: int, torrent: Torrent) -> int:
    elapsed_since_last_announce = torrent.elapsed_since_last_announce()
    upload_delta = speed * elapsed_since_last_announce
    upload_actual = torrent.session_upload_actual + upload_delta

    return upload_actual


@db.connection_context()
def create_announce(torrent: Torrent, req: TrackerRequest, upload_actual: int, upload_speed_actual: int):

    new_announce = Announce.create(torrent=torrent,
                    event=req.event,
                    download=req.downloaded,
                    upload=req.uploaded,
                    # upload_delta=upload_delta,
                    upload_actual=upload_actual,
                    upload_speed=upload_speed_actual,
                    elapsed_since_last_announce=torrent.elapsed_since_last_announce(),
                    leechers=torrent.leechers,
                    )

    torrent.last_announced_at = datetime.now()
    torrent.save()

    return new_announce


def start_new_session(torrent: Torrent):
    torrent.session_upload = 0
    torrent.session_download = 0
    torrent.session_upload_actual = 0

    torrent.last_started_at = datetime.now()
    torrent.last_event = "started"

    torrent.save()


def stop_session(torrent: Torrent):
    torrent.total_download += torrent.session_download
    torrent.total_upload += torrent.session_upload
    torrent.total_upload_actual += torrent.session_upload_actual

    torrent.session_upload = 0
    torrent.session_download = 0
    torrent.session_upload_actual = 0

    torrent.last_event = "stopped"

    torrent.save()


def update_stats_pre_mod(torrent: Torrent, req: TrackerRequest):
    torrent.completed = req.left == 0 or req.event == 'completed'
    torrent.is_seeding = (req.left == 0) or req.event == 'paused'

    torrent.save()

def update_stats_post_mod(torrent: Torrent, req: TrackerRequest, new_announce: Announce):
    torrent.session_upload = req.uploaded
    torrent.session_download = req.downloaded
    torrent.session_upload_actual = new_announce.upload_actual

    torrent.save()

def modify_uri_with_upload_actual(request_uri:str, upload_actual:int) -> str:
    return re.sub(TrackerRequest.RE_UPLOAD, str(int(upload_actual)), request_uri)


def modify_uri_with_supported_peer_id(request_uri:str) -> str:
    # assume that the original client is qBittorrent
    return re.sub(r"qB\d{4}", "qB4670", request_uri)

def modify_headers(headers:dict) -> dict:
    new_headers = deepcopy(headers)
    new_headers["User-Agent"] = "qBittorrent/4.6.7"

    return new_headers


def start_server(host='0.0.0.0', port=9050, config={}):
    global SERVER
    logger.info(f'MRF proxy (socks5) is listening on {host}:{port}')
    SERVER = MRFServer((host, port), MRFHandler)
    SERVER.set_config(config)
    SERVER.serve_forever()


def shutdown_server():
    if SERVER:
        SERVER.shutdown()


# 暂时不考虑总体速度控制
# Consider the overall speed limit
# overall_speed = fcm.overall_speed()
# if overall_speed > MAX_OVERALL_UPLOAD_SPEED:
#     overspeed_by = (overall_speed - MAX_OVERALL_UPLOAD_SPEED) / MAX_OVERALL_UPLOAD_SPEED
#     # -> Reduce to just fit max speed
#     # reduce_ratio = max_allowed_speed / overall_speed

#     # -> Or, let's try to make more fluctuation by canceling this upload
#     reduce_ratio = 0

#     logger.info(f"[Overall] Overspeed by {convert_size(overspeed_by)}/s @{convert_size(overall_speed)}/s reduce to " + str(reduce_ratio))
#     session_uploaded_modified = session_uploaded_modified * reduce_ratio

# Make sure session upload always increase
# session_uploaded_modified = max(session_uploaded_modified, ctx.session_uploaded_modified)
# upload_speed_session = session_uploaded_modified / elapsed_since_start

# if upload_speed_session > MAX_SESSION_UPLOAD_SPEED:
#     logger.warn(f"[Session] Overspeed at {convert_size(upload_speed_session)}/s")


# NOTE: 暂时不做Session上传速度检测
# upload_speed_session =  session_upload_actual / session_elapsed

# if upload_speed_session > MAX_UPLOAD_SPEED:
#     logger.warn(f"[Session:{info_hash}] Overspeed at {convert_size(upload_speed_session)}/s")
#     # 恢复本次汇报前的数据
#     session_upload_actual = ctx.session_uploaded_actual
#     upload_speed_session =  ctx.session_uploaded_actual / session_elapsed

# TODO: 是否要做amount check?
# modified_upload = max(modified_upload, req.uploaded)
# logger.info(f"Mod stats: {info_hash} "
#             f"[THIS] {convert_size(upload_speed_this)}/s [SESSION] {convert_size(upload_speed_session)}/s "
#             f"[ELA_SINCE_LAST_ANN] {round(elapsed_since_last_announce)}s [ELA_SINCE_START] {round(elapsed_since_start)}s "
#             f"[THIS] {convert_size(this_uploaded_modified)} [SESSION] {convert_size(session_uploaded_modified)}")

# logger.info(f"[{info_hash}] Reporting {convert_size(session_uploaded_modified)}")