import os
import datetime
import json


from extract_json_code import extract_json_code
import logging
from urllib.parse import urlparse
import shutil
import argparse
import subprocess
import tempfile

import pathlib
import time

logging.basicConfig(
    level=logging.INFO,
    filename="app.log",
    filemode="w",
    format="%(asctime)s - %(levelname)s - %(message)s",
)

def pretty_json(data) -> str:
    """返回美化后的 JSON 字符串，保持 Unicode 字符不被转义"""
    return json.dumps(data, indent=2, ensure_ascii=False)


def load_from_json(json_file_path):
    with open(json_file_path, "r", encoding="utf-8") as json_file:
        data = json.load(json_file)
    print(f"
    return data


def rollback_all(dir, code):
    """
    :dir: the git repo dir
    :code: should follow this schema
    [
        { "path": "path/to/file", "content": "file content" },
        ...
    ]
    """
    
    for root, dirs, files in os.walk(dir):
        
        top = pathlib.Path(root).relative_to(dir).parts
        if top and top[0] == ".git":
            continue
        for name in files:
            os.remove(os.path.join(root, name))
        for name in dirs:
            if name != ".git":
                shutil.rmtree(os.path.join(root, name))

    
    for file in code:
        file_path = pathlib.Path(os.path.join(dir, file["path"]))
        file_path.parent.mkdir(parents=True, exist_ok=True)
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(file["content"])


def rollback_part(dir, code, chosen_paths=[], except_paths=[]):
    """
    :dir: the git repo dir
    :code: should follow this schema
    [
        { "path": "path/to/file", "content": "file content" },
        ...
    ]
    """
    
    for file in code:
        file_path = pathlib.Path(os.path.join(dir, file["path"]))
        if except_paths:
            flag = True
            for e in except_paths:
                if e in str(file_path):  
                    print(f"Rollback all, and skip {file_path}")
                    flag = False
                    break
        if chosen_paths:
            flag = False
            for c in chosen_paths:
                if c in str(file_path):  
                    print(f"Rollback {file_path}")
                    flag = True
                    break
        if flag is False: continue
        file_path.parent.mkdir(parents=True, exist_ok=True)
        with open(file_path, "w", encoding="utf-8") as f:  
            f.write(file["content"])


def reset_buggy_code(buggy_dir, buggycode):
    rollback_all(buggy_dir, buggycode)


APPLY_SUCCESS = "Patch applied successfully!"
APPLY_FAIL = "Patch Apply Fail"
def apply_patch(patch_name, testjson, repo_path):
    patch_info = testjson[patch_name]
    with open(os.path.join(repo_path, "temp_patch.diff"), "w", encoding="utf-8") as pf:
        pf.write(patch_info)
    try:
        
        result = subprocess.run(
            ['git', 'apply', "temp_patch.diff"],
            cwd=repo_path,
            capture_output=True,
            text=True
        )
        
        if result.returncode == 0:
            
            return APPLY_SUCCESS 
        else:
            
            return result.stderr
    except Exception as e:
        
        return str(e)


TEST_PASS = "Test Pass"
TEST_FAIL = "Test Fail"
def unit_test(repo_name, command): 
    def run_shell_command(command):
        
        result = subprocess.run(command, shell=True,
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        
        output = result.stdout
        error = result.stderr
        
        if error:
            return json.dumps(error)
        
        escaped_output = json.dumps(output)
        return escaped_output
    os.chdir(os.path.join("buggycode/", repo_name))

    if 'munch' in repo_name:
        shutil.copy("../../preparation/Infinidat_munch/README.md", "./")
    elif 'sklearn-pandas' in repo_name:
        plan = 'cp -r ../../preparation/scikit-learn-contrib_sklearn-pandas/test_data/ tests/'
        run_shell_command(plan)
    elif 'pandarallel' in repo_name:
        plan = 'export PYTHONPATH=$(pwd)'
        run_shell_command(plan)

    
    
    escaped_output = run_shell_command(command)
    json_output = json.loads(escaped_output)
    
    return_info = TEST_PASS
    
    for line in json_output.splitlines():
        if 'FAILED' in line or 'ERROR' in line or 'Traceback' in line:
            return_info = TEST_FAIL
            break
    os.chdir("../../")
    return json_output, return_info

if __name__ == "__main__":
    
    
    
    
    
    
    origin_dir = "origincode/"
    buggy_dir = "buggycode/"
    
    gpt_res_path = "/"
    apply_out_path = "/"
    
    for subdir in os.listdir(gpt_res_path):
        print(subdir)
        subdir_path = os.path.join(gpt_res_path, subdir)
        subout_path = os.path.join(apply_out_path, subdir)
        if not os.path.exists(subout_path):
            os.mkdir(subout_path)
        if os.path.isdir(subdir_path):
            res_list = os.listdir(subdir_path)
            res_list = sorted([str(_) for _ in res_list], key=lambda x: int(x.split('_')[0]))
    
            
            for a_file in res_list[:3]:  
                gpt_res = os.path.join(subdir_path, a_file)  
                testjson = load_from_json(gpt_res)
                unit_test_result = {
                    "id": str(a_file),
                    "repoName": testjson.get('RepoName'),
                    "buggyCategory": [_['file'] for _ in testjson.get("BuggyCodeLocation", [])], 
                    "buggyTest": {
                        "IsPass": None,
                        "Messages": "", 
                    },
                    "originTest":{
                        "IsPass": None,
                        "Messages": "", 
                    },
                    "patchApplyINFO": [
                        
                        
                        
                        
                        
                        
                        
                        
                        
                    ]
                }
        
                if testjson["BuggyCodeLocation"]:
                    repo_location = testjson["BuggyCodeLocation"][0]["file"]
                elif testjson["Difficulty"] == "Difficult":
                    repo_location = testjson["OriginCode"][0]["path"]
                else:
                    print("No Location: ", a_file)
                    print(testjson["Difficulty"])
                    continue

                
                repo_name = repo_location.split("/", maxsplit=1)[0]
                command = testjson["Command"] 
                reset_buggy_code(buggy_dir, testjson["BuggyCode"])
                json_output, return_info = unit_test(repo_name, command)
                print('Buggy Code: ', return_info)
                unit_test_result['buggyTest']['IsPass'] = True if return_info == TEST_PASS else False
                unit_test_result['buggyTest']['Messages'] = json_output

                
                repo_name = repo_location.split("/", maxsplit=1)[0]
                command = testjson["Command"] 
                reset_buggy_code(origin_dir, testjson["OriginCode"])
                json_output, return_info = unit_test(repo_name, command)
                print('Origin Code: ', return_info)
                unit_test_result['originTest']['IsPass'] = True if return_info == TEST_PASS else False
                unit_test_result['originTest']['Messages'] = json_output

                command = testjson["Command"] 
                test_code_paths = []
                for cm in command:
                    test_code_paths += [os.path.join(buggy_dir, repo_name, _) for _ in cm.split() \
                                        if 'test' in _ and 'unittest' not in _]
                for patch_key in testjson["Results"].keys():
                    
                    if "patch" in patch_key or "Patch" in patch_key:  
                        print(f'[{patch_key}]')
                        new_item = {
                            "patchKey": patch_key,
                            "applyTest": {
                                "IsApplied": None,
                                "Messages": "",
                            },
                            "unitTest": {
                                
                                
                            },
                        }
                        reset_buggy_code(buggy_dir, testjson["BuggyCode"])  
                        echo = apply_patch(patch_key, testjson["Results"], os.path.join(buggy_dir, repo_name))  
                        new_item['applyTest']['Messages'] = echo
                        if echo == APPLY_SUCCESS:
                            new_item['applyTest']['IsApplied'] = True
                            
                            reset_buggy_code(buggy_dir, testjson["BuggyCode"])
                            apply_patch(patch_key, testjson["Results"], os.path.join(buggy_dir, repo_name))
                            rollback_part(buggy_dir, testjson["BuggyCode"], chosen_paths=test_code_paths)
                            print('Unittest func: ', return_info)
                            json_output, return_info = unit_test(repo_name, command)
                            new_item['unitTest']['IsPass1'] = True if return_info == TEST_PASS else False
                            new_item['unitTest']['Messages1'] = json_output
                            
                            reset_buggy_code(buggy_dir, testjson["BuggyCode"])
                            apply_patch(patch_key, testjson["Results"], os.path.join(buggy_dir, repo_name))
                            rollback_part(buggy_dir, testjson["BuggyCode"], except_paths=test_code_paths)
                            json_output, return_info = unit_test(repo_name, command)
                            print('Unittest test: ', return_info)  
                            new_item['unitTest']['IsPass2'] = True if return_info == TEST_PASS else False
                            new_item['unitTest']['Messages2'] = json_output
                            
                            json_output, return_info = unit_test(repo_name, command)  
                            print('Unittest both: ', return_info)
                            new_item['unitTest']['IsPass3'] = True if return_info == TEST_PASS else False
                            new_item['unitTest']['Messages3'] = json_output
                        else:
                            print(APPLY_FAIL)
                            new_item['applyTest']['IsApplied'] = False
                        unit_test_result["patchApplyINFO"].append(new_item)
                
                with open(os.path.join(subout_path, a_file), "w", encoding="utf-8") as f:
                    print(json.dumps(unit_test_result, indent=2, ensure_ascii=False), file=f)
