"""
Sandboxed Tool Executor

Executes synthesized tools in an isolated environment with
scoped credentials and resource limits.
"""

import asyncio
import os
import sys
import tempfile
import time
from pathlib import Path
from typing import Any

from oats.core.types import (
    GroundableOutput,
    SynthesizedTool,
    ToolOutput,
)


class SandboxConfig:
    """Configuration for the execution sandbox."""

    def __init__(
        self,
        timeout_seconds: int = 30,
        max_memory_mb: int = 512,
        allowed_network: bool = True,
        allowed_domains: list[str] | None = None,
        working_dir: Path | None = None,
        env_allowlist: list[str] | None = None,
    ) -> None:
        self.timeout_seconds = timeout_seconds
        self.max_memory_mb = max_memory_mb
        self.allowed_network = allowed_network
        self.allowed_domains = allowed_domains or []
        self.working_dir = working_dir or Path(tempfile.gettempdir())
        self.env_allowlist = env_allowlist or []


class CredentialProvider:
    """
    Provides credentials for tool execution.

    CRITICAL: Credentials are NEVER embedded in tool code.
    They are injected as environment variables at execution time.
    """

    def __init__(self) -> None:
        self._credentials: dict[str, str] = {}

    def register_credential(self, scope: str, env_var: str) -> None:
        """
        Register a credential for a scope.

        Args:
            scope: The auth scope (e.g., "github:read")
            env_var: Environment variable containing the credential
        """
        self._credentials[scope] = env_var

    def get_env_for_scopes(self, scopes: list[str]) -> dict[str, str]:
        """
        Get environment variables for the requested scopes.

        Returns a dict of env vars to inject into the sandbox.
        """
        env = {}
        for scope in scopes:
            if scope in self._credentials:
                env_var = self._credentials[scope]
                value = os.environ.get(env_var)
                if value:
                    env[env_var] = value
        return env

    def has_scope(self, scope: str) -> bool:
        """Check if a scope is available."""
        return scope in self._credentials


class Executor:
    """
    Executes synthesized tools in a sandbox.

    The executor:
    1. Validates the tool code
    2. Sets up an isolated environment
    3. Injects scoped credentials
    4. Executes with resource limits
    5. Captures output and errors
    6. Extracts groundable content
    """

    def __init__(
        self,
        config: SandboxConfig | None = None,
        credential_provider: CredentialProvider | None = None,
    ) -> None:
        self.config = config or SandboxConfig()
        self.credentials = credential_provider or CredentialProvider()

    async def execute(
        self,
        tool: SynthesizedTool,
        context: dict[str, Any] | None = None,
    ) -> ToolOutput:
        """
        Execute a synthesized tool in a sandbox.

        Args:
            tool: The synthesized tool to execute
            context: Additional context to pass to the tool

        Returns:
            ToolOutput with results or errors
        """
        start_time = time.time()

        try:
            # Step 1: Validate code (basic safety checks)
            validation_error = self._validate_code(tool.code)
            if validation_error:
                return ToolOutput(
                    tool_id=tool.id,
                    success=False,
                    error=f"Code validation failed: {validation_error}",
                    execution_time_ms=int((time.time() - start_time) * 1000),
                )

            # Step 2: Check credential availability
            missing_scopes = self._check_scopes(tool.requested_scopes)
            if missing_scopes:
                return ToolOutput(
                    tool_id=tool.id,
                    success=False,
                    error=f"Missing credentials for scopes: {', '.join(missing_scopes)}",
                    execution_time_ms=int((time.time() - start_time) * 1000),
                )

            # Step 3: Prepare execution environment
            env = self._prepare_environment(tool.requested_scopes)

            # Step 4: Execute in sandbox
            result, stdout, stderr = await self._execute_in_sandbox(
                tool.code,
                context or {},
                env,
            )

            execution_time = int((time.time() - start_time) * 1000)

            # Step 5: Extract groundable output
            groundable = self._extract_groundable(tool, result)

            return ToolOutput(
                tool_id=tool.id,
                success=True,
                result=result,
                stdout=stdout,
                stderr=stderr,
                execution_time_ms=execution_time,
                groundable=groundable,
            )

        except TimeoutError:
            return ToolOutput(
                tool_id=tool.id,
                success=False,
                error=f"Execution timed out after {self.config.timeout_seconds}s",
                execution_time_ms=int((time.time() - start_time) * 1000),
            )
        except Exception as e:
            return ToolOutput(
                tool_id=tool.id,
                success=False,
                error=str(e),
                execution_time_ms=int((time.time() - start_time) * 1000),
            )

    def _validate_code(self, code: str) -> str | None:
        """
        Basic code validation.

        Returns error message if invalid, None if OK.
        """
        # Check for obviously dangerous patterns
        dangerous_patterns = [
            ("os.system(", "Direct system calls not allowed"),
            ("subprocess.call(", "Use subprocess.run with shell=False"),
            ("eval(", "eval() not allowed"),
            # Note: "exec(" is allowed in create_subprocess_exec context
            ("exec(context", "exec() not allowed for arbitrary code"),
            ("__import__(", "Dynamic imports not allowed"),
            ("open('/etc", "System file access not allowed"),
            ("open('/root", "Root directory access not allowed"),
        ]

        for pattern, message in dangerous_patterns:
            if pattern in code:
                return message

        # Verify it has the expected function signature
        if "async def execute" not in code:
            return "Code must define 'async def execute(context: dict)'"

        return None

    def _check_scopes(self, scopes: list[str]) -> list[str]:
        """Check which scopes are missing credentials."""
        import re
        # Filter out empty strings, "none", "(none)", "N/A", etc.
        not_a_scope = {"none", "(none)", "n/a", "(n/a)", "-", "null", "no scopes", "no auth"}
        real_scopes = [
            s for s in scopes
            if s and s.lower().strip("() ") not in not_a_scope
            and not re.match(r'^\W+$', s)  # skip pure punctuation
        ]

        # Scopes that don't require credentials (public access or local)
        no_creds_required = {
            "http:get", "http:post", "http:put", "http:delete",  # Public HTTP
            "http", "http:request",  # Generic HTTP
            "json", "datetime", "data",  # Local processing
            "shell", "shell:execute", "shell:run",  # Local shell (uses env creds)
            "filesystem", "filesystem:read", "filesystem:write",  # Local ops
            # Cloud provider scopes — ambient creds from env/CLI configs
            "compute.readonly", "compute.read_only", "compute",
            "storage.readonly", "storage.read_only", "storage",
            "iam.readonly", "iam", "iam:read",
            "gcp", "gcp:read", "gcp:compute", "gcp:storage", "gcp:iam",
            "aws", "aws:read", "aws:s3", "aws:ec2", "aws:iam",
            "azure", "azure:read",
            "kubernetes", "kubernetes:read", "k8s", "k8s:read",
            "github", "github:read", "github:api",
        }

        # Also treat any scope containing known cloud provider patterns as ambient
        cloud_patterns = {"compute.", "storage.", "iam.", "s3.", "ec2.", "rds.",
                          "lambda.", "sqs.", "sns.", "cloudwatch.", "logs.",
                          "bigquery.", "pubsub.", "cloudsql.", "gke.", "aks.",
                          "eks.", "kms.", "secretsmanager.", "ssm.", "route53."}

        # URL-format scopes from cloud providers — these use ambient CLI credentials
        # (az login, gcloud auth, aws configure) so no explicit creds needed
        cloud_url_patterns = [
            "management.azure.com",
            "graph.microsoft.com",
            "vault.azure.net",
            "database.windows.net",
            "blob.core.windows.net",
            "table.core.windows.net",
            "queue.core.windows.net",
            "servicebus.windows.net",
            "azurecr.io",
            "cognitiveservices.azure.com",
            "search.windows.net",
            "googleapis.com",
            "amazonaws.com",
        ]

        def is_ambient(scope: str) -> bool:
            s = scope.lower()
            if s in no_creds_required:
                return True
            for pattern in cloud_patterns:
                if pattern in s:
                    return True
            # URL-format scopes (e.g. https://management.azure.com/.default)
            for url_pattern in cloud_url_patterns:
                if url_pattern in s:
                    return True
            return False

        needs_creds = [s for s in real_scopes if not is_ambient(s)]
        return [s for s in needs_creds if not self.credentials.has_scope(s)]

    def _prepare_environment(self, scopes: list[str]) -> dict[str, str]:
        """Prepare environment variables for execution."""
        # Start with minimal base environment
        env = {
            "PATH": os.environ.get("PATH", "/usr/bin:/bin"),
            "HOME": os.environ.get("HOME", str(self.config.working_dir)),
            "TMPDIR": str(self.config.working_dir),
        }

        # Always pass through cloud provider credentials if present
        cloud_vars = [
            "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN",
            "AWS_REGION", "AWS_DEFAULT_REGION", "AWS_PROFILE",
            "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET", "AZURE_TENANT_ID", "AZURE_SUBSCRIPTION_ID",
            "AZURE_CONFIG_DIR",  # Azure CLI config (~/.azure)
            "GOOGLE_APPLICATION_CREDENTIALS", "GCLOUD_PROJECT", "CLOUDSDK_CORE_PROJECT",
            "CLOUDSDK_CONFIG",  # gcloud config dir
            "GITHUB_TOKEN", "SLACK_TOKEN",
            "KUBECONFIG",  # kubectl config
        ]
        for var in cloud_vars:
            if var in os.environ:
                env[var] = os.environ[var]

        # Add allowlisted env vars
        for var in self.config.env_allowlist:
            if var in os.environ:
                env[var] = os.environ[var]

        # Add scoped credentials
        env.update(self.credentials.get_env_for_scopes(scopes))

        return env

    async def _execute_in_sandbox(
        self,
        code: str,
        context: dict[str, Any],
        env: dict[str, str],
    ) -> tuple[Any, str, str]:
        """
        Execute code in a sandboxed subprocess.

        For now, uses subprocess with restrictions.
        Future: Could use Deno, Firecracker, or WASM.
        """
        # Write the code to a temp file
        with tempfile.NamedTemporaryFile(
            mode="w",
            suffix=".py",
            delete=False,
            dir=self.config.working_dir,
        ) as f:
            # Wrap the code in a runner
            runner_code = f'''
import asyncio
import json
import sys

{code}

async def main():
    context = json.loads(sys.argv[1]) if len(sys.argv) > 1 else {{}}
    result = await execute(context)
    print("__OAT_RESULT__")
    print(json.dumps(result))

if __name__ == "__main__":
    asyncio.run(main())
'''
            f.write(runner_code)
            script_path = f.name

        try:
            # Execute with timeout
            process = await asyncio.wait_for(
                asyncio.create_subprocess_exec(
                    sys.executable,
                    script_path,
                    __import__("json").dumps(context),
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                    env=env,
                    cwd=str(self.config.working_dir),
                ),
                timeout=self.config.timeout_seconds,
            )

            stdout_bytes, stderr_bytes = await asyncio.wait_for(
                process.communicate(),
                timeout=self.config.timeout_seconds,
            )

            stdout = stdout_bytes.decode("utf-8", errors="replace")
            stderr = stderr_bytes.decode("utf-8", errors="replace")

            # Extract result from stdout
            result = None
            if "__OAT_RESULT__" in stdout:
                parts = stdout.split("__OAT_RESULT__")
                if len(parts) > 1:
                    result_str = parts[1].strip()
                    try:
                        result = __import__("json").loads(result_str)
                    except Exception:
                        result = result_str
                stdout = parts[0]  # Remove result marker from stdout

            return result, stdout, stderr

        finally:
            # Clean up temp file
            try:
                Path(script_path).unlink()
            except Exception:
                pass

    def _extract_groundable(
        self,
        tool: SynthesizedTool,
        result: Any,
    ) -> GroundableOutput | None:
        """Extract groundable content from the result."""
        if result is None:
            return None

        # Basic extraction - can be enhanced with LLM
        import json

        if isinstance(result, dict):
            summary = tool.human_explanation
            embedding_text = f"{tool.intent}\n\n{json.dumps(result, indent=2)}"
        else:
            summary = tool.human_explanation
            embedding_text = f"{tool.intent}\n\n{str(result)}"

        return GroundableOutput(
            summary=summary,
            entities=[],  # Could extract with NER
            embedding_text=embedding_text,
            metadata={
                "tool_id": tool.id,
                "intent": tool.intent,
                "capabilities": tool.capabilities_used,
            },
        )
