"""
DSL Executor - Executes DSL queries using atomic functions

Takes DSL queries generated by LLM and executes them
against the page using atomic functions.
"""

import logging
from typing import Any, Dict, List, Optional
from dataclasses import dataclass

from curllm_core.llm_dsl.atoms import AtomicFunctions, AtomResult
from curllm_core.llm_dsl.generator import DSLQuery, DSLPlan, DSLQueryGenerator

logger = logging.getLogger(__name__)


@dataclass
class ExecutionResult:
    """Result of executing a DSL plan"""
    success: bool
    results: Dict[str, AtomResult]
    errors: List[str]
    final_data: Any = None


class DSLExecutor:
    """
    Executes DSL queries using atomic functions.
    
    Architecture:
    1. Takes DSL queries (from LLM generator)
    2. Executes atomic functions
    3. Handles errors and retries
    4. Returns aggregated results
    """
    
    def __init__(self, page=None, llm=None):
        self.page = page
        self.llm = llm
        self.atoms = AtomicFunctions(page=page, llm=llm)
        self.generator = DSLQueryGenerator(llm=llm)
    
    async def execute_natural_language(self, task: str) -> ExecutionResult:
        """
        Execute a natural language task.
        
        1. LLM generates DSL plan
        2. Executor runs atomic functions
        3. Results are aggregated
        """
        # Get page context
        page_context = None
        if self.page:
            try:
                page_context = await self.atoms.analyze_page_structure()
                if page_context.success:
                    page_context = page_context.data
            except Exception:
                pass
        
        # Generate DSL plan
        plan = await self.generator.generate_plan(task, page_context)
        
        if not plan or not plan.queries:
            return ExecutionResult(
                success=False,
                results={},
                errors=["Failed to generate DSL plan from task"],
            )
        
        # Execute plan
        return await self.execute_plan(plan)
    
    async def execute_plan(self, plan: DSLPlan) -> ExecutionResult:
        """
        Execute a DSL plan.
        """
        results = {}
        errors = []
        
        for i, query in enumerate(plan.queries):
            query_id = f"q{i}_{query.function}"
            
            try:
                result = await self.execute_query(query)
                results[query_id] = result
                
                if not result.success:
                    errors.append(f"{query.function}: No result")
                    
                    # Try to refine and retry
                    refined = await self.generator.refine_query(
                        query, 
                        f"Function returned no result"
                    )
                    if refined:
                        retry_result = await self.execute_query(refined)
                        if retry_result.success:
                            results[query_id] = retry_result
                            errors.pop()  # Remove error if retry succeeded
                            
            except Exception as e:
                error_msg = f"{query.function}: {str(e)}"
                errors.append(error_msg)
                results[query_id] = AtomResult(success=False)
        
        # Determine overall success
        success = any(r.success for r in results.values())
        
        # Extract final data from successful results
        final_data = None
        for result in results.values():
            if result.success and result.data:
                final_data = result.data
                break
        
        return ExecutionResult(
            success=success,
            results=results,
            errors=errors,
            final_data=final_data,
        )
    
    async def execute_query(self, query: DSLQuery) -> AtomResult:
        """
        Execute a single DSL query.
        """
        func_name = query.function
        params = query.params or {}
        
        # Map function name to atomic function
        func_map = {
            'find_input_by_context': self.atoms.find_input_by_context,
            'find_clickable_by_intent': self.atoms.find_clickable_by_intent,
            'find_url_by_intent': self.atoms.find_url_by_intent,
            'analyze_page_structure': self.atoms.analyze_page_structure,
            'find_repeating_pattern': self.atoms.find_repeating_pattern,
            'detect_message_type': self.atoms.detect_message_type,
            'extract_data_pattern': self.atoms.extract_data_pattern,
        }
        
        func = func_map.get(func_name)
        if not func:
            logger.error(f"Unknown DSL function: {func_name}")
            return AtomResult(success=False)
        
        # Execute function with params
        try:
            if params:
                result = await func(**params)
            else:
                result = await func()
            return result
        except TypeError as e:
            # Parameter mismatch
            logger.error(f"Parameter error for {func_name}: {e}")
            return AtomResult(success=False)
        except Exception as e:
            logger.error(f"Execution error for {func_name}: {e}")
            return AtomResult(success=False)
    
    async def find_element(self, description: str) -> Optional[str]:
        """
        High-level: Find element by description, return selector.
        
        Uses LLM to generate query, then executes it.
        """
        query = await self.generator.generate_single_query(
            f"find element for: {description}"
        )
        
        if not query:
            return None
        
        result = await self.execute_query(query)
        
        if result.success and result.data:
            return result.data.get('selector')
        
        return None
    
    async def detect_page_state(self) -> Dict[str, Any]:
        """
        High-level: Detect current page state.
        """
        # Analyze structure
        structure = await self.atoms.analyze_page_structure()
        
        # Detect messages
        messages = await self.atoms.detect_message_type()
        
        return {
            'structure': structure.data if structure.success else None,
            'message': messages.data if messages.success else None,
        }
    
    async def execute_form_fill(self, instruction: str) -> ExecutionResult:
        """
        High-level: Fill a form based on instruction.
        
        LLM interprets instruction and generates DSL queries
        for each field to fill.
        """
        return await self.execute_natural_language(
            f"Fill form with: {instruction}"
        )
