from typing import Dict, List, Any, Optional
from datetime import datetime, UTC
import logging
from .methodBase import MethodBase
from modules.interfaces.serviceChatModel import MethodResult

logger = logging.getLogger(__name__)

class MethodOperator(MethodBase):
    """Operator methods for handling collections and AI operations"""
    
    def __init__(self, serviceContainer: Any):
        super().__init__(serviceContainer)
        self.name = "operator"
        self.description = "Operator methods for handling collections and AI operations"
        
    @property
    def actions(self) -> Dict[str, Dict[str, Any]]:
        """Available actions and their parameters"""
        return {
            "forEach": {
                "description": "Execute an action for each item in a list",
                "parameters": {
                    "items": {
                        "type": "List[Any]",
                        "description": "List of items to process",
                        "required": True
                    },
                    "action": {
                        "type": "Dict[str, Any]",
                        "description": "Action to execute for each item",
                        "required": True,
                        "properties": {
                            "method": {"type": "str", "required": True},
                            "action": {"type": "str", "required": True},
                            "parameters": {"type": "Dict[str, Any]", "required": False}
                        }
                    }
                }
            },
            "aiCall": {
                "description": "Call AI service with document content",
                "parameters": {
                    "prompt": {
                        "type": "str",
                        "description": "Prompt for AI processing",
                        "required": True
                    },
                    "extractedDocumentContent": {
                        "type": "List[Dict[str, str]]",
                        "description": "List of documents and their extraction prompts",
                        "required": True,
                        "items": {
                            "type": "object",
                            "properties": {
                                "document": {"type": "str", "required": True},
                                "promptForContentExtraction": {"type": "str", "required": True}
                            }
                        }
                    }
                }
            }
        }
        
    async def _executeAction(self, action: str, parameters: Dict[str, Any], authData: Optional[Dict[str, Any]] = None) -> MethodResult:
        """Execute operator action"""
        if action == "forEach":
            return await self._executeForEach(parameters)
        elif action == "aiCall":
            return await self._executeAiCall(parameters)
        else:
            raise ValueError(f"Unsupported action: {action}")
            
    async def _executeForEach(self, parameters: Dict[str, Any]) -> MethodResult:
        """Execute forEach operation"""
        try:
            items = parameters.get("items", [])
            action = parameters.get("action", {})
            
            if not items or not action:
                return self._createResult(
                    success=False,
                    data={},
                    error="Missing required parameters"
                )
                
            results = []
            for item in items:
                try:
                    # Execute action for each item
                    method = action.get("method")
                    action_name = action.get("action")
                    action_params = action.get("parameters", {})
                    
                    # Add current item to parameters
                    action_params["item"] = item
                    
                    # Execute method action
                    method_result = await self.service.methods[method][action_name](action_params)
                    results.append(method_result)
                    
                except Exception as e:
                    logger.error(f"Error processing item: {str(e)}")
                    results.append({
                        "success": False,
                        "error": str(e)
                    })
                    
            return self._createResult(
                success=True,
                data={"results": results}
            )
            
        except Exception as e:
            logger.error(f"Error in forEach execution: {str(e)}")
            return self._createResult(
                success=False,
                data={},
                error=str(e)
            )
            
    async def _executeAiCall(self, parameters: Dict[str, Any]) -> MethodResult:
        """Execute AI call with document content"""
        try:
            prompt = parameters.get("prompt")
            documents = parameters.get("extractedDocumentContent", [])
            
            if not prompt:
                return self._createResult(
                    success=False,
                    data={},
                    error="Missing prompt parameter"
                )
                
            # Extract content from documents
            extracted_content = []
            for doc in documents:
                try:
                    doc_ref = doc.get("document")
                    doc_prompt = doc.get("promptForContentExtraction")
                    
                    if not doc_ref or not doc_prompt:
                        continue
                        
                    # Extract content using document manager
                    content = self.service.extractContent(doc_prompt, doc_ref)
                    extracted_content.append({
                        "document": doc_ref,
                        "content": content
                    })
                    
                except Exception as e:
                    logger.error(f"Error extracting document content: {str(e)}")
                    continue
                    
            # Prepare AI prompt with extracted content
            full_prompt = f"{prompt}\n\nExtracted Content:\n"
            for content in extracted_content:
                full_prompt += f"\nDocument: {content['document']}\n{content['content']}\n"
                
            # Call AI service
            response = await self.service.callAiBasic(full_prompt)
            
            return self._createResult(
                success=True,
                data={
                    "response": response,
                    "processedDocuments": len(extracted_content)
                }
            )
            
        except Exception as e:
            logger.error(f"Error in AI call execution: {str(e)}")
            return self._createResult(
                success=False,
                data={},
                error=str(e)
            )