from src.environment.environment_registry import environment_registry

@environment_registry.register_environment_method(
    "custom_environment",
    "execute_node"
)
async def execute_node(self: "BaseEnvironment", **kwargs) -> Dict:
    """Custom environment execution logic"""
    if not self._initialized:
        await self.initialize()
    
    # Custom evaluation workflow
    results = await self._run_custom_workflow(**kwargs)
    
    return {"success": True, "data": results}

# Automatic discovery in environment/ directory
    def discover_environments(self, environments_directory: Optional[str] = None) -> List[str]:
        """
        Discover all environment modules in the specified directory.

        Args:
            environments_directory (Optional[str]): Optional custom directory for environment modules

        Returns:
            List[str]: List of discovered environment module names
        """
        if environments_directory is None:
            # default to the directory containing this file
            environments_directory = Path(__file__).parent
        else:
            environments_directory = Path(environments_directory)

        discovered_environments = []

        # look for environment_*.py files
        for file_path in environments_directory.glob("environment_*.py"):
            if "environment_registry.py" in file_path.name:
                continue
            if file_path.is_file() and not file_path.name.startswith("__"):
                module_name = file_path.stem
                discovered_environments.append(module_name)

        logger.info(f"Discovered environments: {discovered_environments}")
        return discovered_environments

class BaseEnvironment:
    def __init__(self, environment_name: str, **kwargs):
        self.environment_name = environment_name
        self.config = self._load_config()
        self.agents = {}
        self._initialized = False
        self.output_dir = self._create_temp_dir()
    
    async def initialize(self) -> None:
        """Initialize agents and resources"""
        await self._initialize_agents()
        self._initialized = True
    
    async def execute_node(self, **kwargs) -> Dict:
        """Override in specific environments"""
        raise NotImplementedError
    
    async def reset(self) -> None:
        """Cleanup resources"""
        await self._cleanup_temp_files()
        await self._reset_agents()

async def _initialize_agents(self):
    """Initialize all configured agents"""
    for agent_name in self.config["agents"]:
        session_id = await self.interface_client.initialize_session(
            role=agent_name
        )
        self.agents[agent_name] = {
            "session_id": session_id,
            "initialized": True
        }

async def _reset_agents(self):
    """Reset all agent sessions"""
    for agent_name, agent_info in self.agents.items():
        if agent_info.get("session_id"):
            await self.interface_client.cleanup_session(
                agent_info["session_id"]
            )

def run_script(script_path: str) -> Tuple[bool, str]:
    """
    Run a Python script and capture its output.

    Args:
        script_path (str): Path to the script to run

    Returns:
        Tuple[bool, str]: Success status and output/error message
    """
    try:
        result = subprocess.run(
            ["python", script_path],
            capture_output=True,
            text=True,
            check=False,
        )

        if result.returncode == 0:
            return True, result.stdout or "All tests passed."
        else:
            return False, result.stderr or result.stdout

    except Exception as e:
        logger.error(f"Error running script: {e}")
        return False, str(e)

async def interact(self, **kwargs) -> Optional[str]:
        """
        Interact with the LLM service.

        Args:
            **kwargs: Input data for the interaction

        Returns:
            Optional[str]: Response from the LLM
        """
        if not self._initialized:
            await self.initialize()

        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                # Start the interaction
                logger.info(f"Interacting with {self.role} with {kwargs} - {self.session_id}")
                response = await client.post(
                    f"{self.base_url}/interact",
                    json={
                        "session_id": self.session_id,
                        "input_data": kwargs,
                        "use_agent": False,
                    },
                )
                response.raise_for_status()
                data = response.json()
                task_id = data.get("task_id")

                # Poll for results
                while True:
                    await asyncio.sleep(0.5)  # Wait before checking status
                    status_response = await client.get(f"{self.base_url}/task_status/{task_id}")
                    logger.info(f"Status response: {status_response}")
                    status_response.raise_for_status()
                    status_data = status_response.json()

                    if status_data["status"] == "completed":
                        return status_data["result"]["response"]
                    elif status_data["status"] == "failed":
                        logger.error(f"Task failed: {status_data.get('error')}")
                        return None

        except Exception as e:
            logger.opt(exception=e).error(f"Error in LLM interaction: {e}")
            return None

async def _run_agent_sequence(self, workflow_steps: List[Dict]) -> Dict:
    """Execute agents in sequence with data passing"""
    context = {}
    
    for step in workflow_steps:
        agent_name = step["agent"]
        input_template = step["input_template"]
        
        # Format input with previous context
        formatted_input = self._format_template(input_template, context)
        
        # Execute agent
        result = await self._execute_agent(agent_name, formatted_input)
        
        # Update context with result
        context[step["output_key"]] = result
    
    return context

async def _run_parallel_agents(self, parallel_tasks: List[Dict]) -> Dict:
    """Execute multiple agents concurrently"""
    tasks = []
    for task in parallel_tasks:
        agent_task = self._execute_agent(
            task["agent"],
            task["input_data"]
        )
        tasks.append(agent_task)
    
    results = await asyncio.gather(*tasks, return_exceptions=True)
    return dict(zip([task["output_key"] for task in parallel_tasks], results))

class EnvironmentException(Exception):
    """Base environment exception"""
    pass

class EnvironmentExecutionException(EnvironmentException):
    """Execution-related errors"""
    pass

class AgentCommunicationException(EnvironmentException):
    """Agent interaction errors"""
    pass

class ConfigurationException(EnvironmentException):
    """Configuration-related errors"""
    pass

async def health_check(self) -> Dict:
    """Check environment health"""
    health_status = {
        "environment": self.environment_name,
        "status": "healthy",
        "checks": {}
    }
    
    # Check agent availability
    for agent_name in self.config["agents"]:
        try:
            await self._ping_agent(agent_name)
            health_status["checks"][f"agent_{agent_name}"] = "healthy"
        except Exception as e:
            health_status["checks"][f"agent_{agent_name}"] = f"unhealthy: {str(e)}"
            health_status["status"] = "degraded"
    
    # Check resource availability
    if not self._check_disk_space():
        health_status["checks"]["disk_space"] = "low"
        health_status["status"] = "degraded"
    
    return health_status