import random
import re
from collections.abc import Callable
from dataclasses import dataclass

from synthegrator.code_problems import CodeProblem
from synthegrator.transformation_spec import markup_path


@dataclass
class FewShotLibrary:
    """
    A collections of problems to be used for in-context few-shotting.
    This is a basic implementation that randomly samples
    """

    examples: list[CodeProblem]

    def sample_problems(
        self,
        n: int,
        target_problem: CodeProblem,
        seed: int | None = None,
    ) -> list[CodeProblem]:
        """Sample n problems from the library, excluding the target problem."""
        if seed is None:
            seed = random.randint(0, 2**32 - 1)
        random.seed(seed)
        if n < 0:
            msg = f"n must be non-negative, got {n}"
            raise ValueError(msg)
        if n == 0:
            return []
        n_with_over = min(n + 1, len(self.examples))
        samples: list[CodeProblem] = random.sample(self.examples, n_with_over)
        samples = remove_problem(samples, target_problem)
        return samples[: min(n, len(samples))]


def default_bm25_tokenizer(text: str) -> list[str]:
    # Split based on whitespace
    tokens = re.split(r"\s+", text)

    # Further split based on snake_case and camelCase
    refined_tokens = []
    for token in tokens:
        # Split snake_case
        snake_case_split = re.split(r"_", token)
        for sub_token in snake_case_split:
            # Split camelCase
            camel_case_split = re.findall(
                r"[A-Z]?[a-z]+|[A-Z]+(?=[A-Z][a-z]|$)", sub_token
            )
            if camel_case_split:
                refined_tokens.extend(camel_case_split)
            else:
                refined_tokens.append(sub_token)

    # Split off special symbols from tokens, keeping them as separate tokens
    final_tokens = []
    for token in refined_tokens:
        # Splitting special characters, keeping them as separate tokens
        special_char_split = re.findall(r"[^\w\s]|[\w]+", token)
        final_tokens.extend(special_char_split)

    # Filter out empty strings
    final_tokens = [token for token in final_tokens if token.strip() != ""]

    stop_tokens = {
        " ",
        "\n",
        "\t",
        "[",
        "]",
        "(",
        ")",
        "{",
        "}",
        "<",
        ">",
        "'",
        '"',
        "`",
        ":",
        ";",
        ",",
        ".",
        "?",
        "!",
        "-",
        "_",
        "=",
        "+",
        "*",
        "/",
        "\\",
        "|",
        "&",
        "^",
        "%",
        "$",
        "#",
        "@",
        "~",
    }
    final_tokens = [
        token.lower().strip() for token in final_tokens if token not in stop_tokens
    ]

    return final_tokens


def default_render_func(problem: CodeProblem) -> str:
    from synthegrator.problem_rendering import render_markup_to_text_autoregressive

    instructions = problem.instructions
    texts = [
        render_markup_to_text_autoregressive(
            markup_path(problem.working_directory, path, problem.transformation_spec),
            path,
            problem,
        ).code_prefix
        for path, file in problem.working_directory.files.walk()
    ]
    return "\n".join(texts)


class Bm25FewShotLibrary(FewShotLibrary):
    """
    A few-shot library that uses BM25 to select examples.
    """

    def __init__(
        self,
        examples: list[CodeProblem],
        render_func: Callable[[CodeProblem], str] = default_render_func,
        tokenizer: Callable[[str], list[str]] = default_bm25_tokenizer,
    ):
        self.examples = examples
        try:
            from rank_bm25 import BM25Okapi
        except ImportError:
            raise ImportError(
                "BM25FewShotLibrary requires the rank_bm25 package. "
                "You can install it with `pip install rank-bm25`."
            )
        self.tokenizer = tokenizer
        prompt_texts = [render_func(problem) for problem in examples]
        self.bm25 = BM25Okapi([self.tokenizer(text) for text in prompt_texts])
        self._render_func = render_func

    def sample_problems(
        self,
        n: int,
        target_problem: CodeProblem,
        seed: int | None = None,
    ) -> list[CodeProblem]:
        """Sample n problems from the library, excluding the target problem."""
        if seed is None:
            seed = random.randint(0, 2**32 - 1)
        random.seed(seed)
        if n < 0:
            msg = f"n must be non-negative, got {n}"
            raise ValueError(msg)
        if n == 0:
            return []
        n_with_over = min(n + 1, len(self.examples))
        query = self._render_func(target_problem)
        top_values = self.bm25.get_top_n(
            self.tokenizer(query), self.examples, n_with_over
        )
        top_values = remove_problem(top_values, target_problem)
        top_values = top_values[: min(n, len(top_values))]
        top_values.reverse()  # Most relevant last
        return top_values


def remove_problem(
    problems: list[CodeProblem],
    target_problem: CodeProblem,
) -> list[CodeProblem]:
    # Make sure the target is not included
    return [
        x
        for x in problems
        if (
            x is not target_problem
            # Compare a few key values to exclude if match. Might need
            #    to make more specific later
            and not (
                x.working_directory.files.get_hash()
                == target_problem.working_directory.files.get_hash()
                and x.instructions == target_problem.instructions
                and x.transformation_spec == target_problem.transformation_spec
            )
        )
    ]


@dataclass
class FewShotConfig:
    library: FewShotLibrary
    num_examples: int
    change_examples_between_solves: bool
    """When sampling multiple solutions from a solver,
    should the examples change between each solve?"""
