import functools
import hashlib
from .logging_util import logger
from types import ModuleType, FunctionType
import inspect
from typing import Iterable
import ast
import io
import tokenize
import re
import textwrap

import numpy as np

# This variable is used within the checksum function to salt the sha1 sum.
# May be changed to force a different checksum for similar objects.
SALT = 42


def version(version_nr: int, calls: Iterable = ()):
    """Function decorator that assigns a custom checksum to a function."""

    def decorator(func):
        cs = checksum(func.__name__, version_nr, *calls)
        func.__checksum__ = lambda: cs

        @functools.wraps(func)
        def wrapped(*args, **kwargs):
            return func(*args, **kwargs)

        return wrapped

    return decorator


def strip_comments(source: str) -> str:
    """Removes docstrings, comments, and irrelevant whitespace from Python source code."""

    # Step 1: Remove docstrings using AST
    def remove_docstrings(node):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Module)):
            if (doc := ast.get_docstring(node, clean=False)):
                first_stmt = node.body[0]
                if isinstance(first_stmt, ast.Expr) and isinstance(first_stmt.value, ast.Constant):
                    node.body.pop(0)  # Remove the docstring entirely
        for child in ast.iter_child_nodes(node):
            remove_docstrings(child)

    tree = ast.parse(textwrap.dedent(source))
    remove_docstrings(tree)
    code_without_docstrings = ast.unparse(tree)

    # Step 2: Remove comments using tokenize
    tokens = tokenize.generate_tokens(io.StringIO(code_without_docstrings).readline)
    result = []
    last_lineno = -1
    last_col = 0

    for toknum, tokval, (srow, scol), (erow, ecol), line in tokens:
        if toknum == tokenize.COMMENT:
            continue
        if srow > last_lineno:
            last_col = 0
        if scol > last_col:
            result.append(" " * (scol - last_col))
        result.append(tokval)
        last_lineno, last_col = erow, ecol

    code_no_comments = ''.join(result)

    # Step 3: Remove empty lines (whitespace-only or truly blank)
    return "\n".join([line for line in code_no_comments.splitlines() if line.strip() != ""])


def checksum(*args, csum=None, _seen=None):
    """
    Calculate a checksum of any object, by sha1 hash.

    Inputs for the hash are some salt bytes and the byte encoding of a string
    that depends on the object and its type:

    - If a method __checksum__ is available, its return value is converted to bytes
    - str or bytes are used as sha1 input directly
    - modules use the __name__ attribute
    - functions use the function code and any closures of the function
    - functools.partial uses the checksum of the function and any arguments, that were
        defined
    - numpy.ndarray uses bytes representation of the array (arr.tobytes())
    - Anything else is converted to a str
    """
    if csum is None:
        csum = hashlib.sha1()
        csum.update(str(SALT).encode())

    if _seen is None:
        _seen = set()

    for arg in args:
        obj_id = id(arg)
        if obj_id in _seen:
            continue
        _seen.add(obj_id)

        if hasattr(arg, "__checksum__"):
            method = getattr(arg, "__checksum__")
            if callable(method) and not isinstance(arg, type):
                logger.debug("Checksum via __checksum__: %s", str(arg))
                csum.update(str(method()).encode())
            elif isinstance(arg, type):
                try:
                    src = inspect.getsource(arg)
                    csum.update(strip_comments(src).encode())
                    logger.debug("Checksum via class source for %s", arg.__name__)
                except (OSError, TypeError):
                    csum.update(arg.__name__.encode())
                    logger.debug("Checksum via class name for %s", arg.__name__)
            else:
                logger.debug("Skipping unbound __checksum__ on %s", type(arg))
        elif isinstance(arg, bytes):
            csum.update(arg)
        elif isinstance(arg, str):
            csum.update(arg.encode())
        elif isinstance(arg, ModuleType):
            csum.update(arg.__name__.encode())
        elif isinstance(arg, FunctionType):
            csum.update(strip_comments(inspect.getsource(arg)).encode())
            c = inspect.getclosurevars(arg)
            merged = {**c.nonlocals, **c.globals}
            for key in sorted(merged):  # deterministic ordering
                v = merged[key]
                if v is not arg:
                    checksum(v, csum=csum, _seen=_seen)
        elif isinstance(arg, functools.partial):
            logger.debug("Checksum via partial for %s", str(arg))
            checksum(arg.func, csum=csum, _seen=_seen)
            for x in arg.args:
                checksum(x, csum=csum, _seen=_seen)
            for k in sorted(arg.keywords.keys()):
                csum.update(k.encode())
                checksum(arg.keywords[k], csum=csum, _seen=_seen)
        elif isinstance(arg, np.ndarray):
            csum.update(arg.tobytes())
        else:
            logger.debug("Checksum via str for %s", str(arg))
            csum.update(str(arg).encode())

    return int.from_bytes(csum.digest(), "big")
