import os
import subprocess
import tempfile
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Tuple

import yaml


class Policy(ABC):
    """
    Abstract base class for code policies.
    """

    @abstractmethod
    def check(self, python_code: str) -> str | None:
        """
        Returns a string describing the violation, or None if no violation.
        """
        pass


class PltSavefigPolicy(Policy):
    """
    Policy that checks for a plt.savefig call in the source code.
    """

    def check(self, python_code: str) -> str | None:
        code_no_comments_no_strings = remove_comments_and_strings(python_code)
        if ".savefig" not in code_no_comments_no_strings:
            return (
                "No savefig call found in source code - "
                "save your plot to a file using plt.savefig() or fig.savefig()."
            )
        return None


class NFilesPolicy(Policy):
    """
    Policy that checks for 'NFiles=1' outside of comments and strings.
    """

    def check(self, python_code: str) -> str | None:
        import re

        code_no_comments_no_strings = remove_comments_and_strings(python_code)
        # Match NFiles followed by optional spaces, =, optional spaces, and 1
        if not re.search(r"NFiles\s*=\s*1", code_no_comments_no_strings):
            return (
                "NFiles=1 not found in source code - it must be present in the ServiceX "
                "`Sample` definition to assure a quick test run."
            )
        return None


# Global list of policies
POLICIES: List[Policy] = [NFilesPolicy(), PltSavefigPolicy()]


def copy_servicex_yaml_if_exists(target_dir: str):
    """
    Copies servicex.yaml from the user's home directory to target_dir if it exists.
    """

    cwd_servicex = os.path.join(os.getcwd(), "servicex.yaml")
    home_servicex = os.path.expanduser("~/servicex.yaml")
    target_servicex = os.path.join(target_dir, "servicex.yaml")
    source_servicex = None
    if os.path.exists(cwd_servicex):
        source_servicex = cwd_servicex
    elif os.path.exists(home_servicex):
        source_servicex = home_servicex
    if source_servicex:
        with open(source_servicex, "r", encoding="utf-8") as f:
            loaded = yaml.safe_load(f)
        # If loaded is not a dict, make it a dict
        if not isinstance(loaded, dict):
            loaded = {}  # treat as empty config
        loaded["cache_path"] = "/cache"
        with open(target_servicex, "w", encoding="utf-8") as f:
            yaml.safe_dump(loaded, f)


@dataclass
class DockerRunResult:
    """
    Contains the output and results of running a python script in docker.
    png_files: list of (filename, file bytes)
    """

    stdout: str
    stderr: str
    elapsed: float
    png_files: List[Tuple[str, bytes]]
    exit_code: int


def run_python_in_docker(
    python_code: str, docker_image: str = "atlasplotagent:latest"
) -> DockerRunResult:
    """
    Runs the given python_code in a Docker container, captures stdout/stderr, elapsed time,
    and PNG outputs.

    Returns a DockerRunResult dataclass.
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        script_path = os.path.join(temp_dir, "script.py")
        with open(script_path, "w", encoding="utf-8") as f:
            f.write(python_code)

        # Copy servicex.yaml from home directory if it exists
        copy_servicex_yaml_if_exists(temp_dir)

        # Run the docker container
        container_dir = "/app"
        # Mount a docker volume at /cache
        cache_volume = "atlasplotagent_servicex_cache"
        command = [
            "docker",
            "run",
            "--rm",
            "-v",
            f"{temp_dir}:{container_dir}",
            "-v",
            f"{cache_volume}:/cache",
            docker_image,
            "bash",
            "-i",  # so it executes the `.bashrc` file which defines the venv
            "-c",
            f"python {container_dir}/script.py",
        ]

        start = time.time()
        proc = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            encoding="utf-8",
        )

        stdout, stderr = proc.communicate()
        elapsed = time.time() - start
        exit_code = proc.returncode
        # Check for Docker connection error
        if "docker: error during connect" in stderr:
            raise RuntimeError(f"Docker connection error: {stderr}")

        # Find PNG files in temp_dir and load them into memory
        png_files = []
        for p in Path(temp_dir).glob("*.png"):
            with open(p, "rb") as f:
                png_files.append((p.name, f.read()))

        # No need to manually clean up temp_dir; TemporaryDirectory handles it

        return DockerRunResult(
            stdout=stdout,
            stderr=stderr,
            elapsed=elapsed,
            png_files=png_files,
            exit_code=exit_code,
        )


def remove_comments_and_strings(python_code: str) -> str:
    """
    Utility function to remove comments and strings from Python code.

    Args:
        python_code: The Python source code as a string

    Returns:
        The code with comments and strings removed
    """
    import re

    # First, search for any "```python" and only keep the text between those if we find those.
    code_blocks = re.findall(r"```python(.*?)```", python_code, re.DOTALL)
    if code_blocks:
        python_code = "\n".join(code_blocks)

    # Process line by line to handle comments and strings properly
    result_lines = []

    for line in python_code.splitlines():
        stripped = line.strip()

        # Skip empty lines and full-line comments
        if not stripped or stripped.startswith("#"):
            continue

        # For lines that may have inline comments, we need to be careful
        # First, let's try to identify if there's a # outside of strings
        processed_line = line

        # Simple approach: remove inline comments by finding # that's not in quotes
        # This is a simplified approach that works for most cases
        in_single_quote = False
        in_double_quote = False
        escape_next = False
        comment_start = -1

        for i, char in enumerate(line):
            if escape_next:
                escape_next = False
                continue

            if char == "\\":
                escape_next = True
                continue

            if char == "'" and not in_double_quote:
                in_single_quote = not in_single_quote
            elif char == '"' and not in_single_quote:
                in_double_quote = not in_double_quote
            elif char == "#" and not in_single_quote and not in_double_quote:
                comment_start = i
                break

        # Remove inline comment if found
        if comment_start >= 0:
            processed_line = line[:comment_start].rstrip()

        # Skip if the line becomes empty after removing comments
        if not processed_line.strip():
            continue

        result_lines.append(processed_line)

    # Join lines back together
    code_no_comments = "\n".join(result_lines)

    # Now remove strings from the comment-free code
    def remove_strings(s):
        s = re.sub(r"'''[\s\S]*?'''", "", s)
        s = re.sub(r'"""[\s\S]*?"""', "", s)
        s = re.sub(r"'(?:\\.|[^'])*'", "", s)
        s = re.sub(r'"(?:\\.|[^"])*"', "", s)
        return s

    code_no_strings = remove_strings(code_no_comments)

    # Clean up any remaining empty lines
    final_lines = []
    for line in code_no_strings.splitlines():
        line = line.rstrip()
        if line:
            final_lines.append(line)

    return "\n".join(final_lines)


def check_code_policies(
    python_code: str, policies: Optional[List[Policy]] = None
) -> bool | DockerRunResult:
    """
    Checks the provided Python code against a list of policy objects.

    Args:
        python_code: The Python source code to check.
        policies: Optional list of Policy objects to use for checking. If None, uses the default
            POLICIES.

    Returns:
        True if no policy violations are found.
        DockerRunResult with details of violations if any are found.
    """
    if policies is None:
        policies = POLICIES

    violations = []
    for policy in policies:
        violation = policy.check(python_code)
        if violation:
            violations.append(violation)
    if violations:
        markdown = "\n".join([f"- {v}" for v in violations])
        return DockerRunResult(
            stdout="",
            stderr=f"Policy violations found:\n{markdown}",
            elapsed=0,
            png_files=[],
            exit_code=1,
        )
    return True


def print_md_table_for_phased_usage_docker(
    fh_out, usage_info: List[Tuple[str, DockerRunResult]]
) -> float:
    """
    Print a markdown table for the usage information of each phase.
    """
    fh_out.write("\n### Docker Usage\n")
    fh_out.write("| Phase | Elapsed Time (seconds) |\n")
    fh_out.write("|-------|--------------|\n")
    for phase, info in usage_info:
        fh_out.write(f"| {phase} | {info.elapsed:.2f} |\n")

    # Totals
    total_elapsed = sum(info.elapsed for _, info in usage_info)
    fh_out.write(f"| **Total** | **{total_elapsed:.2f}** |\n")

    return total_elapsed
