import difflib
import hashlib
import inspect
import string
from collections import UserString
from collections.abc import Callable, Iterable, Iterator, Sequence
from functools import wraps
from typing import SupportsIndex

import numpy as np


def ensure_newline_at_end(s: str) -> str:
    if s.endswith("\n"):
        return s
    return s + "\n"


def assert_str_eq_with_diff_error(a: str, b: str):
    if a != b:
        # hacky get the two expressions of the caller line
        # https://stackoverflow.com/a/50636091/382912
        import inspect

        caller_frame = inspect.stack()[1]
        caller_line = caller_frame.code_context[0].strip()
        caller_line = caller_line.split("assert_str_eq_with_diff_error(")[1]
        caller_line = caller_line.split(")")[0]
        a_name, b_name = caller_line.split(", ")
        diff = difflib.unified_diff(
            a.splitlines(),
            b.splitlines(),
            fromfile=a_name,
            tofile=b_name,
        )
        diff_str = "\n".join(diff)
        msg = f"Strings not equal:\n{diff_str}"
        raise AssertionError(msg)


def pretty_print_python_code(code: str):
    """Pretty prints python code with syntax highlighting."""
    import pygments
    from pygments.formatters import TerminalFormatter
    from pygments.lexers import PythonLexer

    print(pygments.highlight(code, PythonLexer(), TerminalFormatter()))


def n_lines(code: str, n: int) -> str:
    lines = code.split("\n")
    return "\n".join(lines[: min(n, len(lines))])


def clip(value, min_value, max_value):
    if value < min_value:
        return min_value
    if value > max_value:
        return max_value
    return value


class MultilineString(UserString):
    def __init__(self, string: str, expect_newline: bool = False) -> None:
        super().__init__(string)
        self.lines = string.splitlines(keepends=True)
        self.expect_newline = expect_newline

    def _new_line_guard(self, value: str) -> str:
        if self.expect_newline and not value.endswith("\n"):
            msg = "All strings must end in a newline!"
            raise Exception(msg)

        if not value.endswith("\n"):
            return value + "\n"
        return value

    @property
    def _string(self) -> str:
        return "".join(self.lines)

    def __str__(self) -> str:
        return self._string

    def __getitem__(self, index: SupportsIndex | slice) -> str:
        return "".join(self.lines[index])

    def __setitem__(self, index: SupportsIndex | slice, value: str) -> str:
        if not isinstance(value, str):
            msg = "Currently only a string is a valid value."
            raise TypeError(msg)
        value = self._new_line_guard(value)
        self.lines[index] = value
        self.data = self._string
        return self

    def __delitem__(self, index: SupportsIndex | slice) -> None:
        del self.lines[index]
        self.data = self._string

    def __len__(self) -> int:
        return len(self.lines)

    def insert(self, index: SupportsIndex | slice, lines: str | list) -> str:
        if isinstance(lines, str):
            lines = self._new_line_guard(lines)
            lines = lines.splitlines(keepends=True)

        self.lines[index:index] = lines
        self.data = self._string
        return self

    @property
    def line_count(self) -> int:
        return len(self.lines)


def lazy_prop(func):
    val = None

    def inner():
        nonlocal val
        if val is None:
            val = func()
        return val

    return inner


def estimate_best_overlap_span(
    larger_str: str,
    larger_tokens: list[str] | None,
    mod_string: str,
    tokenizer_func: Callable[[str], list[str]],
) -> tuple[int, int]:
    """
    Returns a slice into a sequence of larger tokens that might best overlap
    with smaller tokens. This is designed for when the smaller string is a
    chopped and slightly modified (such as adding preceeding indent) version
    """
    if larger_tokens is None:
        larger_tokens = tokenizer_func(larger_str)
    if larger_str == mod_string:
        return 0, len(larger_tokens)
    smaller_tokens = tokenizer_func(mod_string)
    large_t = np.array(larger_tokens)
    smaller_tokens = np.array(smaller_tokens)
    # TODO: I think an algorithm that looks like making a edit distance
    #   dp matrix first, then finding the offset as a path through that
    #   can work.
    for start_point in range(len(large_t) - len(smaller_tokens)):
        end_point = start_point + len(smaller_tokens)
        if np.array_equal(large_t[start_point:end_point], smaller_tokens):
            return start_point, end_point
    # For now just give up and return from beginning
    return 0, len(larger_tokens)


def frozen_method_cache(func: Callable) -> Callable:
    """
    A decorator for methods. It is expected this only on non-mutable clases.
    The `self` is ignored in the cache. If you use it on mutable classes that
    depend on the class state, then it might get unexpected results.
    """
    uncached_sentinel = object()  # Unique. Allows a return of None to be cached
    cache = uncached_sentinel

    if len(inspect.signature(func).parameters) == 1:
        uncached_sentinel = object()  # Unique. Allows a return of None to be cached
        cache = uncached_sentinel

        @wraps(func)
        def wrapper(self):
            nonlocal cache
            if cache is uncached_sentinel:
                cache = func(self)
            return cache

    else:
        cache = {}

        @wraps(func)
        def wrapper(self, *args, **kwargs):
            key = (args, tuple(kwargs.items()))
            if key not in cache:
                cache[key] = func(self, *args, **kwargs)
            return cache[key]

    return wrapper


def get_first_line_with_nonwhitespace(value: str) -> str:
    try:
        first_line = next(
            line for line in value.splitlines(keepends=True) if line.strip() != ""
        )
    except StopIteration:
        first_line = value
    return first_line


from typing import TypeVar

T = TypeVar("T")


class IteratorWithLength(Iterator[T]):
    def __init__(self, iterable: Iterator[T], length: int | None) -> None:
        self.iterable = iterable
        self.length = length
        self.num_yielded = 0

    def __iter__(self) -> Iterator[T]:
        return self

    def __next__(self) -> T:
        if self.length is not None and self.num_yielded >= self.length:
            raise StopIteration
        self.num_yielded += 1
        return next(self.iterable)

    def __len__(self) -> int | None:
        if self.length is None:
            return None
        return self.length - self.num_yielded


# Characters used for Base62 encoding
BASE62 = string.digits + string.ascii_letters


def base62_encode(num):
    if num == 0:
        return BASE62[0]
    base62 = []
    while num:
        num, rem = divmod(num, 62)
        base62.append(BASE62[rem])
    return "".join(reversed(base62))


def hash_string_base62(
    s: str,
    digest_size_bytes: int = 8,
) -> str:
    # Using blake2b for a faster, non-cryptographic hash.
    # Adjust the digest_size if you want a shorter output.
    hash_bytes = hashlib.blake2b(s.encode(), digest_size=digest_size_bytes).digest()
    # Convert hash bytes to an integer
    num = int.from_bytes(hash_bytes, byteorder="big")
    # Encode this number in base62
    return base62_encode(num)


def find_after_first_occurrence_of_string(
    string_parts: list[str],
    split_str: str,
) -> tuple[int, int] | tuple[None, None]:
    """
    We are given `string_parts` list, which represents a
    string that is a concatenation of the strings in the list.
    That is, if you were to join all the parts, you would get one
    string. We are searching for the index that begins a string
    after the `split_str`. Because it might split cleanly within
    string, we also return the index within the string part itself
    which begins after the split str

    If `search_str` is not found, returns (None, None)

    Examples
    --------
        >> find_after_first_occurrence_of_string(
            ["My", " Name", " Is"],
            "Name"
        ) == (2, 0)
        # The split string might span between parts
        >> find_after_first_occurrence_of_string(
            ["My", " Name", " Is", " Bob"],
            "Name Is"
        ) == (3, 0)
        # An example where doesn't cleanly split, and have
        #   to within one of the parts
        >> find_after_first_occurrence_of_string(
            ["My", " Name", " Is Bob"],
            "Name Is"
        ) == (2j, 3)
        # When it is not found, we return (None, None)
        >> find_after_first_occurrence_of_string(
            ["My", " Name", " Is Bob"],
            "Random"
        ) == (None, None)

    """
    if len(string_parts) == 0:
        return (None, None)

    full_string = "".join(string_parts)  # Concatenate all parts to form the full string
    split_index = full_string.find(
        split_str,
    )  # Find the index of split_str in the full string

    if split_index == -1:
        # If split_str is not found, return (None, None)
        return (None, None)

    # Adjust split_index to point to the character after split_str
    split_index += len(split_str)

    # Now, find which part and the index within that part corresponds to split_index
    current_index = 0
    for i, part in enumerate(string_parts):
        # If the split_index is within the current part
        if current_index + len(part) >= split_index:
            # Calculate the index within the current part
            index_within_part = split_index - current_index
            if index_within_part == len(part):
                # If the index within the part is the length of the part,
                #   it means the split_str is at the end of the part, and
                #   the next part is the one that starts after the split_str
                return (i + 1, 0)
            return (i, index_within_part)
        current_index += len(part)

    msg = "Expected to have returned"
    raise RuntimeError(msg)


def get_only_item(values: Sequence[T] | Iterable[T]) -> T:
    if not isinstance(values, Sequence):
        values = list(values)
    if len(values) != 1:
        msg = f"Expected exactly one item, got {len(values)}"
        raise ValueError(msg)
    return values[0]


def normalize_probs_dict(
    answer_to_prob: dict[str, float],
    available_answers: Sequence[str],
) -> dict[str, float]:
    """
    Makes a new dict that will sum to 1.0. If there are missing
    answers, distributes any available probability mass to missing answers.
    If all the answers are present, then distributes the mass
    in proportion to the current probabilities.
    """
    # Calculate the total probability of the given answers
    total_prob = sum(answer_to_prob.values())

    # Calculate the number of missing answers
    missing_answers = [
        answer for answer in available_answers if answer not in answer_to_prob
    ]
    num_missing_answers = len(missing_answers)
    if 1.0 - 1e-6 <= total_prob <= 1.0 + 1e-6 and num_missing_answers == 0:
        return answer_to_prob

    # Initialize a new dictionary to store the normalized probabilities
    normalized_probs = {}

    # If there are missing answers, distribute the remaining probability mass equally among them
    if num_missing_answers > 0:
        # Calculate the remaining probability mass to distribute
        remaining_prob = 1.0 - total_prob
        # Distribute the remaining probability equally among the missing answers
        for answer in available_answers:
            if answer in answer_to_prob:
                # If the answer is not missing, use its original probability
                normalized_probs[answer] = answer_to_prob[answer]
            else:
                # If the answer is missing, assign it an equal share of the remaining probability
                normalized_probs[answer] = remaining_prob / num_missing_answers
    else:
        # If there are no missing answers, normalize the probabilities so they sum to 1.0
        for answer in available_answers:
            normalized_probs[answer] = answer_to_prob[answer] / total_prob

    # Return in a standard order
    return {answer: normalized_probs[answer] for answer in sorted(available_answers)}


if __name__ == "__main__":
    print("hello world")