import os
import subprocess
import time
import psutil
import tempfile
import shutil

class CppCompiler:
    def __init__(self, workdir: str = None):
        self.workdir = workdir or tempfile.mkdtemp(prefix="algo_finder_")
        os.makedirs(self.workdir, exist_ok=True)

    def compile(self, filename: str, source_code: str) -> dict:
        if shutil.which("g++") is None:
            return {
                "compiled": False,
                "binary_path": None,
                "compile_error": "g++ not found. On Hugging Face Spaces add 'g++' to packages.txt; locally install a C++ toolchain (e.g. 'apt install g++').",
                "compile_stderr": "",
            }

        source_path = os.path.join(self.workdir, f"{filename}.cpp")
        binary_path = os.path.join(self.workdir, filename)
        if os.name == "nt" and not binary_path.lower().endswith(".exe"):
            binary_path = binary_path + ".exe"

        with open(source_path, "w", encoding="utf-8") as f:
            f.write(source_code)

        compile_cmd = ["g++", "-std=c++17", "-O2", source_path, "-o", binary_path]
        try:
            compile_result = subprocess.run(
                compile_cmd,
                capture_output=True,
                text=True,
                timeout=60,
            )
        except subprocess.TimeoutExpired:
            return {
                "compiled": False,
                "binary_path": None,
                "compile_error": "Compilation timed out (60s limit)",
                "compile_stderr": "",
            }
        except FileNotFoundError as e:
            return {
                "compiled": False,
                "binary_path": None,
                "compile_error": f"g++ executable not found: {e}. Install g++ (packages.txt on Spaces).",
                "compile_stderr": "",
            }
        except OSError as e:
            return {
                "compiled": False,
                "binary_path": None,
                "compile_error": f"Compilation failed to start: {e}",
                "compile_stderr": "",
            }

        if compile_result.returncode != 0:
            return {
                "compiled": False,
                "binary_path": None,
                "compile_error": compile_result.stderr.strip(),
                "compile_stderr": compile_result.stderr.strip(),
            }

        return {
            "compiled": True,
            "binary_path": binary_path,
            "compile_error": None,
            "compile_stderr": None,
        }

    def run(self, binary_path: str, timeout: int = 30) -> dict:
        if not binary_path or not os.path.exists(binary_path):
            return {
                "ran": False,
                "returncode": None,
                "stdout": "",
                "stderr": "Binary not found",
                "execution_time_s": None,
                "memory_kb": None,
            }

        process = psutil.Process()
        start_time = time.perf_counter()
        try:
            run_result = subprocess.run(
                [binary_path],
                capture_output=True,
                text=True,
                timeout=timeout,
            )
            end_time = time.perf_counter()
            execution_time = end_time - start_time

            try:
                mem_info = process.memory_info()
                memory_kb = mem_info.rss / 1024
            except Exception:
                memory_kb = None

            return {
                "ran": True,
                "returncode": run_result.returncode,
                "stdout": run_result.stdout.strip(),
                "stderr": run_result.stderr.strip(),
                "execution_time_s": round(execution_time, 6),
                "memory_kb": round(memory_kb, 2) if memory_kb is not None else None,
            }
        except subprocess.TimeoutExpired:
            return {
                "ran": False,
                "returncode": None,
                "stdout": "",
                "stderr": f"Execution timed out ({timeout}s limit)",
                "execution_time_s": timeout,
                "memory_kb": None,
            }
        except Exception as e:
            return {
                "ran": False,
                "returncode": None,
                "stdout": "",
                "stderr": str(e),
                "execution_time_s": None,
                "memory_kb": None,
            }

    def cleanup(self):
        if os.path.exists(self.workdir):
            shutil.rmtree(self.workdir, ignore_errors=True)