"""
AsyncWorkflowBuilder - Async-first workflow development with enhanced ergonomics.

This module provides an async-optimized workflow builder with built-in patterns,
resource management integration, and type-safe construction helpers.
"""

import ast
import asyncio
import inspect
import textwrap
import uuid
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Set, TypeVar, Union

from ..nodes.base import Node
from ..resources.registry import ResourceFactory, ResourceRegistry
from .builder import WorkflowBuilder
from .graph import Workflow

T = TypeVar("T")


@dataclass
class RetryPolicy:
    """Retry policy configuration for async nodes."""

    max_attempts: int = 3
    initial_delay: float = 1.0
    backoff_factor: float = 2.0
    max_delay: float = 60.0
    retry_exceptions: Optional[List[str]] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "max_attempts": self.max_attempts,
            "initial_delay": self.initial_delay,
            "backoff_factor": self.backoff_factor,
            "max_delay": self.max_delay,
            "retry_exceptions": self.retry_exceptions,
        }


@dataclass
class ErrorHandler:
    """Error handler configuration for async nodes."""

    handler_type: str  # 'log', 'ignore', 'fallback', 'custom'
    fallback_value: Optional[Any] = None
    custom_handler: Optional[str] = None  # Code string for custom handler
    log_level: str = "error"


class AsyncWorkflowBuilder(WorkflowBuilder):
    """Async-optimized workflow builder with enhanced features."""

    def __init__(
        self,
        name: Optional[str] = None,
        resource_registry: Optional[ResourceRegistry] = None,
        description: Optional[str] = None,
    ):
        super().__init__()
        self.name = name or f"async_workflow_{uuid.uuid4().hex[:8]}"
        self.description = description
        self._resource_registry = resource_registry or ResourceRegistry()
        self._resource_requirements: Set[str] = set()
        self._error_handlers: Dict[str, ErrorHandler] = {}
        self._retry_policies: Dict[str, RetryPolicy] = {}
        self._node_metadata: Dict[str, Dict[str, Any]] = {}
        self._workflow_metadata: Dict[str, Any] = {
            "async_workflow": True,
            "builder_version": "1.0",
            "name": self.name,
            "description": description,
        }

    def add_async_code(
        self,
        node_id: str,
        code: str,
        *,
        timeout: int = 30,
        max_concurrent_tasks: int = 10,
        retry_policy: Optional[RetryPolicy] = None,
        error_handler: Optional[ErrorHandler] = None,
        required_resources: Optional[List[str]] = None,
        description: Optional[str] = None,
        **kwargs,
    ) -> "AsyncWorkflowBuilder":
        """Add async Python code node with enhanced configuration."""
        # Clean up code indentation
        code = textwrap.dedent(code).strip()

        # Validate code
        self._validate_async_code(code)

        # Track resource requirements
        if required_resources:
            self._resource_requirements.update(required_resources)
            self._node_metadata.setdefault(node_id, {})[
                "required_resources"
            ] = required_resources

        # Configure node
        config = {
            "code": code,
            "timeout": timeout,
            "max_concurrent_tasks": max_concurrent_tasks,
            **kwargs,
        }

        # Add description if provided
        if description:
            config["description"] = description
            self._node_metadata.setdefault(node_id, {})["description"] = description

        # Add node using base builder
        self.add_node("AsyncPythonCodeNode", node_id, config)

        # Configure error handling
        if retry_policy:
            self._retry_policies[node_id] = retry_policy
            self._node_metadata.setdefault(node_id, {})[
                "retry_policy"
            ] = retry_policy.to_dict()
        if error_handler:
            self._error_handlers[node_id] = error_handler
            self._node_metadata.setdefault(node_id, {})["error_handler"] = {
                "type": error_handler.handler_type,
                "fallback_value": error_handler.fallback_value,
            }

        return self  # Fluent interface

    def add_parallel_map(
        self,
        node_id: str,
        map_function: str,
        *,
        input_field: str = "items",
        output_field: str = "results",
        max_workers: int = 10,
        batch_size: Optional[int] = None,
        timeout_per_item: int = 5,
        continue_on_error: bool = False,
        description: Optional[str] = None,
    ) -> "AsyncWorkflowBuilder":
        """Add node that processes items in parallel using asyncio.gather."""
        # Validate function
        self._validate_async_function(map_function)

        code = f"""
import asyncio
from asyncio import Semaphore
import time

# Define processing function
{map_function}

# Validate function is defined
if 'process_item' not in locals():
    raise ValueError("map_function must define 'process_item' function")

# Create semaphore for concurrency control
semaphore = Semaphore({max_workers})

async def process_with_timeout(item, index):
    async with semaphore:
        start_time = time.time()
        try:
            # Check if process_item is async
            if asyncio.iscoroutinefunction(process_item):
                result = await asyncio.wait_for(
                    process_item(item),
                    timeout={timeout_per_item}
                )
            else:
                result = await asyncio.wait_for(
                    asyncio.create_task(asyncio.to_thread(process_item, item)),
                    timeout={timeout_per_item}
                )
            return {{
                "index": index,
                "success": True,
                "result": result,
                "duration": time.time() - start_time
            }}
        except asyncio.TimeoutError:
            return {{
                "index": index,
                "success": False,
                "error": "timeout",
                "item": item,
                "duration": time.time() - start_time
            }}
        except Exception as e:
            return {{
                "index": index,
                "success": False,
                "error": str(e),
                "error_type": type(e).__name__,
                "item": item,
                "duration": time.time() - start_time
            }}

# Get input items - check both direct field name and generate_output
input_items = None
if '{input_field}' in locals():
    input_items = {input_field}
elif 'generate_output' in locals() and isinstance(generate_output, dict):
    # When connected from another node, input might be in generate_output
    input_items = generate_output.get('{input_field}')
elif 'generate_output' in locals() and '{input_field}' == 'items':
    # Special case: if the output is directly the items list
    input_items = generate_output

if input_items is None:
    available_vars = list(locals().keys())
    raise ValueError(f"Input field '{input_field}' not found. Available: {{available_vars}}")

if not isinstance(input_items, (list, tuple)):
    raise ValueError(f"'{input_field}' must be a list or tuple, got {{type(input_items).__name__}}")

# Process items
total_start = time.time()
all_results = []

# Process in batches if specified
batch_size_val = {batch_size}
if batch_size_val:
    for i in range(0, len(input_items), batch_size_val):
        batch = input_items[i:i+batch_size_val]
        batch_results = await asyncio.gather(
            *[process_with_timeout(item, i+j) for j, item in enumerate(batch)],
            return_exceptions=True
        )
        # Filter out exceptions and convert to error results
        for j, result in enumerate(batch_results):
            if isinstance(result, Exception):
                all_results.append({{
                    "index": i+j,
                    "success": False,
                    "error": str(result),
                    "error_type": type(result).__name__,
                    "item": batch[j],
                    "duration": 0
                }})
            else:
                all_results.append(result)
else:
    batch_results = await asyncio.gather(
        *[process_with_timeout(item, i) for i, item in enumerate(input_items)],
        return_exceptions=True
    )
    # Filter out exceptions and convert to error results
    for i, result in enumerate(batch_results):
        if isinstance(result, Exception):
            all_results.append({{
                "index": i,
                "success": False,
                "error": str(result),
                "error_type": type(result).__name__,
                "item": input_items[i],
                "duration": 0
            }})
        else:
            all_results.append(result)

# Organize results
successful = [r for r in all_results if r.get("success", False)]
failed = [r for r in all_results if not r.get("success", False)]

# Extract processed items
processed_items = [r["result"] for r in successful]

# Continue on error flag
if not {continue_on_error} and failed:
    error_summary = {{
        "total_errors": len(failed),
        "error_types": {{}}
    }}
    for f in failed:
        error_type = f.get("error", "unknown")
        error_summary["error_types"][error_type] = error_summary["error_types"].get(error_type, 0) + 1

    raise RuntimeError(f"Processing failed for {{len(failed)}} items: {{error_summary}}")

result = {{
    "{output_field}": processed_items,
    "statistics": {{
        "total": len(input_items),
        "successful": len(successful),
        "failed": len(failed),
        "total_duration": time.time() - total_start,
        "average_duration": sum(r["duration"] for r in all_results) / len(all_results) if all_results else 0
    }},
    "errors": failed if failed else []
}}
"""

        return self.add_async_code(
            node_id,
            code,
            max_concurrent_tasks=max_workers,
            timeout=(
                timeout_per_item * len(input_field)
                if hasattr(input_field, "__len__")
                else 300
            ),
            description=description
            or f"Parallel map processing with {max_workers} workers",
        )

    def add_resource_node(
        self,
        node_id: str,
        resource_name: str,
        operation: str,
        params: Optional[Dict[str, Any]] = None,
        *,
        output_field: str = "result",
        description: Optional[str] = None,
        **kwargs,
    ) -> "AsyncWorkflowBuilder":
        """Add node that interacts with a registered resource."""
        # Track resource requirement
        self._resource_requirements.add(resource_name)

        # Build parameter string for operation call
        param_parts = []
        if params:
            for key, value in params.items():
                if isinstance(value, str):
                    param_parts.append(f'{key}="{value}"')
                else:
                    param_parts.append(f"{key}={repr(value)}")

        param_str = ", ".join(param_parts)

        code = f"""
# Access resource (function is provided by runtime)
if 'get_resource' in globals():
    resource = await get_resource("{resource_name}")
else:
    # Fallback for testing - resource should be in inputs
    resource = locals().get("{resource_name}")
    if resource is None:
        raise RuntimeError(f"Resource '{resource_name}' not available")

# Validate resource has the operation
if not hasattr(resource, "{operation}"):
    raise AttributeError(f"Resource '{resource_name}' does not have operation '{operation}'")

# Execute operation
operation_result = await resource.{operation}({param_str})

# Wrap result
result = {{
    "{output_field}": operation_result,
    "resource": "{resource_name}",
    "operation": "{operation}"
}}
"""

        return self.add_async_code(
            node_id,
            code,
            required_resources=[resource_name],
            description=description or f"Execute {operation} on {resource_name}",
            **kwargs,
        )

    def add_scatter_gather(
        self,
        scatter_id: str,
        process_id_prefix: str,
        gather_id: str,
        process_function: str,
        *,
        worker_count: int = 4,
        scatter_field: str = "items",
        gather_field: str = "results",
        description: Optional[str] = None,
    ) -> "AsyncWorkflowBuilder":
        """Add scatter-gather pattern for parallel processing."""
        # Use parallel_map which is simpler and more reliable
        return self.add_parallel_map(
            scatter_id,
            process_function,
            input_field=scatter_field,
            output_field=gather_field,
            max_workers=worker_count,
            description=description
            or f"Scatter-gather processing with {worker_count} workers",
        )

    def _validate_async_code(self, code: str):
        """Validate async Python code."""
        try:
            # Try to compile the code - but allow module-level await
            # by wrapping in an async function for validation
            wrapped_code = "async def __validate_wrapper():\n"
            for line in code.split("\n"):
                wrapped_code += f"    {line}\n"

            # Try to compile the wrapped version first
            try:
                compile(wrapped_code, "<string>", "exec")
            except SyntaxError:
                # If wrapped version fails, try original (might be valid module-level code)
                compile(code, "<string>", "exec")

        except SyntaxError as e:
            # Only reject if it's a true syntax error, not await-related
            if "await" not in str(e):
                raise ValueError(f"Invalid Python code: {e}")
            # For await-related errors, we'll allow them since AsyncPythonCodeNode handles module-level await

    def _validate_async_function(self, function_code: str):
        """Validate async function definition."""
        # Check if it defines process_item function
        if "def process_item" not in function_code:
            raise ValueError(
                "Function must define 'def process_item(item)' or 'async def process_item(item)'"
            )

        # Validate syntax
        self._validate_async_code(function_code)

    # Resource management methods
    def require_resource(
        self,
        name: str,
        factory: ResourceFactory,
        health_check: Optional[Callable] = None,
        cleanup_handler: Optional[Callable] = None,
        description: Optional[str] = None,
    ) -> "AsyncWorkflowBuilder":
        """Declare a required resource for this workflow."""
        # Register with resource registry
        self._resource_registry.register_factory(
            name, factory, health_check=health_check, cleanup_handler=cleanup_handler
        )

        # Track requirement
        self._resource_requirements.add(name)

        # Add to workflow metadata
        self._workflow_metadata.setdefault("resources", {})[name] = {
            "factory_type": type(factory).__name__,
            "description": description or f"Resource: {name}",
            "has_health_check": health_check is not None,
            "has_cleanup": cleanup_handler is not None,
        }

        return self

    def with_database(
        self,
        name: str = "db",
        host: str = "localhost",
        port: int = 5432,
        database: Optional[str] = None,
        user: Optional[str] = None,
        password: Optional[str] = None,
        min_size: int = 10,
        max_size: int = 20,
        **kwargs,
    ) -> "AsyncWorkflowBuilder":
        """Add database resource requirement."""
        from ..resources.factory import DatabasePoolFactory

        config = {
            "host": host,
            "port": port,
            "min_size": min_size,
            "max_size": max_size,
            **kwargs,
        }

        # Only add non-None values
        if database:
            config["database"] = database
        if user:
            config["user"] = user
        if password:
            config["password"] = password

        factory = DatabasePoolFactory(**config)

        # Health check for PostgreSQL
        async def pg_health_check(pool):
            try:
                async with pool.acquire() as conn:
                    await conn.fetchval("SELECT 1")
                return True
            except Exception:
                return False

        # Cleanup handler
        async def pg_cleanup(pool):
            await pool.close()

        return self.require_resource(
            name,
            factory,
            health_check=pg_health_check,
            cleanup_handler=pg_cleanup,
            description=f"PostgreSQL database connection pool to {host}:{port}/{database or 'default'}",
        )

    def with_http_client(
        self,
        name: str = "http",
        base_url: Optional[str] = None,
        headers: Optional[Dict[str, str]] = None,
        timeout: int = 30,
        **kwargs,
    ) -> "AsyncWorkflowBuilder":
        """Add HTTP client resource requirement."""
        from ..resources.factory import HttpClientFactory

        config = {"timeout": timeout, **kwargs}

        if headers:
            config["headers"] = headers

        factory = HttpClientFactory(base_url=base_url, **config)

        # Cleanup handler for aiohttp
        async def http_cleanup(session):
            await session.close()

        return self.require_resource(
            name,
            factory,
            cleanup_handler=http_cleanup,
            description="HTTP client session"
            + (f" for {base_url}" if base_url else ""),
        )

    def with_cache(
        self,
        name: str = "cache",
        backend: str = "redis",
        host: str = "localhost",
        port: int = 6379,
        **kwargs,
    ) -> "AsyncWorkflowBuilder":
        """Add cache resource requirement."""
        if backend == "redis":
            from ..resources.factory import CacheFactory

            factory = CacheFactory(backend=backend, host=host, port=port, **kwargs)

            # Health check for Redis
            async def redis_health_check(cache):
                try:
                    await cache.ping() if hasattr(cache, "ping") else True
                    return True
                except Exception:
                    return False

            # Cleanup handler
            async def redis_cleanup(cache):
                if hasattr(cache, "close"):
                    cache.close()
                if hasattr(cache, "wait_closed"):
                    await cache.wait_closed()

            return self.require_resource(
                name,
                factory,
                health_check=redis_health_check,
                cleanup_handler=redis_cleanup,
                description=f"Redis cache connection to {host}:{port}",
            )
        else:
            raise ValueError(f"Unsupported cache backend: {backend}")

    def build(self, workflow_id: str | None = None, **kwargs: Any) -> Workflow:
        """Build the async workflow with enhanced metadata."""
        # Add resource requirements to workflow metadata
        self._workflow_metadata["required_resources"] = list(
            self._resource_requirements
        )
        self._workflow_metadata["node_metadata"] = self._node_metadata

        # Build base workflow
        workflow = super().build(workflow_id, **kwargs)

        # Enhance workflow with async metadata
        if hasattr(workflow, "metadata"):
            workflow.metadata.update(self._workflow_metadata)
        else:
            workflow.metadata = self._workflow_metadata

        # Attach resource registry to workflow (stored in metadata for type safety)
        workflow.metadata["_resource_registry"] = self._resource_registry

        return workflow

    def get_resource_registry(self) -> ResourceRegistry:
        """Get the resource registry for this workflow."""
        return self._resource_registry

    def list_required_resources(self) -> List[str]:
        """List all required resources for this workflow."""
        return list(self._resource_requirements)

    def get_node_metadata(self, node_id: str) -> Dict[str, Any]:
        """Get metadata for a specific node."""
        return self._node_metadata.get(node_id, {})

    def add_connection(
        self, from_node: str, from_output: str, to_node: str, to_input: str
    ) -> "AsyncWorkflowBuilder":
        """Connect two nodes in the workflow (fluent interface version)."""
        super().add_connection(from_node, from_output, to_node, to_input)
        return self
