import hashlib
import os
import re
import json
import platform
from pathlib import Path

from dataclasses import dataclass
from typing import Any, Optional, Union, cast

from swebench.harness.constants import (
    DEFAULT_DOCKER_SPECS,
    KEY_INSTANCE_ID,
    LATEST,
    MAP_REPO_TO_EXT,
    MAP_REPO_VERSION_TO_SPECS,
    USE_X86,
    SWEbenchInstance,
    UTF8,
)
from swebench.harness.dockerfiles import (
    get_dockerfile_base,
    get_dockerfile_env,
    get_dockerfile_instance,
)
from swebench.harness.test_spec.create_scripts import (
    make_repo_script_list,
    make_env_script_list,
    make_eval_script_list,
)
from swebench.harness.path_utils import safe_component


@dataclass
class TestSpec:
    """
    A dataclass that represents a test specification for a single instance of SWE-bench.
    """

    instance_id: str
    repo: str
    version: str
    repo_script_list: list[str]
    eval_script_list: list[str]
    env_script_list: list[str]
    arch: str
    FAIL_TO_PASS: list[str]
    PASS_TO_PASS: list[str]
    language: str
    docker_specs: dict
    namespace: Optional[str]
    base_image_tag: str = LATEST
    env_image_tag: str = LATEST
    instance_image_tag: str = LATEST
    # number of newly injected TE tests for progress reporting
    new_te_tests_count: int = 0

    @property
    def setup_env_script(self):
        return (
            "\n".join(["#!/bin/bash", "set -euxo pipefail"] + self.env_script_list)
            + "\n"
        )

    @property
    def eval_script(self):
        return (
            "\n".join(["#!/bin/bash", "set -uxo pipefail"] + self.eval_script_list)
            + "\n"
        )
        # Don't exit early because we need to revert tests at the end

    @property
    def install_repo_script(self):
        return (
            "\n".join(["#!/bin/bash", "set -euxo pipefail"] + self.repo_script_list)
            + "\n"
        )

    @property
    def base_image_key(self):
        """
        If docker_specs are present, the base image key includes a hash of the specs.
        """
        if self.docker_specs != {}:
            hash_key = str(self.docker_specs)
            hash_object = hashlib.sha256()
            hash_object.update(hash_key.encode("utf-8"))
            hash_value = hash_object.hexdigest()
            val = hash_value[
                :10
            ]  # 10 characters is still likely to be unique given only a few base images will be created
            return f"sweb.base.{MAP_REPO_TO_EXT[self.repo]}.{self.arch}.{val}:{self.base_image_tag}"
        return (
            f"sweb.base.{MAP_REPO_TO_EXT[self.repo]}.{self.arch}:{self.base_image_tag}"
        )

    @property
    def env_image_key(self):
        """
        The key for the environment image is based on the hash of the environment script list.
        If the environment script list changes, the image will be rebuilt automatically.

        Note that old images are not automatically deleted, so consider cleaning up old images periodically.
        """
        hash_key = str(self.env_script_list)
        if self.docker_specs != {}:
            hash_key += str(self.docker_specs)
        hash_object = hashlib.sha256()
        hash_object.update(hash_key.encode("utf-8"))
        hash_value = hash_object.hexdigest()
        val = hash_value[:22]  # 22 characters is still very likely to be unique
        return f"sweb.env.{MAP_REPO_TO_EXT[self.repo]}.{self.arch}.{val}:{self.env_image_tag}"

    @property
    def instance_image_key(self):
        key = f"sweb.eval.{self.arch}.{self.instance_id.lower()}:{self.instance_image_tag}"
        if self.is_remote_image:
            key = f"{self.namespace}/{key}".replace("__", "_1776_")
        return key

    @property
    def is_remote_image(self):
        return self.namespace is not None

    def get_instance_container_name(self, run_id=None):
        base = "sweb.eval"
        if not run_id:
            iid = safe_component(self.instance_id.lower(), max_len=64)
            name = f"{base}.{iid}"
        else:
            iid = safe_component(self.instance_id.lower(), max_len=48)
            rid = safe_component(str(run_id), max_len=48)
            name = f"{base}.{iid}.{rid}"
        # Docker name limit is 128 bytes; keep a margin
        if len(name) <= 120:
            return name
        import hashlib as _hashlib
        h = _hashlib.sha1(name.encode("utf-8")).hexdigest()[:8]
        # keep prefix and append short hash
        return f"{name[:110]}_{h}"

    @property
    def base_dockerfile(self):
        return get_dockerfile_base(
            self.platform,
            self.arch,
            self.language,
            **{**DEFAULT_DOCKER_SPECS, **self.docker_specs},
        )

    @property
    def env_dockerfile(self):
        return get_dockerfile_env(
            self.platform,
            self.arch,
            self.language,
            self.base_image_key,
            **{**DEFAULT_DOCKER_SPECS, **self.docker_specs},
        )

    @property
    def instance_dockerfile(self):
        return get_dockerfile_instance(self.platform, self.language, self.env_image_key)

    @property
    def platform(self):
        if self.arch == "x86_64":
            return "linux/x86_64"
        elif self.arch == "arm64":
            return "linux/arm64/v8"
        else:
            raise ValueError(f"Invalid architecture: {self.arch}")


def get_test_specs_from_dataset(
    dataset: Union[list[SWEbenchInstance], list[TestSpec]],
    namespace: Optional[str] = None,
    instance_image_tag: str = LATEST,
) -> list[TestSpec]:
    """
    Idempotent function that converts a list of SWEbenchInstance objects to a list of TestSpec objects.
    """
    if isinstance(dataset[0], TestSpec):
        return cast(list[TestSpec], dataset)
    return list(
        map(
            lambda x: make_test_spec(x, namespace, instance_image_tag),
            cast(list[SWEbenchInstance], dataset),
        )
    )

import ast

def extract_nodes(file_content):
    tree = ast.parse(file_content)

    def _is_func(node):  # regular or async function
        return isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))

    def _class_info(cls_node: ast.ClassDef):
        methods = []
        nested_classes = []
        for n in cls_node.body:
            if _is_func(n):
                methods.append(n.name)
            elif isinstance(n, ast.ClassDef):
                nested_classes.append(_class_info(n))
        methods = [ f'{cls_node.name}.{method}' for method in methods ]
        nested_classes = [  f'{cls_node.name}.{method}' for method in nested_classes ]
        methods.extend(nested_classes)
        return methods

    classes= []
    functions= []

    for node in tree.body:
        if isinstance(node, ast.ClassDef):
            classes.extend(_class_info(node))
        elif _is_func(node):
            functions.append(node.name)

    return classes + functions

def extract_test_headers(repo, test_file, test_content):
    if repo == 'django/django':
        test_module = '.'.join(test_file.replace('.py','').split('/')[1:])
        nodes = extract_nodes(test_content)
        tests = []
        for node in nodes:
            if '.' in node:
                nodesplit = node.split('.')
                head, tail = '.'.join(nodesplit[:-1]), nodesplit[-1]
                tests.append(f"{tail} ({test_module}.{head})")
            else:
                tests.append(f"{node} ({test_module})")
        return tests
    elif repo == 'sympy/sympy':
        tests = []
        nodes = extract_nodes(test_content)
        for node in nodes:
            tests.append(f"{node}")
        return tests
    elif repo in ['astropy/astropy', 'matplotlib/matplotlib', 'mwaskom/seaborn',
                  'pydata/xarray', 'pytest-dev/pytest', 'scikit-learn/scikit-learn',
                  'sphinx-doc/sphinx', 'psf/requests', 'pallets/flask', 'pylint-dev/pylint' ]:
        nodes = extract_nodes(test_content)
        tests = []
        for node in nodes:
            if '.' in node:
                nodesplit = node.split('.')
                head, tail = '.'.join(nodesplit[:-1]), nodesplit[-1]
                tests.append(f"{test_file}::{head}::{tail}")
            else:
                tests.append(f"{test_file}::{node}")
        return tests
    else:
        print(f"ERROR: extract_test_headers: repo {repo} not recognized")
        return []

def get_node(repo, test_file, entry):
    if repo == 'django/django':
        test_module = '.'.join(test_file.replace('.py','').split('/')[1:])
        pattern = f"([\\d\\w_\\.]+) \\({test_module}.([\\d\\w_\\.]+)\\)"
        match = re.search(pattern, entry)
        if not match:
            pattern = f"([\\d\\w_\\.]+) \\({test_module}\\)"
            match = re.search(pattern, entry)
            if not match:
                print(f"get_node: No match: {entry} -> {test_module}")
                return None
            node = match.groups()
            return node
        tail, head = match.groups()
        node = f"{head}.{tail}"
        return node
    elif repo == 'sympy/sympy':
        return entry
    elif repo in ['astropy/astropy', 'matplotlib/matplotlib', 'mwaskom/seaborn',
                  'pydata/xarray', 'pytest-dev/pytest', 'scikit-learn/scikit-learn',
                  'sphinx-doc/sphinx', 'psf/requests' ]:
        node = ".".join(entry.split("::")[1:])
        return node
    else: return None

def make_test_spec(
    instance: SWEbenchInstance,
    namespace: Optional[str] = None,
    base_image_tag: str = LATEST,
    env_image_tag: str = LATEST,
    instance_image_tag: str = LATEST,
) -> TestSpec:
    if isinstance(instance, TestSpec):
        return instance
    assert base_image_tag is not None, "base_image_tag cannot be None"
    assert env_image_tag is not None, "env_image_tag cannot be None"
    assert instance_image_tag is not None, "instance_image_tag cannot be None"
    instance_id = instance[KEY_INSTANCE_ID]
    repo = instance["repo"]
    version = instance.get("version")
    base_commit = instance["base_commit"]
    problem_statement = instance.get("problem_statement")
    hints_text = instance.get("hints_text")  # Unused
    test_patch = instance["test_patch"]

    def _from_json_or_obj(key: str) -> Any:
        """If key points to string, load with json"""
        if key not in instance:
            # If P2P, F2P keys not found, it's a validation instance
            return []
        if isinstance(instance[key], str):
            return json.loads(instance[key])
        return instance[key]

    pass_to_pass = _from_json_or_obj("PASS_TO_PASS")
    fail_to_pass = _from_json_or_obj("FAIL_TO_PASS")

    ## TODO: 
        ## read the json file containing tests
        ## add those tests to pass_to_pass & fail_to_pass test lists
    # if isinstance(fail_to_pass, list):
    #     fail_to_pass.append("astropy/io/fits/tests/test_connect.py::test_testenhancer_failing")

    env_name = "testbed"
    repo_directory = f"/{env_name}"
    specs = MAP_REPO_VERSION_TO_SPECS[repo][version]
    docker_specs = specs.get("docker_specs", {})

    repo_script_list = make_repo_script_list(
        specs, repo, repo_directory, base_commit, env_name
    )
    env_script_list = make_env_script_list(instance, specs, env_name)
    eval_script_list = make_eval_script_list(
        instance, specs, env_name, repo_directory, base_commit, test_patch
    )
    if platform.machine() in {"aarch64", "arm64"}:
        # use arm64 unless explicitly specified
        arch = "arm64" if instance_id not in USE_X86 else "x86_64"
    else:
        arch = "x86_64"

    test_files = re.findall(r'^diff --git a/(.*?) b/', instance['test_patch'], flags=re.MULTILINE)
    # Only use concrete Python test files; diffs may include directories like tests/test_utils
    test_files = [f for f in test_files if f.endswith('.py')]
    # Helper: identify real test files (under tests/ or named test_*.py)
    def _is_test_file(path: str) -> bool:
        name = Path(path).name
        # Exclude actual module files that aren't tests
        if name.startswith('_') and not name.startswith('test_'):
            return False
        return (
            path == 'accepted_tests.py'
            or path == 'accepted_tests_model_any.py'
            or '/tests/' in path
            or path.startswith('tests/')
            or name.startswith('test_')
            or name.endswith('_test.py')
        ) and not name.endswith('.pyc') and path.endswith('.py')
    # Filter to test files only
    test_files = [f for f in test_files if _is_test_file(f)]

    import os
    add_test_enhancer_patches = os.environ.get('TE', None) is None
    TE_Id = os.environ.get('TE_ID')
    # Always initialize; may remain empty if TE is disabled or no new tests
    new_fail_to_pass: list[str] = []
    new_te_tests_count: int = 0
    # Only inject Test Enhancer patches when explicitly enabled (TE unset) AND TE_ID is provided
    update_test_file_command: list[str] = []
    if add_test_enhancer_patches and TE_Id is not None:
        HEREDOC_DELIMITER = "EOF_114329324912"
        testgen_patch_dir = Path(f"logs/test_enhancer/{TE_Id}/{instance['instance_id']}")
        # Optional: restrict to accepted tests listed in CSV (instance_id,test_header)
        accepted_filter: set[str] | None = None
        acc_csv = os.environ.get('TE_ACCEPTED_CSV')
        if acc_csv and Path(acc_csv).is_file():
            try:
                import csv as _csv
                accepted_filter = set()
                with open(acc_csv, 'r', encoding=UTF8, newline='') as _f:
                    reader = _csv.DictReader(_f)
                    for row in reader:
                        if row.get('instance_id') == instance_id and row.get('test_header'):
                            accepted_filter.add(row['test_header'])
            except Exception:
                accepted_filter = None
        # Also include any baseline files present under TE dir (reverse '__' back to '/')
        try:
            extra_files = []
            if testgen_patch_dir.exists():
                for child in testgen_patch_dir.iterdir():
                    if child.is_file() and child.name.endswith('.py') and not child.name.startswith('out_'):
                        # reverse the naming convention used when saving
                        logical = child.name.replace('__', '/')
                        if _is_test_file(logical):
                            extra_files.append(logical)
            # Merge and deduplicate, filter out non-test files
            all_files = {*test_files, *extra_files}
            test_files = sorted([f for f in all_files if _is_test_file(f)])
        except Exception:
            pass
        # Include any discovered .py in digit-named subfolders (handle both out_*.py and nested real paths)
        try:
            discovered = []
            subs = [p for p in testgen_patch_dir.iterdir() if p.is_dir() and p.name.isdigit()]
            for sub in subs:
                for f in sub.rglob('*.py'):
                    if f.name.startswith('out_'):
                        logical = f.name[len('out_'):].replace('__', '/')
                    else:
                        # Use path relative to subfolder
                        logical = f.relative_to(sub).as_posix()
                    if _is_test_file(logical):
                        discovered.append(logical)
            # Merge and deduplicate, filter out non-test files
            discovered = sorted(set(discovered))
            all_files = set(test_files + discovered)
            test_files = sorted([f for f in all_files if _is_test_file(f)])
        except Exception:
            pass
        # Track global union of new headers across all files for accurate counting
        union_new_all: set[str] = set()
        def _normalize_repo_path(name: str) -> str:
            s = name
            if s.startswith('llm_'):
                s = s[len('llm_'):]
            s = s.replace('__', '/')
            return s

        # Deduplicate to avoid duplicate create/append operations
        test_files = sorted(set(test_files))
        processed_files: set[str] = set()
        for test_file in test_files:
            if test_file in processed_files:
                continue
            processed_files.add(test_file)
            # TE dir may store as-is or with '__' flattening and optional 'llm_' prefix
            te_name = test_file
            if not (testgen_patch_dir / te_name).is_file():
                te_name = test_file.replace('/', '__')
            test_file_path = testgen_patch_dir / te_name
            baseline_headers = []
            if test_file_path.is_file():
                test_content = test_file_path.read_text(encoding=UTF8)
                # Skip shell error outputs captured from previous runs
                if not (test_content.lstrip().startswith("cat: ") or "No such file or directory" in test_content):
                    try:
                        baseline_headers = extract_test_headers(repo, test_file, test_content)
                    except Exception:
                        baseline_headers = []
            # Scan all numeric subfolders and collect candidates
            candidates: list[tuple[int, str, set[str]]] = []  # (idx, content, new_headers)
            union_headers: set[str] = set()
            subs = [p for p in testgen_patch_dir.iterdir() if p.is_dir() and p.name.isdigit()]
            for sub in subs:
                try:
                    idx = int(sub.name)
                except Exception:
                    idx = -1
                # Option A: flat out_ mapping file (use TE naming)
                out_path = sub / f"out_{te_name}"
                direct_content = None
                if out_path.is_file():
                    direct_content = out_path.read_text(encoding=UTF8)
                else:
                    # Option B: nested real path (use normalized repo path)
                    direct_path = sub.joinpath(*_normalize_repo_path(test_file).split('/'))
                    if direct_path.is_file():
                        direct_content = direct_path.read_text(encoding=UTF8)
                if direct_content is None:
                    continue
                if direct_content.lstrip().startswith("cat: ") or "No such file or directory" in direct_content:
                    continue
                try:
                    repo_file = _normalize_repo_path(test_file)
                    headers_i = extract_test_headers(repo, repo_file, direct_content)
                except Exception:
                    headers_i = []
                new_i = set(headers_i) - set(baseline_headers)
                # Apply accepted filter if provided
                if accepted_filter is not None:
                    new_i = {h for h in new_i if h in accepted_filter}
                if new_i:
                    candidates.append((idx, direct_content, new_i))
                    union_headers.update(new_i)
            # Also consider top-level TE file itself as a candidate (useful when no numeric subfolders exist)
            if test_file_path.is_file():
                try:
                    content_top = test_file_path.read_text(encoding=UTF8)
                    # Treat all headers in the top-level TE file as new
                    repo_file = _normalize_repo_path(test_file)
                    headers_top = extract_test_headers(repo, repo_file, content_top)
                    new_top = set(headers_top)
                    if accepted_filter is not None:
                        new_top = {h for h in new_top if h in accepted_filter}
                    if new_top:
                        candidates.append((-1, content_top, new_top))
                        union_headers.update(new_top)
                except Exception:
                    pass

            # If accepted filter provided but no candidates survived, skip this file
            # Choose the candidate that contributes the most new tests; tie-break by highest index
            best_idx = -1
            best_content = None
            best_headers: set[str] = set()
            if candidates:
                candidates.sort(key=lambda t: (len(t[2]), t[0]))
                # If we have an accepted filter, avoid overwriting the whole file; only append accepted tests later
                if accepted_filter is None:
                    best_idx, best_content, best_headers = candidates[-1]
                    repo_file = _normalize_repo_path(test_file)
                    # Only create the file with best content if it does not already exist
                    update_test_file_command.append(
                        f"if [ -f {repo_file} ]; then :; else cat > {repo_file} <<'{HEREDOC_DELIMITER}'\n{best_content}\n{HEREDOC_DELIMITER}; fi"
                    )
                # Avoid noisy prints in quiet mode
                if os.environ.get('TE_QUIET', '').lower() not in ('1','true','yes'):
                    print(best_headers if best_headers else union_headers)
                # Extend with ALL union headers (already filtered if accepted_filter present) for grading
                new_fail_to_pass.extend(list(union_headers))
                # Accumulate across files for accurate total count
                try:
                    union_new_all.update(union_headers)
                except Exception:
                    pass
                # If union headers exceed best_headers (or if we skipped best), append the missing ones
                # Always plan to append all union headers to preserve baseline plus new tests
                missing = set(union_headers)
                if missing:
                    try:
                        # Build union additions content by extracting nodes from the candidates that contain them
                        def _extract_source_segment(src: str, node) -> str:
                            try:
                                import ast as _ast
                                # Python >=3.8 has lineno/end_lineno
                                if hasattr(node, 'lineno') and hasattr(node, 'end_lineno'):
                                    lines = src.splitlines()
                                    seg = "\n".join(lines[node.lineno-1:node.end_lineno])
                                    return seg
                                return ''
                            except Exception:
                                return ''

                        # Collect imports and code chunks
                        union_imports: list[str] = []
                        union_chunks: list[str] = []
                        seen_classes: set[str] = set()
                        # Map header -> (idx, content)
                        header_source_map: dict[str, tuple[int, str]] = {}
                        for idx, content_i, headers_i in candidates:
                            for h in headers_i:
                                if h in missing and h not in header_source_map:
                                    header_source_map[h] = (idx, content_i)
                        import re as _re
                        for h in sorted(missing):
                            idx_src, src = header_source_map[h]
                            try:
                                import ast as _ast
                                tree = _ast.parse(src)
                                # gather imports
                                for n in tree.body:
                                    if isinstance(n, (_ast.Import, _ast.ImportFrom)):
                                        seg = _extract_source_segment(src, n)
                                        if seg and seg not in union_imports:
                                            union_imports.append(seg)
                                node_key = get_node(repo, test_file, h)
                                if node_key and '.' in node_key:
                                    cls_name, meth_name = node_key.split('.', 1)
                                    cls_node = next((n for n in tree.body if isinstance(n, _ast.ClassDef) and n.name == cls_name), None)
                                    if cls_node is not None:
                                        # Append entire class definition; rename if collision risk
                                        seg = _extract_source_segment(src, cls_node)
                                        if seg:
                                            # rename class if name already seen to avoid duplicate definitions
                                            new_name = cls_name + "_TEUnion" if cls_name in seen_classes else cls_name
                                            if new_name != cls_name:
                                                seg = _re.sub(r"(^|\n)(\s*)class\s+" + _re.escape(cls_name) + r"\b", r"\\1\\2class " + new_name, seg)
                                            seen_classes.add(new_name)
                                            union_chunks.append(seg)
                                else:
                                    # top-level function
                                    fn_node = next((n for n in tree.body if isinstance(n, (_ast.FunctionDef, _ast.AsyncFunctionDef)) and n.name == (node_key or h).split('.')[-1]), None)
                                    if fn_node is not None:
                                        seg = _extract_source_segment(src, fn_node)
                                        if seg:
                                            union_chunks.append(seg)
                            except Exception:
                                continue

                        if union_chunks:
                            union_module = "\n# Auto-generated TE union additions (appended)\n" + "\n".join(union_imports) + ("\n\n" if union_imports else "") + "\n\n".join(union_chunks) + "\n"
                            # Append to the same test file so pytest directives still include these tests
                            repo_file = _normalize_repo_path(test_file)
                            update_test_file_command.append(
                                f"cat >> {repo_file} <<'{HEREDOC_DELIMITER}'\n{union_module}\n{HEREDOC_DELIMITER}"
                            )
                    except Exception:
                        pass
                # Persist selection summary and merged copy for traceability
                try:
                    summary = {
                        "test_file": test_file,
                        "baseline_count": len(baseline_headers),
                        "candidates": [
                            {"dir_index": idx, "new_count": len(hdrs)} for idx, _c, hdrs in candidates
                        ],
                        "best": {"dir_index": best_idx, "new_count": len(best_headers)},
                        "union_new_count": len(union_headers),
                        "union_missing_in_best": sorted(list(set(union_headers) - set(best_headers))),
                    }
                    (testgen_patch_dir / "merge_selection.json").write_text(
                        json.dumps(summary, indent=2), encoding=UTF8
                    )
                    if best_content is not None:
                        merged_dir = testgen_patch_dir / "merged"
                        merged_dir.mkdir(parents=True, exist_ok=True)
                        (merged_dir / f"out_{test_file.replace('/','__')}").write_text(best_content, encoding=UTF8)
                except Exception:
                    pass

        if update_test_file_command:
            # Insert right after apply_test_patch_command line for robustness
            try:
                anchor = next(
                    i for i, cmd in enumerate(eval_script_list)
                    if cmd.startswith("git apply -v - <<'EOF_114329324912'")
                )
                eval_script_list = (
                    eval_script_list[:anchor+1]
                    + update_test_file_command
                    + eval_script_list[anchor+1:]
                )
            except StopIteration:
                # Fallback to previous behavior if anchor not found
                eval_script_list = eval_script_list[:-4] + update_test_file_command + eval_script_list[-4:]
        if new_fail_to_pass:
            fail_to_pass.extend(new_fail_to_pass)
        # Count new TE tests for this instance (unique) using the union seen across all candidates/files
        new_te_tests_count = len(union_new_all)

    # fail_to_pass.append('astropy/io/fits/tests/test_connect.py::test_testenhancer_failing')
    # fail_to_pass.extend([
    #     'astropy/io/fits/tests/test_connect.py::test_is_fits_with_invalid_extensions',
    #     'astropy/io/fits/tests/test_connect.py::test_is_fits_with_no_extension',
    # ])

    return TestSpec(
        instance_id=instance_id,
        repo=repo,
        env_script_list=env_script_list,
        repo_script_list=repo_script_list,
        eval_script_list=eval_script_list,
        version=version,
        arch=arch,
        FAIL_TO_PASS=fail_to_pass,
        PASS_TO_PASS=pass_to_pass,
        language=MAP_REPO_TO_EXT[repo],
        docker_specs=docker_specs,
        namespace=namespace,
        base_image_tag=base_image_tag,
        env_image_tag=env_image_tag,
        instance_image_tag=instance_image_tag,
        new_te_tests_count=new_te_tests_count,
    )
