import os
import pickle
import sys
from importlib import import_module
from io import StringIO
from typing import Any, Dict, List

import coverage
from rich.progress import track

from evalplus.eval.utils import swallow_io
from tools.tsr.utils import get_problems, get_task_ids, to_path


class Capturing(list):
    def __enter__(self):
        self._stdout = sys.stdout
        sys.stdout = self._stringio = StringIO()
        return self

    def __exit__(self, *args):
        self.extend(self._stringio.getvalue().splitlines())
        del self._stringio
        sys.stdout = self._stdout


def parse_lcov(outputs: List[str]):
    switch, extracted_outputs = False, []
    for line in outputs:
        if switch == False and "tmp_src" in line:
            switch = True
        if switch == True and "end_of_record" in line:
            switch = False
        if switch:
            extracted_outputs.append(line)

    branch, branch_covered = [], []
    for line in extracted_outputs:
        if line.startswith("BRDA"):
            # BRDA format: BR:<lineno>,<blockno>,<branchno>,<taken>
            lineno, blockno, branchno, taken = line[5:].split(",")
            branch_sig = f"BR:{lineno},{blockno},{branchno}"
            branch.append(branch_sig)
            if taken not in ["0", "-"]:
                branch_covered.append(branch_sig)
    per = 1.0 if len(branch) == 0 else len(branch_covered) / len(branch)
    return per, branch, branch_covered


def test_code_coverage(
    identifier: str, code: str, inputs: List[List[Any]], entry_point: str
):
    module_name = f"tmp_src_{identifier}"
    with open(f"{module_name}.py", "w") as f:
        f.write(code)

    mod = import_module(module_name)
    func = getattr(mod, entry_point, None)
    assert func != None, f"entry_point = {entry_point} not exist, code: {code}"

    cov = coverage.Coverage(branch=True)
    cov.start()
    with swallow_io():
        for input_list in inputs:
            func(*input_list)
    cov.stop()
    with Capturing() as outputs:
        cov.lcov_report(outfile="-")

    ret = parse_lcov(outputs)

    os.remove(f"{module_name}.py")
    return ret


def collect_coverage_info(coverage_dir: str, dataset: str) -> Dict[str, Dict[str, Any]]:
    os.makedirs(coverage_dir, exist_ok=True)
    problems = get_problems(dataset)
    task_ids = get_task_ids(dataset)
    coverage_info = {task_id: {} for task_id in task_ids}
    for task_id in track(task_ids, description="Testing gt coverage..."):
        coverage_cache_path = os.path.join(coverage_dir, f"{to_path(task_id)}.pkl")
        if os.path.isfile(coverage_cache_path):
            with open(coverage_cache_path, "rb") as f:
                coverage_info[task_id] = pickle.load(f)
            continue
        groundtruth_code = (
            problems[task_id]["prompt"] + problems[task_id]["canonical_solution"]
        )
        plus_tests = problems[task_id]["plus_input"]
        entry_point = problems[task_id]["entry_point"]
        for i, plus_test in enumerate(plus_tests):
            per, branch, branch_covered = test_code_coverage(
                to_path(task_id), groundtruth_code, [plus_test], entry_point
            )
            test_id = f"plus_{i}"
            coverage_info[task_id].setdefault(test_id, []).extend(
                [(br, "gt") for br in branch_covered]
            )
        with open(coverage_cache_path, "wb") as f:
            pickle.dump(coverage_info[task_id], f)

    return coverage_info


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str, choices=["humaneval", "mbpp"])
    parser.add_argument("--report_dir", required=True, type=str)
    args = parser.parse_args()

    coverage_dir = os.path.join(args.report_dir, "coverage_cache")
    collect_coverage_info(coverage_dir, args.dataset)