"""
Parallel Runner Library for Robot Framework

This library provides a `Run Parallel` keyword that:
1. Takes parallel branches as JSON
2. Creates a temporary Robot JSON file with branches as separate tests
3. Runs pabot with --testlevelsplit for true parallel execution
4. Returns results back to the main test
"""

import json
import os
import sys
import subprocess
import tempfile
import time
from robot.api.deco import keyword
from robot.api import logger
from robot.libraries.BuiltIn import BuiltIn
from typing import List, Dict, Any


class ParallelRunner:
    """Library to run parallel branches using pabot."""
    
    ROBOT_LIBRARY_SCOPE = 'SUITE'
    
    def __init__(self):
        self.temp_files = []
        self.temp_dirs = []
        self.results = {}
        self._project_root = None
    
    def _get_project_root(self) -> str:
        """Get the project root directory."""
        if self._project_root:
            return self._project_root
            
        # Try to get from environment
        pythonpath = os.environ.get("PYTHONPATH", "")
        if pythonpath:
            self._project_root = pythonpath.split(os.pathsep)[0]
            return self._project_root
        
        # Fall back to current directory
        self._project_root = os.getcwd()
        return self._project_root
    
    @keyword("Run Parallel")
    def run_parallel(self, branches_json: str, variables: str = "{}", listener: str = "", imports: str = "[]"):
        """
        Execute multiple branches in parallel using pabot.
        
        Args:
            branches_json: JSON string containing branches to run in parallel.
            variables: JSON string containing variables to pass to branches.
            listener: Optional listener class name (e.g. 'probe_listener.ProbeListener')
            imports: JSON string containing library imports for branches.
        
        Returns:
            Dictionary with results from all branches.
        """
        logger.info("=" * 50)
        logger.info("Starting parallel execution...")
        logger.info("=" * 50)
        
        # Parse inputs
        try:
            branches = json.loads(branches_json) if isinstance(branches_json, str) else branches_json
            vars_dict = json.loads(variables) if isinstance(variables, str) else variables
            imports_list = json.loads(imports) if isinstance(imports, str) else imports
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse JSON: {e}")
            raise
        
        logger.info(f"Running {len(branches)} branches in parallel")
        for i, branch in enumerate(branches):
            branch_name = branch.get("name", f"Branch_{i+1}")
            body_count = len(branch.get("body", []))
            logger.info(f"  - {branch_name}: {body_count} keywords")
        
        # Get project root
        project_root = self._get_project_root()
        logger.info(f"Project root: {project_root}")
        
        # Create temporary JSON file IN PROJECT ROOT (not system temp)
        # This ensures pabot can find the RPA libraries
        temp_json = self._create_temp_test_file(branches, vars_dict, imports_list, project_root)
        self.temp_files.append(temp_json)
        
        # Get listener from environment if not provided
        if not listener:
            listener = os.environ.get("ROBOT_LISTENER", "")
        
        try:
            # Run pabot
            result = self._run_pabot(temp_json, project_root, listener)
            
            logger.info("=" * 50)
            logger.info(f"Parallel execution completed. Return code: {result['returncode']}")
            logger.info(f"Elapsed time: {result['elapsed_time']:.2f}s")
            logger.info("=" * 50)
            
            if result['returncode'] != 0:
                logger.warn("Some parallel branches may have failed.")
                if result.get('stderr'):
                    logger.warn(f"Error details: {result['stderr'][:500]}")
            
            return result
            
        finally:
            # Cleanup temp file
            self._cleanup_temp_files()
    
    def _create_temp_test_file(self, branches: List[Dict], variables: Dict, imports: List[Dict], project_root: str) -> str:
        """Create a temporary Robot JSON file with branches as tests."""
        
        tests = []
        for i, branch in enumerate(branches):
            branch_name = branch.get("name", f"Branch_{i+1}")
            branch_body = branch.get("body", [])
            
            # Convert body items to proper keyword format
            converted_body = []
            for item in branch_body:
                keyword_item = {
                    "type": "keyword",
                    "name": item.get("name", ""),
                    "args": item.get("args", [])
                }
                if "assign" in item:
                    keyword_item["assign"] = item["assign"]
                converted_body.append(keyword_item)
            
            tests.append({
                "name": branch_name,
                "body": converted_body
            })
        
        # Use imports from parameter, add BuiltIn if not present
        final_imports = list(imports) if imports else []
        import_names = [imp.get("name") for imp in final_imports]
        if "BuiltIn" not in import_names:
            final_imports.insert(0, {"type": "LIBRARY", "name": "BuiltIn"})
        
        # Build the complete JSON structure (flat, no nested suites)
        json_content = {
            "name": "Parallel Branches",
            "tests": tests,
            "resource": {
                "imports": final_imports,
                "variables": [
                    {
                        "name": name, 
                        "value": value if isinstance(value, list) else [str(value)]
                    }
                    for name, value in variables.items()
                ] if variables else []
            }
        }
        
        # Create temp file IN PROJECT ROOT so libraries can be found
        temp_file = os.path.join(project_root, f"_parallel_temp_{int(time.time() * 1000)}.json")
        
        with open(temp_file, "w", encoding="utf-8") as f:
            json.dump(json_content, f, indent=2, ensure_ascii=False)
        
        logger.info(f"Created temp file: {temp_file}")
        logger.debug(f"Temp file content:\n{json.dumps(json_content, indent=2)}")
        
        return temp_file
    
    def _run_pabot(self, test_file: str, project_root: str, listener: str = "") -> Dict[str, Any]:
        """Run pabot with the given test file."""
        
        # Build pabot command - use UNIQUE temp directory for output files to prevent conflicts
        unique_id = int(time.time() * 1000)
        temp_output_dir = os.path.join(project_root, f"_parallel_output_{unique_id}")
        os.makedirs(temp_output_dir, exist_ok=True)
        self.temp_dirs.append(temp_output_dir)  # Track for cleanup
        
        executable = sys.executable
        
        cmd = [
            executable, "-m", "pabot.pabot",
            "--testlevelsplit",
            "--outputdir", temp_output_dir,
        ]
        
        # Add listener if provided
        if listener:
            cmd.extend(["--listener", listener])
            logger.info(f"Using listener: {listener}")
        
        cmd.append(test_file)
        
        logger.info(f"Command: {' '.join(cmd)}")
        logger.info(f"Working directory: {project_root}")
        
        # Set PYTHONPATH to include project root for libraries
        env = os.environ.copy()
        env["PYTHONPATH"] = project_root
        
        # Execute pabot
        start_time = time.time()
        
        try:
            result = subprocess.run(
                cmd,
                cwd=project_root,
                env=env,
                capture_output=True,
                text=True,
                timeout=300  # 5 minute timeout
            )
            
            elapsed_time = time.time() - start_time
            
            # Log output
            if result.stdout:
                # Truncate long output
                stdout = result.stdout
                if len(stdout) > 1000:
                    stdout = stdout[:500] + "\n...[truncated]...\n" + stdout[-500:]
                logger.info(f"STDOUT:\n{stdout}")
            
            if result.stderr:
                stderr = result.stderr
                if len(stderr) > 1000:
                    stderr = stderr[:500] + "\n...[truncated]...\n" + stderr[-500:]
                logger.warn(f"STDERR:\n{stderr}")
            
            return {
                "returncode": result.returncode,
                "stdout": result.stdout,
                "stderr": result.stderr,
                "elapsed_time": elapsed_time
            }
            
        except subprocess.TimeoutExpired:
            logger.error("Pabot execution timed out after 5 minutes")
            return {
                "returncode": -1,
                "stdout": "",
                "stderr": "Timeout expired",
                "elapsed_time": 300
            }
        except Exception as e:
            logger.error(f"Failed to run pabot: {e}")
            return {
                "returncode": -1,
                "stdout": "",
                "stderr": str(e),
                "elapsed_time": time.time() - start_time
            }
    
    def _cleanup_temp_files(self):
        """Remove temporary files and directories."""
        import shutil
        
        # Cleanup temp JSON files
        for temp_file in self.temp_files:
            try:
                if os.path.exists(temp_file):
                    os.remove(temp_file)
                    logger.info(f"Cleaned up temp file: {os.path.basename(temp_file)}")
            except Exception as e:
                logger.warn(f"Failed to cleanup {temp_file}: {e}")
        
        self.temp_files = []
        
        # Cleanup temp output directories
        for output_dir in self.temp_dirs:
            if os.path.exists(output_dir):
                try:
                    shutil.rmtree(output_dir)
                    logger.info(f"Cleaned up parallel output directory: {os.path.basename(output_dir)}")
                except Exception as e:
                    logger.warn(f"Failed to cleanup output dir {output_dir}: {e}")
        
        self.temp_dirs = []
    
    @keyword("Get Parallel Results")
    def get_parallel_results(self) -> Dict:
        """Get results from the last parallel execution."""
        return self.results
