import os
import logging
from pathlib import Path
from typing import Text, List, Dict, Optional, Union, Tuple, Any
from cota.dst import DST
from cota.actions.action import Action
from cota.dpl.dpl import DPL
from cota.utils.io import read_yaml_from_path
from cota.utils.common import hash_str
from cota.constant import DEFAULT_DIALOGUE_MAX_TOKENS

logger = logging.getLogger(__name__)


class LLMDPL(DPL):
    """LLM-based Dialogue Policy Learning class that uses LLM with knowledge retrieval to generate thoughts."""
    
    def __init__(
        self,
        path: Union[Text, Path],
        llms: Union[Text, Dict, List, None] = None
    ) -> None:
        """Initialize LLMDPL with policy data."""
        # Support llms list configuration
        self.llms = llms or []
        
        # Build action to LLM mapping and find default LLM (similar to LLMKnowledge)
        self.action_llm_map = {}
        self.default_llm = None
        
        if isinstance(self.llms, list):
            self._build_llm_mappings()
        
        self.build(path)
    
    def _build_llm_mappings(self):
        """Build action to LLM mapping and find default LLM from llms list configuration."""
        for llm_config in self.llms:
            if not isinstance(llm_config, dict):
                logger.warning(f"Invalid LLM config format: {llm_config}")
                continue
                
            llm_name = llm_config.get('name')
            action = llm_config.get('action')
            
            if not llm_name:
                logger.warning(f"LLM configuration missing name: {llm_config}")
                continue
                
            if action:
                # LLM with specific action
                self.action_llm_map[action] = llm_name
            else:
                # LLM without specific action serves as default
                self.default_llm = llm_name
    
    def get_llm_for_action(self, action_name: str) -> Optional[str]:
        """Get LLM configuration for specific action type."""
        if not self.llms:
            return None
            
        # Handle list format (primary configuration method)
        if isinstance(self.llms, list):
            # Check if we have a specific LLM for this action
            if action_name in self.action_llm_map:
                return self.action_llm_map[action_name]
            
            # Use default LLM if available
            if self.default_llm:
                return self.default_llm
                
            # Fallback to first available LLM
            if self.llms:
                first_llm = self.llms[0]
                if isinstance(first_llm, dict):
                    return first_llm.get('name')
                    
            return None
            
        # Handle string format (single LLM for all actions) - backward compatibility
        elif isinstance(self.llms, str):
            return self.llms
            
        # Handle dictionary format (legacy direct action mapping) - backward compatibility
        elif isinstance(self.llms, dict):
            # First try exact match
            if action_name in self.llms:
                return self.llms[action_name]
            # Fallback to 'default' key if available
            return self.llms.get('default')
            
        else:
            logger.warning(f"Unsupported LLM config type: {type(self.llms)}")
            return None
    
    
    def build(self, path: Union[Text, Path]):
        """Prepare data for LLM with streaming processing."""
        # Load policy data from path
        policies = self.load_data(path)

        if not policies:
            logger.warning(f"No policies found in {path}")
            return {}

        # Stream process: read policies and write to file directly
        self.process_policies(policies, path)

        # Log reminder for LLM knowledge platform submission
        logger.info("=" * 60)
        logger.info("📋 LLM DPL Data Preparation Complete (Auto-Generated by Action Type)")
        logger.info("=" * 60)
        logger.info(f"📁 Source path: {path}")
        logger.info("")
        logger.info("⚠️  IMPORTANT: Please submit the following data to your LLM knowledge platform:")
        logger.info(f"  1. Upload all generated knowledge files to your knowledge base")
        logger.info("   2. Configure LLM knowledge retrieval for each action type")
        logger.info("   3. Test LLM knowledge query functionality for all action types")
        logger.info("📤 Ready-to-upload files automatically separated by action type")
        logger.info("=" * 60)

    def process_policies(self, policies: List[Dict], policies_path: Union[Text, Path]):
        """Process policies: read and write to file directly for all action types."""
        import time
        
        policies_path = Path(policies_path)
        timestamp = int(time.time())
        
        # First pass: collect all action types that have thoughts
        action_types = set()
        for policy in policies:
            actions = policy.get("actions", [])
            for action in actions:
                if 'thought' in action:
                    action_name = action.get('name', '')
                    if action_name:
                        action_types.add(action_name)
        
        if not action_types:
            logger.warning("No actions with thoughts found in policies")
            return
        
        # Create output files for each action type
        action_files = {}
        action_hash_keys = {}
        file_handles = {}
        
        try:
            # Initialize files and data structures for each action type
            for action_type in action_types:
                safe_action_name = action_type.lower().replace(' ', '_')
                output_path = policies_path / f"llm_dpl_{safe_action_name}_{timestamp}.md"
                output_path.parent.mkdir(parents=True, exist_ok=True)
                
                action_hash_keys[action_type] = set()
                file_handles[action_type] = open(output_path, 'w', encoding='utf-8')
                
                logger.info(f"Created knowledge file for {action_type}: {output_path.name}")
            
            # Process all policies
            for policy_idx, policy in enumerate(policies):
                actions = policy.get("actions", [])
                if not actions:
                    continue
                
                policy_title = policy.get("title", f"policy_{policy_idx}")
                user_utter_index = self.build_user_utter_index(actions)
                
                for i, action in enumerate(actions):
                    if 'thought' in action:
                        action_name = action.get('name', '')
                        
                        # Skip if action type not found or no file handle
                        if action_name not in file_handles:
                            continue
                        
                        target_file = file_handles[action_name]
                        target_hash_keys = action_hash_keys[action_name]
                        
                        segments = self.trace_back_to_user_utter(actions, i, user_utter_index)
                        
                        for segment in segments:
                            hash_key = '_'.join(f"{action.get('name')}:{action.get('thought', '')}" for action in segment)
                            if hash_str(hash_key) not in target_hash_keys:
                                target_hash_keys.add(hash_str(hash_key))
                                
                                content = self.segment_to_llm_content(segment)
                                if content:
                                    target_file.write(f"# {policy_title}\n")
                                    target_file.write(f"{content}\n")
                                    target_file.write(f"\n\n")
        
        finally:
            # Close all file handles
            for file_handle in file_handles.values():
                file_handle.close()
            
        logger.info(f"Generated knowledge files for {len(action_types)} action types: {list(action_types)}")


    def segment_to_llm_content(self, segment: List[Dict]) -> str:
        """Convert action segment to LLM knowledge document content."""
        content_lines = []
        
        for action in segment:
            action_name = action.get('name')
            if not action_name:
                continue
            
            # Add thought if available
            if action.get('thought'):
                content_lines.append(f"thought:{action.get('thought')}")

            # Add action result if available
            if action.get('result'):
                content_lines.append(f"{action_name}:{action.get('result')}")
        
        return "\n".join(content_lines)
    
    async def generate_thoughts(self, dst: DST, action: Action) -> Text:
        """Generate thoughts using LLM with action-specific configuration."""
        # Determine action name
        action_name = action.name if hasattr(action, 'name') else action.__class__.__name__
        
        # Get LLM for this action type
        llm_name = self.get_llm_for_action(action_name)
        if not llm_name:
            logger.warning(f"No LLM configured for action type: {action_name}")
            return ""
        
        # Build query from current context
        query_text = self.build_llm_query(dst, action)
        if not query_text.strip():
            logger.debug("Empty query text for LLM knowledge retrieval")
            return ""
        
        logger.debug(f"Using LLM '{llm_name}' for action '{action_name}'")
        
        try:
            result = await dst.agent.llm_instance(llm_name).generate_chat(
                messages=[
                    {
                        "role":"user", 
                        "content": query_text
                    }
                ],
                max_tokens=dst.agent.dialogue.get('max_tokens', DEFAULT_DIALOGUE_MAX_TOKENS)
            )
            return result["content"]
        except Exception as e:
            logger.error(f"Failed to generate thoughts for {action_name} using {llm_name}: {e}")
            return ""

    def build_llm_query(self, dst: DST, action: Action) -> Text:
        """Build query text for LLM knowledge retrieval using the longest context segment."""
        actions = dst.formless_actions
        segment_parts = []
        for act in actions:
            if act and act.result:
                # result is a list, need to search for text in the result items
                text_parts = []
                for result_item in act.result:
                    if isinstance(result_item, dict) and result_item.get('text'):
                        text_parts.append(result_item.get('text', ''))
                
                # Combine all text parts for this action into one entry
                if text_parts:
                    combined_text = ' '.join(text_parts)
                    segment_parts.append(f"{act.name}:{combined_text}")
        
        # Return the complete segment as query
        return "\n".join(segment_parts)
        
    def load_data(self, path: Union[Text, Path]) -> List[Dict]:
        """Load policy data from YAML files, similar to MatchDPL."""
        policies = []
        path_obj = Path(path)

        if not path_obj.is_dir():
            logger.warning(f"Path {path} is not a directory")
            return policies

        # Process all YAML files in directory
        for yml_file in path_obj.glob('*.yml'):
            try:
                data = read_yaml_from_path(yml_file)
                
                # Only extract data from 'policies' sections
                if isinstance(data, dict) and 'policies' in data:
                    policies.extend(data['policies'])
            except Exception as e:
                logger.error(f"Failed to load {yml_file}: {e}")
                
        return policies

    def build_user_utter_index(self, actions):
        return [i for i, action in enumerate(actions) if action.get('name') == 'UserUtter']

    def trace_back_to_user_utter(self, actions, index, user_utter_index):
        segments = []
        for q_index in reversed(user_utter_index):
            if q_index < index:
                segment = [action for action in actions[q_index:index] if action.get('name') and action.get('name') not in ('Selector','Updater')]
                segment.append(actions[index])
                if len(segment) > 1:
                    segments.append(segment)
        return segments
