"""
Base classes for DataFlow-aware Kaizen agents.

Provides mixin and base classes for integrating DataFlow database
capabilities into Kaizen agents while maintaining independence.

Architecture:
- DataFlowOperationsMixin: Database capabilities for any agent
- DataFlowAwareAgent: Pre-configured agent with DataFlow support
- Optional activation: Works with or without DataFlow instance
"""

from typing import Any, Dict, List, Optional

try:
    from kailash_dataflow import DataFlow

    DATAFLOW_AVAILABLE = True
except ImportError:
    DATAFLOW_AVAILABLE = False
    DataFlow = None

from kaizen.core.base_agent import BaseAgent

from .connection import DataFlowConnection


class DataFlowOperationsMixin:
    """
    Mixin providing DataFlow operations to Kaizen agents.

    Add database capabilities to any Kaizen agent:
    - Query operations
    - Insert/update/delete operations
    - Bulk operations
    - Transaction support

    Usage:
        class CustomAgent(BaseAgent, DataFlowOperationsMixin):
            def __init__(self, config, db=None):
                super().__init__(config)
                if db is not None:
                    self.connect_dataflow(db)

            def process_data(self):
                results = self.query_database(table='users', filter={'active': True})
                return results

    Attributes:
        db_connection: Optional DataFlowConnection instance
    """

    db_connection: Optional[DataFlowConnection] = None

    def connect_dataflow(self, db: "DataFlow"):
        """
        Connect this agent to a DataFlow instance.

        Args:
            db: DataFlow instance to connect to

        Raises:
            TypeError: If db is not a DataFlow instance

        Example:
            >>> from kailash_dataflow import DataFlow
            >>> db = DataFlow("postgresql://localhost/mydb")
            >>> agent.connect_dataflow(db)
        """
        # Allow mock objects for testing
        is_mock = hasattr(db, "_mock_name") or type(db).__name__ == "MagicMock"

        if not DATAFLOW_AVAILABLE and not is_mock:
            raise RuntimeError(
                "DataFlow is not installed. Install with: pip install kailash[dataflow]"
            )

        if not is_mock and not isinstance(db, DataFlow):
            raise TypeError(
                f"Expected DataFlow instance, got {type(db).__name__}. "
                f"Ensure you're passing a DataFlow object."
            )

        self.db_connection = DataFlowConnection(db=db, lazy_init=True)

    def query_database(
        self,
        table: str,
        filter: Optional[Dict[str, Any]] = None,
        limit: Optional[int] = None,
        order_by: Optional[List[str]] = None,
    ) -> List[Dict[str, Any]]:
        """
        Execute database query via DataFlow.

        Args:
            table: Table name to query
            filter: MongoDB-style filter conditions (optional)
            limit: Maximum number of results (optional)
            order_by: Sort order (optional)

        Returns:
            List of result dictionaries

        Raises:
            RuntimeError: If no DataFlow connection established

        Example:
            >>> results = agent.query_database(
            ...     table='users',
            ...     filter={'age': {'$gte': 18}},
            ...     limit=10,
            ...     order_by=['-created_at']
            ... )
        """
        if not self.db_connection:
            raise RuntimeError(
                "No DataFlow connection. Call connect_dataflow() first or "
                "pass db parameter during agent initialization."
            )

        # Build query parameters
        query_params = {}
        if filter is not None:
            query_params["filter"] = filter
        if limit is not None:
            query_params["limit"] = limit
        if order_by is not None:
            query_params["order_by"] = order_by

        # Execute query via DataFlow
        # This is a simplified implementation - actual implementation
        # would use DataFlow's ListNode through workflow execution
        if hasattr(self.db_connection.db, "query"):
            return self.db_connection.db.query(table, **query_params)

        # Fallback: Return empty list
        # In real implementation, this would build and execute a workflow
        return []

    def insert_record(self, table: str, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Insert a record into the database.

        Args:
            table: Table name
            data: Record data to insert

        Returns:
            Inserted record with generated ID

        Raises:
            RuntimeError: If no DataFlow connection established

        Example:
            >>> result = agent.insert_record(
            ...     table='users',
            ...     data={'name': 'Alice', 'email': 'alice@example.com'}
            ... )
        """
        if not self.db_connection:
            raise RuntimeError("No DataFlow connection. Call connect_dataflow() first.")

        # This would use DataFlow's CreateNode via workflow execution
        # Simplified placeholder implementation
        return data

    def update_record(
        self, table: str, record_id: Any, data: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Update a record in the database.

        Args:
            table: Table name
            record_id: ID of record to update
            data: Updated field values

        Returns:
            Updated record

        Raises:
            RuntimeError: If no DataFlow connection established

        Example:
            >>> result = agent.update_record(
            ...     table='users',
            ...     record_id=123,
            ...     data={'email': 'newemail@example.com'}
            ... )
        """
        if not self.db_connection:
            raise RuntimeError("No DataFlow connection. Call connect_dataflow() first.")

        # This would use DataFlow's UpdateNode via workflow execution
        return {**data, "id": record_id}

    def delete_record(self, table: str, record_id: Any) -> bool:
        """
        Delete a record from the database.

        Args:
            table: Table name
            record_id: ID of record to delete

        Returns:
            True if deleted successfully

        Raises:
            RuntimeError: If no DataFlow connection established

        Example:
            >>> success = agent.delete_record(table='users', record_id=123)
        """
        if not self.db_connection:
            raise RuntimeError("No DataFlow connection. Call connect_dataflow() first.")

        # This would use DataFlow's DeleteNode via workflow execution
        return True

    def bulk_insert(
        self, table: str, records: List[Dict[str, Any]], batch_size: int = 1000
    ) -> List[Dict[str, Any]]:
        """
        Bulk insert records into the database.

        Args:
            table: Table name
            records: List of records to insert
            batch_size: Number of records per batch

        Returns:
            List of inserted records with IDs

        Raises:
            RuntimeError: If no DataFlow connection established

        Example:
            >>> results = agent.bulk_insert(
            ...     table='users',
            ...     records=[
            ...         {'name': 'Alice', 'email': 'alice@example.com'},
            ...         {'name': 'Bob', 'email': 'bob@example.com'}
            ...     ]
            ... )
        """
        if not self.db_connection:
            raise RuntimeError("No DataFlow connection. Call connect_dataflow() first.")

        # This would use DataFlow's BulkCreateNode via workflow execution
        return records


class DataFlowAwareAgent(BaseAgent, DataFlowOperationsMixin):
    """
    Base class for Kaizen agents with DataFlow integration.

    Combines BaseAgent capabilities with database operations.
    Works with or without DataFlow instance.

    Features:
    - All BaseAgent capabilities (strategies, memory, etc.)
    - Database query operations
    - Database write operations
    - Bulk operations support
    - Optional DataFlow connection

    Usage:
        # Without DataFlow (standard agent)
        agent = DataFlowAwareAgent(config=config)

        # With DataFlow (database-enabled agent)
        from kailash_dataflow import DataFlow
        db = DataFlow("postgresql://localhost/mydb")
        agent = DataFlowAwareAgent(config=config, db=db)

        # Query database
        users = agent.query_database(table='users', filter={'active': True})

    Args:
        config: Agent configuration (can be domain-specific or BaseAgentConfig)
        db: Optional DataFlow instance for database operations
    """

    def __init__(self, config, db: Optional["DataFlow"] = None, **kwargs):
        """
        Initialize DataFlow-aware agent.

        Args:
            config: Agent configuration
            db: Optional DataFlow instance
            **kwargs: Additional arguments passed to BaseAgent
        """
        # Initialize BaseAgent
        super().__init__(config=config, **kwargs)

        # Connect to DataFlow if provided
        if db is not None:
            self.connect_dataflow(db)

    def get_database_info(self) -> Dict[str, Any]:
        """
        Get information about connected database.

        Returns:
            Dictionary with database connection info

        Example:
            >>> info = agent.get_database_info()
            >>> print(info['tables'])
        """
        if not self.db_connection:
            return {"connected": False, "message": "No DataFlow connection"}

        return {
            "connected": True,
            "tables": self.db_connection.list_tables(),
            "database": "connected",
        }

    def list_available_tables(self) -> List[str]:
        """
        List all available database tables.

        Returns:
            List of table names

        Raises:
            RuntimeError: If no DataFlow connection

        Example:
            >>> tables = agent.list_available_tables()
            >>> print(tables)  # ['users', 'products', 'orders']
        """
        if not self.db_connection:
            raise RuntimeError("No DataFlow connection. Call connect_dataflow() first.")

        return self.db_connection.list_tables()

    def get_table_schema(self, table_name: str) -> Dict[str, Any]:
        """
        Get schema for a specific table.

        Args:
            table_name: Name of the table

        Returns:
            Dictionary with schema information

        Raises:
            RuntimeError: If no DataFlow connection

        Example:
            >>> schema = agent.get_table_schema('users')
            >>> print(schema['columns'])
        """
        if not self.db_connection:
            raise RuntimeError("No DataFlow connection. Call connect_dataflow() first.")

        return self.db_connection.get_table_schema(table_name)
