import inspect
import re
import base64
import hashlib
import struct
import time
from typing import Type, Union

from exception import CheckFailed, InvalidModel

# Add these constants at the top of the file
KONASUTE_CLIENT_SALT = b"d4BK3JFREkH5WuyTVEJQ2jbS9h2-df4D"
KONASUTE_SERVER_SALT = b"tGhCtgLuTjV7cZ2phWuCpQ8iwSypVn4W"
INFINITAS_CLIENT_SALT = b"fAwHp6G2FLPHN_ZGBhREJG5flt3hNu"
INFINITAS_SERVER_SALT = b"NH_P-urkCV9npxR90kaAR7YnqDTRL-"


def repair_base64(malformed_base64):
    # Replace Base64URL characters with standard Base64 characters
    base64_str = malformed_base64.replace("-", "+").replace("_", "/")
    # Remove invalid Base64 characters
    base64_str = re.sub(r"[^A-Za-z0-9+/=]", "", base64_str)

    # Add padding if necessary
    padding = len(base64_str) % 4
    if padding == 2:
        base64_str += "=="
    elif padding == 3:
        base64_str += "="
    # No action needed for padding == 0

    return base64_str


def repair_base64IIDX(malformed_base64):
    # Replace Base64URL characters with standard Base64 characters
    base64_str = malformed_base64.replace("-", "+").replace("_", "/")
    base64_str = base64_str.replace(" ", "+")
    # Remove invalid Base64 characters
    base64_str = re.sub(r"[^A-Za-z0-9+/=]", "", base64_str)

    # Add padding if necessary
    padding = len(base64_str) % 4
    if padding == 2:
        base64_str += "=="
    elif padding == 3:
        base64_str += "="
    # No action needed for padding == 0

    return base64_str


def assert_true(check: bool, reason: str, exc: Type[Exception] = CheckFailed):
    if not check:
        line = inspect.stack()[1].code_context
        if line:
            print()
            print("\n".join(line))
        raise exc(reason)


def py_encoding(name: str) -> str:
    if name.startswith("shift-jis"):
        return "shift-jis"
    return name


def parse_model(model: str) -> tuple[str, str, str, str, str]:
    # e.g. KFC:J:A:A:2019020600
    match = re.match(r"^([A-Z0-9]{3}):([A-Z]):([A-Z]):([A-Z])(?::(\d{10}))?$", model)
    if match is None:
        raise InvalidModel
    gamecode, dest, spec, rev, datecode = match.groups()
    return gamecode, dest, spec, rev, datecode


def pack(data, width: int) -> bytes:
    assert_true(1 <= width <= 8, "Invalid pack size")
    assert_true(all(i < (1 << width) for i in data), "Data too large for packing")
    bit_buf = in_buf = 0
    output = bytearray()
    for i in data:
        bit_buf |= i << (8 - width)
        shift = min(8 - in_buf, width)
        bit_buf <<= shift
        in_buf += shift
        if in_buf == 8:
            output.append(bit_buf >> 8)
            in_buf = width - shift
            bit_buf = (bit_buf & 0xFF) << in_buf

    if in_buf:
        output.append(bit_buf >> in_buf)

    return bytes(output)


def unpack(data, width: int) -> bytes:
    assert_true(1 <= width <= 8, "Invalid pack size")
    bit_buf = in_buf = 0
    output = bytearray()
    for i in data:
        bit_buf |= i
        bit_buf <<= width - in_buf
        in_buf += 8
        while in_buf >= width:
            output.append(bit_buf >> 8)
            in_buf -= width
            bit_buf = (bit_buf & 0xFF) << min(width, in_buf)

    if in_buf:
        output.append(bit_buf >> (8 + in_buf - width))

    return bytes(output)


def sign_with_salt(salt: Union[str, bytes], data: bytes, timestamp: float = None) -> bytes:
    """
    Sign data with a salt using SHA256.
    
    Args:
        salt: The salt to use for signing
        data: The data to sign
        timestamp: Optional timestamp in milliseconds (defaults to current time)
    
    Returns:
        bytes: The SHA256 signature
    """
    if isinstance(salt, str):
        salt = salt.encode('utf-8')
    
    if timestamp is None:
        timestamp = time.time() * 1000  # Convert to milliseconds
    
    # Convert timestamp to minutes (from milliseconds)
    timestamp_minutes = int(timestamp / (60 * 1000))
    packed_time = struct.pack('>Q', timestamp_minutes)
    
    # Create hash
    hash_obj = hashlib.sha256()
    hash_obj.update(data)
    hash_obj.update(packed_time)
    hash_obj.update(salt)
    
    return hash_obj.digest()


def find_timestamp_from_signature(known_salt: bytes, known_data: bytes, target_signature: bytes, 
                                time_range_start: int, time_range_end: int) -> int:
    """
    Try to find the timestamp used in a signature by brute force.
    
    Args:
        known_salt: The salt used for signing
        known_data: The data that was signed
        target_signature: The signature we're trying to match
        time_range_start: Start of timestamp range to try (in milliseconds)
        time_range_end: End of timestamp range to try (in milliseconds)
    
    Returns:
        int: The timestamp that produced the matching signature, or None if not found
    """
    for test_time in range(time_range_start, time_range_end, 1000):  # Step by 1 second
        test_signature = sign_with_salt(known_salt, known_data, timestamp=test_time)
        if test_signature == target_signature:
            return test_time
    return None


__all__ = ("assert_true", "py_encoding", "parse_model", "pack", "unpack", "sign_with_salt")