import random
import time
import base64
import hashlib
from typing import Dict, List, Any, Optional
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
import json

# =========================
# DEFAULT TEMPLATE
# =========================

DEFAULT_TEMPLATE: Dict[str, Any] = {
    "deviceId": "84985177a19a010dea49",
    "sdkVersion": "websdk-2.3.15d",
    "initTimestamp": "1765348410850",
    "field3": "91",
    "field4": "1|15",
    "language": "zh-CN",
    "timezoneOffset": "-480",
    "colorDepth": "16705151|12791",
    "screenInfo": "1470|956|283|797|158|0|1470|956|1470|798|0|0",
    "field9": "5",
    "platform": "MacIntel",
    "field11": "10",
    "webglRenderer": (
        "ANGLE (Apple, ANGLE Metal Renderer: Apple M4, Unspecified Version)"
        "|Google Inc. (Apple)"
    ),
    "field13": "30|30",
    "field14": "0",
    "field15": "28",
    "pluginCount": "5",
    "vendor": "Google Inc.",
    "field29": "8",
    "touchInfo": "-1|0|0|0|0",
    "field32": "11",
    "field35": "0",
    "mode": "P",
}


# =========================
# BX-UA GENERATOR
# =========================


class BXUAGenerator:
    def __init__(self):
        self.version = "231"
        self.aes_key = None
        self.aes_iv = None

    def _generate_key_iv(self, seed_data: str) -> tuple:
        """Generate AES key and IV from seed data"""
        # Create deterministic key/IV from seed
        seed_hash = hashlib.sha256(seed_data.encode()).digest()
        key = seed_hash[:16]  # 128-bit key
        iv = seed_hash[16:32]  # 128-bit IV
        return key, iv

    def _encrypt_aes_cbc(self, data: bytes, key: bytes, iv: bytes) -> bytes:
        """Encrypt data using AES-CBC"""
        cipher = AES.new(key, AES.MODE_CBC, iv)
        padded_data = pad(data, AES.block_size)
        encrypted = cipher.encrypt(padded_data)
        return encrypted

    def _create_payload(
        self, fingerprint: str, timestamp: Optional[int] = None
    ) -> Dict[str, Any]:
        """Create the payload structure to be encrypted"""
        if timestamp is None:
            timestamp = int(time.time() * 1000)

        # Extract components from fingerprint
        fields = fingerprint.split("^")

        payload = {
            "v": self.version,
            "ts": timestamp,
            "fp": fingerprint,
            "d": {
                "deviceId": fields[0],
                "sdkVer": fields[1],
                "lang": fields[5],
                "tz": fields[6],
                "platform": fields[10],
                "renderer": fields[12],
                "mode": fields[23],
                "vendor": fields[28],
            },
            "rnd": random.randint(1000, 9999),
            "seq": 1,
        }

        # Add checksum
        checksum_str = f"{fingerprint}{timestamp}{payload['rnd']}"
        payload["cs"] = hashlib.md5(checksum_str.encode()).hexdigest()[:8]

        return payload

    def generate(
        self, fingerprint: str, options: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Generate bx-ua header value

        Args:
            fingerprint: The fingerprint string generated by generate_fingerprint()
            options: Optional configuration
                - timestamp: Custom timestamp (default: current time)
                - seed: Custom seed for key generation
        """
        if options is None:
            options = {}

        # Get timestamp
        timestamp = options.get("timestamp")
        if timestamp is None:
            timestamp = int(time.time() * 1000)

        # Create payload
        payload = self._create_payload(fingerprint, timestamp)

        # Convert to JSON
        payload_json = json.dumps(payload, separators=(",", ":"))

        # Generate key and IV
        seed = options.get("seed", fingerprint)
        key, iv = self._generate_key_iv(seed)

        # Encrypt
        encrypted = self._encrypt_aes_cbc(payload_json.encode(), key, iv)

        # Base64 encode
        encrypted_b64 = base64.b64encode(encrypted).decode()

        # Return in format: version!base64_encoded_data
        return f"{self.version}!{encrypted_b64}"

    def batch_generate(
        self, fingerprints: List[str], options: Optional[Dict[str, Any]] = None
    ) -> List[str]:
        """Generate multiple bx-ua values"""
        return [self.generate(fp, options) for fp in fingerprints]


# =========================
# FINGERPRINT GENERATOR (From your code)
# =========================


def generate_device_id() -> str:
    """Generate a 20-character hex device ID"""
    return "".join(random.choice("0123456789abcdef") for _ in range(20))


def generate_hash() -> int:
    """Generate a 32-bit unsigned random hash"""
    return random.randint(0, 0xFFFFFFFF)


def generate_fingerprint(options: Dict[str, Any] = None) -> str:
    if options is None:
        options = {}

    config = DEFAULT_TEMPLATE.copy()

    # platform preset
    platform = options.get("platform")
    if platform:
        # Handle platform presets if needed
        pass

    # screen preset
    screen = options.get("screen")
    if screen:
        # Handle screen presets if needed
        pass

    # language preset
    locale = options.get("locale")
    if locale:
        # Handle language presets if needed
        pass

    # custom overrides
    if "custom" in options and isinstance(options["custom"], dict):
        config.update(options["custom"])

    device_id = options.get("deviceId") or generate_device_id()
    current_timestamp = int(time.time() * 1000)

    plugin_hash = generate_hash()
    canvas_hash = generate_hash()
    ua_hash1 = generate_hash()
    ua_hash2 = generate_hash()
    url_hash = generate_hash()
    doc_hash = random.randint(10, 100)

    fields: List[Any] = [
        device_id,
        config["sdkVersion"],
        config["initTimestamp"],
        config["field3"],
        config["field4"],
        config["language"],
        config["timezoneOffset"],
        config["colorDepth"],
        config["screenInfo"],
        config["field9"],
        config["platform"],
        config["field11"],
        config["webglRenderer"],
        config["field13"],
        config["field14"],
        config["field15"],
        f'{config["pluginCount"]}|{plugin_hash}',
        canvas_hash,
        ua_hash1,
        "1",
        "0",
        "1",
        "0",
        config["mode"],
        "0",
        "0",
        "0",
        "416",
        config["vendor"],
        config["field29"],
        config["touchInfo"],
        ua_hash2,
        config["field32"],
        current_timestamp,
        url_hash,
        config["field35"],
        doc_hash,
    ]

    return "^".join(map(str, fields))


# =========================
# USAGE EXAMPLE
# =========================


def example_usage():
    """Example of how to generate bx-ua headers"""

    # Initialize generators
    fp_gen = BXUAGenerator()

    # Generate a fingerprint
    fingerprint = generate_fingerprint(
        {
            "deviceId": "84985177a19a010dea49",
            "custom": {
                "language": "zh-CN",
                "platform": "MacIntel",
            },
        }
    )

    print("Generated Fingerprint:")
    print(fingerprint[:100] + "...")
    print()

    # Generate bx-ua header
    bx_ua = fp_gen.generate(
        fingerprint,
        {
            "timestamp": int(time.time() * 1000),
            "seed": "consistent_seed_for_deterministic_output",
        },
    )
    print(bx_ua)
    print("Generated bx-ua header:")
    print(bx_ua[:100] + "...")
    print(f"Total length: {len(bx_ua)}")
    print()

    # Parse the fingerprint (for debugging)
    fields = fingerprint.split("^")
    parsed = {
        "deviceId": fields[0],
        "sdkVersion": fields[1],
        "language": fields[5],
        "timezoneOffset": fields[6],
        "platform": fields[10],
        "webglRenderer": fields[12],
        "mode": fields[23],
        "vendor": fields[28],
        "timestamp": fields[33],
    }

    print("Parsed fingerprint info:")
    for key, value in parsed.items():
        print(f"  {key}: {value}")

    return bx_ua


def batch_example():
    """Example of batch generation"""
    fp_gen = BXUAGenerator()

    # Generate multiple fingerprints
    fingerprints = [
        generate_fingerprint({"deviceId": generate_device_id()}) for _ in range(3)
    ]

    # Generate bx-ua for each
    bx_ua_values = fp_gen.batch_generate(fingerprints, {"seed": "batch_seed"})

    print("Batch Generation Results:")
    for i, (fp, bx_ua) in enumerate(zip(fingerprints, bx_ua_values)):
        print(f"\n{i + 1}. Device ID: {fp.split('^')[0]}")
        print(f"   bx-ua: {bx_ua[:80]}...")

    return bx_ua_values


if __name__ == "__main__":
    print("=" * 60)
    print("BX-UA Header Generator")
    print("=" * 60)

    # Run single example
    print("\n1. Single Generation Example:")
    print("-" * 40)
    example_ua = example_usage()

    print("\n" + "=" * 60)

    # Run batch example
    print("\n2. Batch Generation Example:")
    print("-" * 40)
    batch_ua = batch_example()

    print("\n" + "=" * 60)
    print("\nTo use in requests:")
    print("```python")
    print("import requests")
    print("")
    print("headers = {")
    print('    "bx-ua": f"{example_ua}",')
    print('    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) ...",')
    print('    "Accept": "application/json, text/plain, */*",')
    print('    "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",')
    print('    "Accept-Encoding": "gzip, deflate, br",')
    print('    "Connection": "keep-alive",')
    print("}")
    print("")
    print('response = requests.get("https://example.com/api", headers=headers)')
    print("```")