import ast
import contextlib
import glob
import os
import re
import shutil
import subprocess
from collections.abc import Callable
from functools import wraps
from os import PathLike
from os.path import dirname as pdirname
from os.path import join as pjoin
from pathlib import Path
from subprocess import CalledProcessError
import tree_sitter_python as tspython
from tree_sitter import Language, Parser
import pdb
import docker
from loguru import logger
from app.log import log_and_print, log_exception


@contextlib.contextmanager
def cd(newdir):
    """
    Context manager for changing the current working directory
    :param newdir: path to the new directory
    :return: None
    """
    prevdir = os.getcwd()
    os.chdir(os.path.expanduser(newdir))
    try:
        yield
    finally:
        os.chdir(prevdir)


def run_command(cmd: list[str], **kwargs) -> subprocess.CompletedProcess:
    """
    Run a command in the shell.
    Args:
        - cmd: command to run
    """
    try:
        cp = subprocess.run(cmd, check=True, **kwargs)
    except subprocess.CalledProcessError as e:
        log_and_print(f"Error running command: {cmd}, {e}")
        raise e
    return cp


def is_git_repo() -> bool:
    """
    Check if the current directory is a git repo.
    """
    git_dir = ".git"
    return os.path.isdir(git_dir)


def initialize_git_repo_and_commit(logger=None):
    """
    Initialize the current directory as a git repository and make an initial commit.
    """
    init_cmd = ["git", "init"]
    add_all_cmd = ["git", "add", "."]
    commit_cmd = ["git", "commit", "-m", "Temp commit made by ACR."]
    run_command(init_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    run_command(add_all_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    run_command(commit_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def get_current_commit_hash() -> str:
    command = ["git", "rev-parse", "HEAD"]
    cp = subprocess.run(command, text=True, capture_output=True)
    try:
        cp.check_returncode()
        return cp.stdout.strip()
    except CalledProcessError as e:
        raise RuntimeError(f"Failed to get SHA-1 of HEAD: {cp.stderr}") from e


def repo_commit_current_changes():
    """
    Commit the current active changes so that it's safer to do git reset later on.
    Use case: for storing the changes made in pre_install and test_patch in a commit.
    """
    add_all_cmd = ["git", "add", "."]
    commit_cmd = [
        "git",
        "commit",
        "--allow-empty",
        "-m",
        "Temporary commit for storing changes",
    ]
    run_command(add_all_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    run_command(commit_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def clone_repo(clone_link: str, dest_dir: str, cloned_name: str):
    """
    Clone a repo to dest_dir.

    Returns:
        - path to the newly cloned directory.
    """
    clone_cmd = ["git", "clone", clone_link, cloned_name]
    create_dir_if_not_exists(dest_dir)
    with cd(dest_dir):
        run_command(clone_cmd)
    cloned_dir = pjoin(dest_dir, cloned_name)
    return cloned_dir


def clone_repo_and_checkout(
    clone_link: str, commit_hash: str, dest_dir: str, cloned_name: str
):
    """
    Clone a repo to dest_dir, and checkout to commit `commit_hash`.

    Returns:
        - path to the newly cloned directory.
    """
    cloned_dir = clone_repo(clone_link, dest_dir, cloned_name)
    checkout_cmd = ["git", "checkout", commit_hash]
    with cd(cloned_dir):
        run_command(checkout_cmd)
    return cloned_dir


def repo_clean_changes() -> None:
    """
    Reset repo to HEAD. Basically clean active changes and untracked files on top of HEAD.

    A lightweight version of `repo_reset_and_clean_checkout`. This is also used in different scenarios.
    """
    reset_cmd = ["git", "reset", "--hard"]
    clean_cmd = ["git", "clean", "-fd"]
    run_command(reset_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    run_command(clean_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    # pass # change to this in case it will recover the target function ground truth


def repo_reset_and_clean_checkout(commit_hash: str) -> None:
    """
    Run commands to reset repo to the original commit state.
    Cleans both the uncommited changes and the untracked files, and submodule changes.
    Assumption: The current directory is the git repository.
    """
    # NOTE: do these before `git reset`. This is because some of the removed files below
    # may actually be in version control. So even if we deleted such files here, they
    # will be brought back by `git reset`.
    # Clean files that might be in .gitignore, but could have been created by previous runs
    if os.path.exists(".coverage"):
        os.remove(".coverage")
    if os.path.exists("tests/.coveragerc"):
        os.remove("tests/.coveragerc")
    other_cov_files = glob.glob(".coverage.TSS.*", recursive=True)
    for f in other_cov_files:
        os.remove(f)

    reset_cmd = ["git", "reset", "--hard", commit_hash]
    clean_cmd = ["git", "clean", "-fd"]
    checkout_cmd = ["git", "checkout", commit_hash]
    run_command(reset_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    run_command(clean_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    # need to checkout before submodule init. Otherwise submodule may init to another version
    run_command(checkout_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    # this is a fail-safe combo to reset any changes to the submodule: first unbind all submodules
    # and then make a fresh checkout of them.
    # Reference: https://stackoverflow.com/questions/10906554/how-do-i-revert-my-changes-to-a-git-submodule
    submodule_unbind_cmd = ["git", "submodule", "deinit", "-f", "."]
    submodule_init_cmd = ["git", "submodule", "update", "--init"]
    run_command(
        submodule_unbind_cmd,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    run_command(
        submodule_init_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
    )


def run_script_in_conda(
    args: list[str], env_name: str, **kwargs
) -> subprocess.CompletedProcess:
    """
    Run a python command in a given conda environment.
    """
    if(env_name != ''):
        cmd = ["conda", "run", "-n", env_name, "python", *args]
    else:
        cmd = ["/opt/conda/bin/python", *args]
    return subprocess.run(cmd, **kwargs)


def run_string_cmd_in_conda(
    command: str, env_name: str, **kwargs
) -> subprocess.CompletedProcess:
    """
    Run a complete command in a given conda environment, where the command is a string.

    This is useful when the command to be run contains &&, etc.

    NOTE: use `conda activate` instead of `conda run` in this verison, so that we can
          run commands that contain `&&`, etc.
    """
    conda_bin_path = os.getenv("CONDA_EXE")  # for calling conda
    if conda_bin_path is None:
        raise RuntimeError("Env variable CONDA_EXE is not set")
    conda_root_dir = pdirname(pdirname(conda_bin_path))
    conda_script_path = pjoin(conda_root_dir, "etc", "profile.d", "conda.sh")
    conda_cmd = f"source {conda_script_path} ; conda activate {env_name} ; {command} ; conda deactivate"
    log_and_print(f"Running command: {conda_cmd}")
    return subprocess.run(conda_cmd, shell=True, **kwargs)


def create_dir_if_not_exists(dir_path: str):
    """
    Create a directory if it does not exist.
    Args:
        dir_path (str): Path to the directory.
    """
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)


def create_fresh_dir(dir_path: str | PathLike):
    """
    Create a dir from fresh.
    If there is existing dir with the same name, remove it.
    Args:
        dir_path (str): Path to the directory.
    """
    if os.path.exists(dir_path):
        shutil.rmtree(dir_path)
    os.makedirs(dir_path)


def to_relative_path(file_path: str, project_root: str) -> str:
    """Convert an absolute path to a path relative to the project root.

    Args:
        - file_path (str): The absolute path.
        - project_root (str): Absolute path of the project root dir.

    Returns:
        The relative path.
    """
    if Path(file_path).is_absolute():
        return str(Path(file_path).relative_to(project_root))
    else:
        return file_path


def to_absolute_path(file_path: str, project_root: str) -> str:
    """Convert a relative path to an absolute path.

    Args:
        - file_path (str): The relative path.
        - project_root (str): Absolute path of a root dir.
    """
    return pjoin(project_root, file_path)


def find_file(directory, filename) -> str | None:
    """
    Find a file in a directory. filename can be short name, relative path to the
    directory, or an incomplete relative path to the directory.
    Returns:
        - the relative path to the file if found; None otherwise.
    """

    # Helper method one
    def find_file_exact_relative(directory, filename) -> str | None:
        if os.path.isfile(os.path.join(directory, filename)):
            return filename
        return None

    # Helper method two
    def find_file_shortname(directory, filename) -> str | None:
        for root, dirs, files in os.walk(directory):
            for file in files:
                if file == filename:
                    return os.path.relpath(os.path.join(root, file), directory)
        return None

    # if the filename is exactly the relative path.
    found = find_file_exact_relative(directory, filename)
    if found is not None:
        return found

    # if the filename is a short name without any directory
    found = find_file_shortname(directory, filename)
    if found is not None:
        return found

    # if the filename has some directory, but is not a relative path to
    # the directory
    parts = filename.split(os.path.sep)
    shortname = parts[-1]
    found = find_file_shortname(directory, shortname)
    if found is None:
        # really cannot find this file
        return None
    # can find this shortname, but we also need to check whether the intermediate
    # directories match
    if filename in found:
        return found
    else:
        return None


def parse_function_invocation(
    invocation_str: str,
) -> tuple[str, list[str]]:
    # edited to handle the case of search_method_relevance
    # if("search_method_relevance" in invocation_str):
    #     function_name = "search_method_relevance"
    #     params = invocation_str[len(function_name):] # (XXXX, XXXX)
        
    #     first_para = params.split(',')[0][1:].strip() # the method name: 'label'
    #     second_para = ','.join(params.split(',')[1:])[:-1].strip() # the docstring, which might have wrong usage of ' or "
    #     if('"""' not in second_para and "'''" not in second_para):
    #         second_para = '"""' + second_para[1:-1] + '"""'
    #     if("'''" in second_para):
    #         second_para = second_para.replace("'''", '"""')
    #     # return function_name, arguments
    #     invocation_str = f'{function_name}({first_para}, {second_para})'

    try:
        tree = ast.parse(invocation_str)
        expr = tree.body[0]
        assert isinstance(expr, ast.Expr)
        call = expr.value
        assert isinstance(call, ast.Call)
        func = call.func
        assert isinstance(func, ast.Name)
        function_name = func.id
        raw_arguments = [ast.unparse(arg) for arg in call.args]
        # clean up spaces or quotes, just in case
        arguments = [arg.strip().strip("'").strip('"') for arg in raw_arguments]
        try:
            new_arguments = [ast.literal_eval(x) for x in raw_arguments]
            if new_arguments != arguments:
                # log_and_print( # : this log is usually printed when some arguments are handled as string in old but int in new, such as: old result is ['seaborn/_core/scales.py', '632', '10'], new result  is ['seaborn/_core/scales.py', 632, 10]
                #     f"Refactored invocation argument parsing gives different result on "
                #     f"{invocation_str!r}: old result is {arguments!r}, new result "
                #     f" is {new_arguments!r}",
                # )
                log_and_print( # : this log is usually printed when some arguments are handled as string in old but int in new, such as: old result is ['seaborn/_core/scales.py', '632', '10'], new result  is ['seaborn/_core/scales.py', 632, 10]
                    f"Refactored invocation argument parsing gives different results",
                    # f"{invocation_str!r}: old result is {arguments!r}, new result "
                    # f" is {new_arguments!r}",
                )
        except Exception as e:
            log_and_print(
                f"Refactored invocation argument parsing failed on {invocation_str!r}: {e!s}",
            )
    except Exception as e:
        raise ValueError(f"Invalid function invocation: {invocation_str}") from e

    return function_name, arguments


def catch_all_and_log(func: Callable) -> Callable:
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            log_exception(e)
            err = str(e)
            # TODO: this string is there for legacy reasons
            summary = "The tool returned error message."
            return err, summary, False

    return wrapper


def coroutine(func):
    @wraps(func)
    def primer(*args, **kwargs):
        gen = func(*args, **kwargs)
        next(gen)
        return gen

    return primer


def build_container(container_name: str):
    client = docker.from_env()  # Docker client to interact with Docker
    container = client.containers.get(container_name)
    if container.status != 'running':
        container.start()
        container.reload()  # update status again
    return container
    
    

def copy_file_to_docker(container, src_path, dest_path):
    """Copy a file from the local system to the Docker container."""
    command = f"docker cp {src_path} {container.id}:{dest_path}"
    try:
        subprocess.run(command, shell=True, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        logger.info(f"Successfully copied {src_path} to {dest_path} in container {container}")
    except subprocess.CalledProcessError as e:
        print(f"Failed to copy file to Docker container: {e}")
        exit()
        
        
def copy_file_from_docker(container_name: str, container_file_path: str, local_destination_path: str):
    """
    Copy a file from a Docker container to the host machine.

    Parameters:
        container_name (str): The name or ID of the container.
        container_file_path (str): The full path to the file inside the container.
        local_destination_path (str): The path on the host where the file should be copied.
    """
    try:
        subprocess.run(
            ["docker", "cp", f"{container_name}:{container_file_path}", local_destination_path],
            check=True
        )
        # print(f"Copied {container_file_path} from container '{container_name}' to {local_destination_path}")
    except subprocess.CalledProcessError as e:
        print(f"Failed to copy file: {e}")
    



def execute_command_in_docker(container, command):
    """Execute command inside the Docker container"""
    try:
        exec_result = container.exec_run(command)
        output = exec_result.output.decode('utf-8')
        return output, exec_result.exit_code
    except Exception as e:
        print(f"Error executing command in Docker: {str(e)}")
        return None, -1



    
def remove_function_from_repo(function_name, file_path, replace_contents = None):
    """
    Remove the specified function from the repo.
    Return start_line, end_line, bool_value
    """
    def find_subnode_with_name_for_decorated_calls(decorated_call_node, function_name):
        for child_node in decorated_call_node.children:
            if child_node.type == "function_definition":
                child_name_node = child_node.child_by_field_name("name")
                if child_name_node.text.decode() == function_name:
                    return child_node
        return None
    
    def find_child_with_name_for_class(current_node, function_name):
        cursor = current_node.walk()
        while(True):
            descentent_node = cursor.node
            if descentent_node.type == "function_definition":
                function_name_node = descentent_node.child_by_field_name("name")
                if function_name_node.text.decode() == function_name:
                    return descentent_node
            elif descentent_node.type == "decorated_definition":
                result_node = find_subnode_with_name_for_decorated_calls(descentent_node, function_name)
                if result_node:
                    return result_node
            if cursor.goto_first_child():
                    continue
            if cursor.goto_next_sibling():
                continue
            while cursor.goto_parent():
                if cursor.goto_next_sibling():
                    break
            else:
                break
        return None
    
    start_line, end_line = None, None
    with open(file_path, 'r') as f:
        code = f.read()
    class_name = None
    if "." in function_name:
        [class_name, function_name] = function_name.split(".")[-2:]
    
    # Parse and modify the file if the function is found
    if not function_name in code:
        if os.path.basename(file_path).replace(".py","") == function_name:
            return -1, -1, False
        raise Exception(f"ERROR: function {function_name} not found in file: {file_path}")
    elif class_name and not class_name in code and class_name != os.path.basename(file_path).replace(".py",""):
        raise Exception(f"ERROR: There should be class name, but class name {class_name} not found in file: {file_path}")

    PY_LANGUAGE = Language(tspython.language())
    # Initialize the parser with the Python language
    parser = Parser(PY_LANGUAGE)
    tree = parser.parse(bytes(code, 'utf8'))
    root_node = tree.root_node
    cursor = root_node.walk()
    new_code_lines = code.splitlines(keepends=True)
    while True:
        node = cursor.node
        if class_name and node.type == "class_definition":
            # if this is the target class
            if class_name == node.child_by_field_name("name").text.decode():
                function_node = find_child_with_name_for_class(node, function_name)
                if function_node:
                    start_line = function_node.start_point[0]
                    end_line = function_node.end_point[0] + 1
                    assert f"def {function_name}" in "".join(new_code_lines[start_line:end_line])
                    break
                else:
                    pdb.set_trace()
        elif (not class_name or class_name == os.path.basename(file_path).replace(".py","")) and node.type == "function_definition":
            function_name_node = node.child_by_field_name("name")
            if function_name_node.text.decode() == function_name:
                start_line = node.start_point[0]
                end_line = node.end_point[0] + 1
                assert f"def {function_name}" in "".join(new_code_lines[start_line:end_line])
                break
        if cursor.goto_first_child():
            continue
        if cursor.goto_next_sibling():
            continue

        while cursor.goto_parent():
            if cursor.goto_next_sibling():
                break
        else:
            break
    if start_line and end_line and start_line < end_line and len(new_code_lines) >= 0:
        with open(file_path, 'w') as f:
            new_code_lines[start_line:end_line] = replace_contents
            f.writelines(new_code_lines)
        return start_line+1, start_line+1+len(replace_contents.strip().split('\n')), True
    return -1, -1, False



def clean_ansi(text):
    ansi_escape = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]')
    return ansi_escape.sub('', text)