import json
import shutil
import subprocess
import os
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
from .registry import ModeRegistry

class TestManager:
    """
    Core orchestrator for the stdcodetest framework.
    """

    def __init__(self, tests_root=None, mode="default"):
        """Initializes the manager with a root directory and a specific testing mode."""
        self.tests_base = Path(tests_root or "TEST_SUITE").resolve()
        self.tests_base.mkdir(parents=True, exist_ok=True)
        self.mode = mode

    def create_std_test(self, test_name: str, replace: bool = False) -> Path:
        """Scaffolds a new test directory structure relative to the tests_base."""
        target_path = self.tests_base / test_name
        return self._build_structure(target_path, replace)

    def create_std_test_from_path(self, absolute_path: str | Path, replace: bool = False) -> Path:
        """Scaffolds a new test directory structure at a specific absolute location."""
        target_path = Path(absolute_path).resolve()
        return self._build_structure(target_path, replace)

    def _build_structure(self, test_dir: Path, replace: bool) -> Path:
        """Shared internal logic to create folders and inject config.json."""
        if replace and test_dir.exists():
            shutil.rmtree(test_dir)
        
        folders = ModeRegistry.get_folders(self.mode)
        for folder in folders:
            (test_dir / folder).mkdir(parents=True, exist_ok=True)

        meta_folder = next((f for f in folders if "META_CONFIG" in f), "2-META_CONFIG")
        config = {
            "std_version": "0.2.1",
            "mode": self.mode,
            "execution_meta": {"engine": "stdcodetest-core-v1", "env_requirements": self._get_env_hint()}
        }
        (test_dir / meta_folder / "config.json").write_text(json.dumps(config, indent=4))
        return test_dir

    def execute_test(self, test_dir: Path):
        """Navigates to '0_GEN' and runs 'EXECUTE.py' or 'EXECUTE.sh'."""
        gen_path = test_dir / "0_GEN"
        if not gen_path.exists():
            return f"Error: {gen_path} not found."

        py_script = gen_path / "EXECUTE.py"
        sh_script = gen_path / "EXECUTE.sh"

        try:
            if py_script.exists():
                return subprocess.run(["python", str(py_script)], cwd=gen_path, capture_output=True, text=True)
            elif sh_script.exists():
                return subprocess.run(["bash", str(sh_script)], cwd=gen_path, capture_output=True, text=True)
            else:
                return f"No EXECUTE script found in {gen_path}"
        except Exception as e:
            return f"Execution failed: {str(e)}"

    def get_mode_definitions(self, modes: str | list | None = None) -> dict:
            """
            Retrieves the folder structure definitions from the registry.
            
            :param modes: 
                - If None: Returns all mode definitions.
                - If str: Returns the definition for that specific mode.
                - If list: Returns definitions for all modes in the list.
            :return: A dictionary mapping mode names to their folder lists.
            """
            if modes is None:
                # Return everything in the registry
                all_modes = ModeRegistry.list_modes()
                return {m: ModeRegistry.get_folders(m) for m in all_modes}
            
            if isinstance(modes, str):
                # Return a single mode
                return {modes: ModeRegistry.get_folders(modes)}
            
            if isinstance(modes, list):
                # Return a subset of modes
                return {m: ModeRegistry.get_folders(m) for m in modes}
            
            return {}

    def run_all_tests(self, parallel: bool = False, max_workers: int = 4):
        """Discovers all test cases in the suite and runs them."""
        test_dirs = [d for d in self.tests_base.iterdir() if d.is_dir()]
        
        if parallel:
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                results = list(executor.map(self.execute_test, test_dirs))
        else:
            results = [self.execute_test(d) for d in test_dirs]
        
        return results

    def get_help(self):
        """
        Returns an illustrative help guide with usage examples.
        """
        modes = ", ".join(ModeRegistry.list_modes())
        help_text = f"""
stdcodetest Framework v0.2.1 - Illustrative Guide
{'='*50}
The TestManager automates the scaffolding and execution of AI evaluation pipelines.
It ensures every test follows a strict directory sequence (0_GEN -> 1-INPUT -> etc).

CURRENT CONFIGURATION:
- Mode: {self.mode}
- Root: {self.tests_base}
- Available Standards: {modes}

QUICK START EXAMPLE:
-------------------
    from stdcodetest.manager import TestManager

    # 1. Initialize the manager
    mgr = TestManager(tests_root="eval_suite", mode="pytorch2.0")

    # 2. Create a new test case
    test_path = mgr.create_std_test("nlp_classification_v1")
    print(f"Folders created at: {{test_path}}")

    # 3. Run all tests in the suite (Parallelized)
    results = mgr.run_all_tests(parallel=True, max_workers=8)
    
    # 4. Access individual execution
    mgr.execute_test(test_path)
    
    # 5. get the mode
    mgr.get_mode_definitions(['default', 'pytorch'])
    mgr.get_mode_definitions('default')

KEY METHODS:
- .create_std_test(name): Builds the directory tree in the root folder.
- .execute_test(path): Looks inside '0_GEN' for EXECUTE.py/sh and runs it.
- .run_all_tests(): Batch processor for the entire '{{self.tests_base}}' directory.
{'='*50}
"""
        return help_text

    def _get_env_hint(self) -> str:
        hints = {"vue": "node-18", "pytorch": "python-3.10"}
        return next((v for k, v in hints.items() if k in self.mode), "python-3.8")