import ast
import glob
import re
from os.path import join as pjoin
from pathlib import Path
from rank_bm25 import BM25Okapi
from nltk.tokenize import word_tokenize
import textwrap
import tree_sitter_python as tspython
from tree_sitter import Language, Parser
# import nltk
# nltk.download('punkt_tab') # do it in advance

def is_test_file(file_path: str) -> bool:
    """Check if a file is a test file.

    This is a simple heuristic to check if a file is a test file.
    """
    return (
        "test" in Path(file_path).parts
        or "tests" in Path(file_path).parts
        or file_path.endswith("_test.py")
    )


def find_python_files(dir_path: str, include_test_files=True) -> list[str]:
    """Get all .py files recursively from a directory.

    Skips files that are obviously not from the source code, such third-party library code.

    Args:
        dir_path (str): Path to the directory.
        include_test_files (bool): Whether to include test files during search.
    Returns:
        List[str]: List of .py file paths. These paths are ABSOLUTE path!
    """

    py_files = glob.glob(pjoin(dir_path, "**/*.py"), recursive=True)
    res = []
    for file in py_files:
        rel_path = file[len(dir_path) + 1 :]
        if not include_test_files and is_test_file(rel_path): # exclude test cases if include_test_files = False
            continue
        res.append(file)
    return res


def parse_class_def_args(source: str, node: ast.ClassDef) -> list[str]:
    # TODO this is simple enough to cover a lot of cases but can be improvied
    super_classes = []
    for base in node.bases:
        if isinstance(base, ast.Name):
            if base.id in ["type", "object"]:
                continue
            super_classes.append(ast.get_source_segment(source, base))
        if (
            isinstance(base, ast.Call)
            and ast.get_source_segment(source, base.func) == "type"
        ):
            super_classes.append(ast.get_source_segment(source, base.args[0]))
    return super_classes


def parse_python_file(
    file_full_path: str,
) -> (
    tuple[
        list[tuple[str, int, int]],
        dict[str, list[tuple[str, int, int]]],
        list[tuple[str, int, int]],
        dict[tuple[str, int, int], list[str]],
    ]
    | None
):
    """
    Main method to parse AST and build search index.
    Handles complication where python ast module cannot parse a file.
    """
    try:
        file_content = Path(file_full_path).read_text()
        tree = ast.parse(file_content)
    except Exception:
        # failed to read/parse one file, we should ignore it
        return None

    # (1) get all classes defined in the file
    classes: list[tuple[str, int, int]] = []
    # (2) for each class in the file, get all functions defined in the class.
    class_to_funcs: dict[str, list[tuple[str, int, int]]] = {}
    # (3) get top-level functions in the file (exclues functions defined in classes)
    top_level_funcs: list[tuple[str, int, int]] = []
    # (4) get class relations
    class_relation_map: dict[tuple[str, int, int], list[str]] = {}

    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            ## class part (1): collect class info
            class_name = node.name
            start_lineno = node.lineno
            end_lineno = node.end_lineno
            assert end_lineno is not None, "class should have end_lineno in AST."
            # line numbers are 1-based
            classes.append((class_name, start_lineno, end_lineno))
            class_relation_map[(class_name, start_lineno, end_lineno)] = (
                parse_class_def_args(file_content, node)
            )

            ## class part (2): collect function info inside this class
            class_funcs = [
                (n.name, n.lineno, n.end_lineno)
                for n in ast.walk(node)
                if isinstance(n, ast.FunctionDef) and n.end_lineno is not None
            ]
            class_to_funcs[class_name] = class_funcs

        elif isinstance(node, ast.FunctionDef):
            function_name = node.name
            start_lineno = node.lineno
            end_lineno = node.end_lineno
            assert end_lineno is not None, "function should have end_lineno in AST."
            # line numbers are 1-based
            top_level_funcs.append((function_name, start_lineno, end_lineno))

    return classes, class_to_funcs, top_level_funcs, class_relation_map


def get_code_region_containing_code(
    file_full_path: str, code_str: str, with_lineno=True
) -> list[tuple[int, str]]:
    """In a file, get the region of code that contains a specific string.

    Args:
        - file_full_path: Path to the file. (absolute path)
        - code_str: The string that the function should contain.
    Returns:
        - A list of tuple, each of them is a pair of (line_no, code_snippet).
        line_no is the starting line of the matched code; code snippet is the
        source code of the searched region.
    """
    with open(file_full_path) as f:
        file_content = f.read()

    context_size = 3
    # since the code_str may contain multiple lines, let's not split the source file.

    # we want a few lines before and after the matched string. Since the matched string
    # can also contain new lines, this is a bit trickier.
    pattern = re.compile(re.escape(code_str))
    # each occurrence is a tuple of (line_no, code_snippet) (1-based line number)
    occurrences: list[tuple[int, str]] = []
    for match in pattern.finditer(file_content):
        matched_start_pos = match.start()
        # first, find the line number of the matched start position (0-based)
        matched_line_no = file_content.count("\n", 0, matched_start_pos)

        file_content_lines = file_content.splitlines()

        window_start_index = max(0, matched_line_no - context_size)
        window_end_index = min(
            len(file_content_lines), matched_line_no + context_size + 1
        )

        if with_lineno:
            context = ""
            for i in range(window_start_index, window_end_index):
                context += f"{i+1} {file_content_lines[i]}\n"
        else:
            context = "\n".join(file_content_lines[window_start_index:window_end_index])
        occurrences.append((matched_line_no, context))

    return occurrences


def get_func_snippet_with_code_in_file(file_full_path: str, code_str: str) -> list[str]:
    """In a file, get the function code, for which the function contains a specific string.

    Args:
        file_full_path (str): Path to the file. (absolute path)
        code_str (str): The string that the function should contain.

    Returns:
        A list of code snippets, each of them is the source code of the searched function.
    """
    with open(file_full_path) as f:
        file_content = f.read()

    tree = ast.parse(file_content)
    all_snippets = []
    for node in ast.walk(tree):
        if not isinstance(node, ast.FunctionDef):
            continue
        func_start_lineno = node.lineno
        func_end_lineno = node.end_lineno
        assert func_end_lineno is not None
        func_code = get_code_snippets(
            file_full_path, func_start_lineno, func_end_lineno
        )
        # This func code is a raw concatenation of source lines which contains new lines and tabs.
        # For the purpose of searching, we remove all spaces and new lines in the code and the
        # search string, to avoid non-match due to difference in formatting.
        stripped_func = " ".join(func_code.split())
        stripped_code_str = " ".join(code_str.split())
        if stripped_code_str in stripped_func:
            all_snippets.append(func_code)

    return all_snippets


def get_code_snippets(
    file_full_path: str, start: int, end: int, with_lineno=True
) -> str:
    """Get the code snippet in the range in the file,with or without line numbers.

    Args:
        file_path (str): Full path to the file.
        start (int): Start line number. (1-based)
        end (int): End line number. (1-based)
    """
    with open(file_full_path) as f:
        file_content = f.readlines()
    snippet = ""
    for i in range(start - 1, end):
        if with_lineno:
            snippet += f"{i+1} {file_content[i]}"
        else:
            snippet += file_content[i]
    return snippet


def extract_func_sig_from_ast(func_ast: ast.FunctionDef) -> list[int]:
    """Extract the function signature from the AST node.

    Includes the decorators, method name, and parameters.

    Args:
        func_ast (ast.FunctionDef): AST of the function.

    Returns:
        The source line numbers that contains the function signature.
    """
    func_start_line = func_ast.lineno
    if func_ast.decorator_list:
        # has decorators
        decorator_start_lines = [d.lineno for d in func_ast.decorator_list]
        decorator_first_line = min(decorator_start_lines)
        func_start_line = min(decorator_first_line, func_start_line)
    # decide end line from body
    if func_ast.body:
        # has body
        body_start_line = func_ast.body[0].lineno
        end_line = body_start_line - 1
    else:
        # no body
        end_line = func_ast.end_lineno
    assert end_line is not None
    return list(range(func_start_line, end_line + 1))


def extract_class_sig_from_ast(class_ast: ast.ClassDef) -> list[int]:
    """Extract the class signature from the AST.

    Args:
        class_ast (ast.ClassDef): AST of the class.

    Returns:
        The source line numbers that contains the class signature.
    """
    # STEP (1): extract the class signature
    sig_start_line = class_ast.lineno
    if class_ast.body:
        # has body
        body_start_line = class_ast.body[0].lineno
        sig_end_line = body_start_line - 1
    else:
        # no body
        sig_end_line = class_ast.end_lineno
    assert sig_end_line is not None
    sig_lines = list(range(sig_start_line, sig_end_line + 1))

    # STEP (2): extract the function signatures and assign signatures
    for stmt in class_ast.body:
        if isinstance(stmt, ast.FunctionDef):
            sig_lines.extend(extract_func_sig_from_ast(stmt))
        elif isinstance(stmt, ast.Assign):
            # for Assign, skip some useless cases where the assignment is to create docs
            stmt_str_format = ast.dump(stmt)
            if "__doc__" in stmt_str_format:
                continue
            # otherwise, Assign is easy to handle
            assert stmt.end_lineno is not None
            assign_range = list(range(stmt.lineno, stmt.end_lineno + 1))
            sig_lines.extend(assign_range)

    return sig_lines


def get_class_signature(file_full_path: str, class_name: str) -> str:
    """Get the class signature.

    Args:
        file_path (str): Path to the file.
        class_name (str): Name of the class.
    """
    with open(file_full_path) as f:
        file_content = f.read()

    tree = ast.parse(file_content)
    relevant_lines = []
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef) and node.name == class_name:
            # we reached the target class
            relevant_lines = extract_class_sig_from_ast(node)
            break
    if not relevant_lines:
        return ""
    else:
        with open(file_full_path) as f:
            file_content = f.readlines()
        result = ""
        for line in relevant_lines:
            line_content: str = file_content[line - 1]
            if line_content.strip().startswith("#"):
                # this kind of comment could be left until this stage.
                # reason: # comments are not part of func body if they appear at beginning of func
                continue
            result += line_content
        return result


def get_code_region_around_line(
    file_full_path: str, line_no: int, window_size: int = 10, with_lineno=True
) -> str | None:
    """Get the code region around a specific line number in a file.

    Args:
        file_full_path (str): Path to the file. (absolute path)
        line_no (int): The line number to search around. (1-based)
    Returns:
        str: The code snippet around the line number.
    """
    with open(file_full_path) as f:
        file_content = f.readlines()

    if line_no < 1 or line_no > len(file_content):
        return None

    # start and end should also be 1-based valid line numbers
    start = max(1, line_no - window_size)
    end = min(len(file_content), line_no + window_size)
    snippet = ""
    for i in range(start, end):
        if with_lineno:
            snippet += f"{i} {file_content[i - 1]}"
        else:
            snippet += file_content[i]
    return snippet



def remove_line_numbers(function_code_with_line_numbers: str):
    """
    Remove line numbers from the function code and normalize indentation.
    Parameters
    ----------
    function_code_with_line_numbers : str
        The function code with line numbers at the beginning of each line.

    Returns
    -------
    str
        The function code with line numbers removed and normalized indentation.
    """
    # Remove line numbers
    code_without_line_numbers = re.sub(r'^\s*\d+\s', '', function_code_with_line_numbers, flags=re.MULTILINE)
    # Normalize indentation by detecting the minimum leading spaces in non-empty lines
    lines = code_without_line_numbers.splitlines()
    non_empty_lines = [line for line in lines if line.strip()]
    if not non_empty_lines:
        return ""  # If all lines are empty, return an empty string

    min_indent = min(len(line) - len(line.lstrip()) for line in non_empty_lines)
    normalized_code = "\n".join(line[min_indent:] if line.strip() else "" for line in lines)
    normalized_code = textwrap.dedent(normalized_code)
    return normalized_code


def extract_docstr(func_content: str):
    """
    Extract the docstring from a given function code string.

    Parameters
    ----------
    function_code : str
        The complete function definition as a string.

    Returns
    -------
    str or "None"
        The extracted docstring, or "None" if no docstring is found.
    """
    # Use a regular expression to remove leading line numbers and whitespace
    func_content = remove_line_numbers(func_content)
    # Parse the function code string into an AST node
    function_node = ast.parse(func_content).body[0]
    # Ensure it's a function definition
    if isinstance(function_node, ast.FunctionDef):
        # Extract and return the docstring
        return ast.get_docstring(function_node)
    return "None"


def extract_target_func_from_Task(task):
    """
    Extract the docstring from a task class
    """
    task_str = task.problem_statement
    func_content = re.findall(r"```(.*?)```", task_str, re.DOTALL)
    return func_content[0].strip()
     
    

def compute_bm25(std_str: str, search_result_list):
    """compute the bm25 between std_str and every ele in search_result_list.
    For every ele in search_result_list, we only extract the function docstring for comparison.
    Args:
        std_str (str): _description_
        str_list (list[SearchResult]): _description_

    Returns:
        ranked_results (list): ranked list, from high to low, [(SearchResult, scores)]
    """
    # TODO: should we exclude the function with the same docstring as std_str? they are likely to be the same function——include them first
    tokenized_str_list =[]
    for search_result in search_result_list:
        try:
            func_code = search_result.code
            tokenized_str_list.append(word_tokenize(func_code.lower()))
        except Exception as e:
            print(f'Exception happened in compute_bm25: {e}')
            continue
    # Initialize BM25
    bm25 = BM25Okapi(tokenized_str_list)
    tokenized_query = word_tokenize(std_str.lower())
    # Get similarity scores
    scores = bm25.get_scores(tokenized_query)
    # Rank and display results
    ranked_results = sorted(zip(search_result_list, scores), key=lambda x: x[1], reverse=True)
    return ranked_results[1:35] # we assume the first one is the target func



def extract_functions(files_content):
    """Extracts all functions and their identifiers from the nested files_content."""
    function_list = []
    function_identifiers = []
    for file_path, classes_dict in files_content.items():
        for class_name, functions_dict in classes_dict.items():
            for function_name, function_code in functions_dict.items():
                function_list.append(function_code)
                function_identifiers.append({
                    'file_path': file_path,
                    'class_name': class_name,
                    'function_name': function_name,
                    'function_body': function_code
                })
    return function_list, function_identifiers



def is_func_calls_target(candidate, target_func_name: str, target_class_name: str | None) -> bool:
    """given a candidate function, determine whether it calls the target funtions
    we consider the following situation: 
    the test must from a py file;
    when target is top-level:
        f(), ✅
        module.f(), ✅
    when target is within a class: 
        __init__ ✅
        Plot().add(...)	✅
        obj.add(...) where obj = Plot()	✅
        Plot(...).on().add(...)	✅
        obj.on(...).add(...) where obj = Plot()	✅
        self.add(...), cls.add(...) ✅
    
    Not covering:
    super().f(), 
    Decorators: @foo or @A.foo, 
    Indirect via wrapper/helper: wrapper() → foo(),
    init function,
    Args:
        candidate (SearchResult): SearchResult(file_path: str, start: int, end: int, class_name: None | str, func_name: str, code: whole function with line number)
        target_func_name (str)
        target_class_name (str | None): the class the target belongs to
    Returns:
        bool: true if called, false otherwise
    """
    if candidate.file_path.endswith(".py") == False:
        return False
    # remove line numbers and normalize indentation
    candidate_src_code = remove_line_numbers(candidate.code)
    # print(candidate_src_code)
    try:
        tree = ast.parse(candidate_src_code)
    except SyntaxError:
        # print('HUYIRAN: error happened in is_func_calls_target, when converting candidate into ast')
        return False
    class CallVisitor(ast.NodeVisitor):
        def __init__(self, target_func, target_class):
            self.target_func = target_func
            self.target_class = target_class
            self.var_to_class = {}   # Maps variable name → class name
            self.found = False

        def visit_Assign(self, node):
            # Detect: var = ClassName(...)
            if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name):
                class_name = node.value.func.id
                if self.target_class is None or class_name == self.target_class:
                    for target in node.targets:
                        if isinstance(target, ast.Name):
                            self.var_to_class[target.id] = class_name
            self.generic_visit(node)

        def visit_Call(self, node):
            # Check for attribute-style call: obj.func(...)
            # Special case: calling class constructor = calling __init__
            if self.target_func == "__init__":
                if isinstance(node.func, ast.Name) and node.func.id == self.target_class:
                    self.found = True
            # Check for attribute-style call: obj.func(...)
            if isinstance(node.func, ast.Attribute):
                func_name = node.func.attr
                if func_name == self.target_func:
                    caller = node.func.value
                    if isinstance(caller, ast.Name):
                        var_name = caller.id
                        if self.target_class is None or self.var_to_class.get(var_name) == self.target_class:
                            self.found = True
                    elif isinstance(caller, ast.Call):
                        # support Plot(...).add(...)
                        if self.contains_class_instantiation(caller, self.target_class):
                            self.found = True
                    elif isinstance(caller, ast.Attribute):
                        # support p1.on(...).add(...)
                        if self.contains_class_instantiation(caller, self.target_class):
                            self.found = True
            self.generic_visit(node)

        def contains_class_instantiation(self, node, class_name):
            if isinstance(node, ast.Call):
                if isinstance(node.func, ast.Name) and node.func.id == class_name:
                    return True
                return self.contains_class_instantiation(node.func, class_name)
            elif isinstance(node, ast.Attribute):
                return self.contains_class_instantiation(node.value, class_name)
            return False

    # visitor = CallVisitor()
    visitor = CallVisitor(target_func_name, target_class_name)
    visitor.visit(tree)
    return visitor.found

