import json
import os
from typing import Dict, List, Any, Optional
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_openai import ChatOpenAI

class KnowledgeBaseConverterAgent:
    """
    Simple knowledge base converter without complex state management
    Converts admin input into structured knowledge base for roleplay training
    """
    
    def __init__(self, llm_model: str = "gpt-3.5-turbo"):
        self.llm = ChatOpenAI(model=llm_model, temperature=0.1, api_key="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")
    
    def extract_content(self, admin_input: str) -> Dict[str, Any]:
        """Step 1: Extract key components from admin input"""
        
        extraction_prompt = ChatPromptTemplate.from_template("""
        You are a content extraction specialist for safety training scenarios. 
        Analyze the admin input and extract key components systematically.
        
        Admin Input: {admin_input}
        
        Extract and categorize the following components:
        1. Scenario Information: Setting, situation, incident type, environment
        2. Characters: Names, roles, relationships, personality traits mentioned
        3. Safety Rules: Protocols, procedures, mandatory actions, restrictions
        4. Learning Objectives: Explicit goals or implied training outcomes
        5. Additional Context: Any other relevant details
        
        Provide a structured analysis in JSON format:
        {{
            "scenario_info": "extracted scenario details",
            "characters": ["character1: role and details", "character2: role and details"],
            "safety_rules": ["rule1", "rule2", "rule3"],
            "learning_objectives": ["objective1", "objective2"],
            "additional_context": "other relevant information"
        }}
        
        Be thorough and extract all relevant information.
        """)
        
        parser = JsonOutputParser()
        chain = extraction_prompt | self.llm | parser
        
        try:
            result = chain.invoke({"admin_input": admin_input})
            return {"success": True, "data": result}
        except Exception as e:
            return {"success": False, "error": f"Content extraction failed: {str(e)}"}
    
    def structure_mapping(self, extracted_content: Dict[str, Any]) -> Dict[str, Any]:
        """Step 2: Map extracted content to knowledge base structure"""
        
        structure_prompt = ChatPromptTemplate.from_template("""
        You are a knowledge base architect. Map extracted content to our structured format.
        Create realistic workplace safety training scenarios with proper protocols.
        
        Extracted Content: {extracted_content}
        
        Create a structured knowledge base with these sections:
        
        1. SCENARIO: Include name, setting (location, environment, equipment), initial situation (incident type, severity, dangers)
        2. CHARACTERS: For each character, define role, personality traits, speech patterns, knowledge areas, authority level, decision-making scope
        3. SAFETY PROTOCOLS: Immediate actions, supervisor actions, escalation procedures with step-by-step details
        4. LEARNING OBJECTIVES: Primary goals and key takeaways
        
        Make it realistic for workplace safety training. Include proper emergency response procedures.
        Focus on safety-first culture and proper incident management.
        
        Return a complete JSON structure following this format:
        {{
            "scenario": {{
                "name": "scenario_name",
                "setting": {{
                    "location": "setting_location",
                    "environment": "environment_type",
                    "safety_equipment_available": ["equipment1", "equipment2"]
                }},
                "initial_situation": {{
                    "incident_type": "type_of_incident",
                    "severity": "moderate",
                    "immediate_danger": "danger_description"
                }}
            }},
            "characters": {{
                "vikram": {{
                    "role": "experienced_employee",
                    "personality_traits": {{
                        "helpful": true,
                        "safety_focused": true,
                        "mentoring": true
                    }},
                    "speech_patterns": {{
                        "tone": "firm_but_caring",
                        "language_style": "uses_we_language"
                    }},
                    "knowledge_areas": ["safety_protocols", "chemical_handling"],
                    "authority_level": "peer_support",
                    "can_make_decisions": ["isolate_area", "guide_response"],
                    "cannot_decide": ["approve_cleanup", "authorize_reopening"]
                }},
                "mr_singh": {{
                    "role": "supervisor",
                    "personality_traits": {{
                        "calm": true,
                        "educational": true,
                        "no_blame_culture": true
                    }},
                    "speech_patterns": {{
                        "tone": "professional_and_reassuring",
                        "language_style": "clear_instructions"
                    }},
                    "knowledge_areas": ["all_safety_protocols", "incident_management"],
                    "authority_level": "supervisor",
                    "can_make_decisions": ["call_HSE", "approve_cleanup", "stop_operations"],
                    "cannot_decide": []
                }}
            }},
            "safety_protocols": {{
                "immediate_actions": [
                    {{
                        "step": 1,
                        "action": "ensure_personal_safety",
                        "details": "Check for chemical contact with skin/clothing"
                    }},
                    {{
                        "step": 2,
                        "action": "do_not_touch_spill",
                        "details": "Avoid contact, do not attempt cleanup"
                    }},
                    {{
                        "step": 3,
                        "action": "isolate_area",
                        "details": "Place warning signs, prevent entry"
                    }},
                    {{
                        "step": 4,
                        "action": "report_immediately",
                        "details": "Inform supervisor without delay"
                    }}
                ],
                "supervisor_actions": [
                    {{
                        "step": 1,
                        "action": "assess_situation",
                        "details": "Determine spill size, chemical type, risks"
                    }},
                    {{
                        "step": 2,
                        "action": "call_HSE_team",
                        "details": "Contact trained cleanup specialists"
                    }}
                ],
                "mandatory_procedures": {{
                    "cannot_skip": ["report_to_supervisor", "isolate_area", "use_proper_PPE"],
                    "consequences": "Safety risk and regulatory violations"
                }}
            }},
            "learning_objectives": {{
                "primary_goals": ["immediate_reporting", "proper_response_sequence", "safety_culture"],
                "key_takeaways": ["Report incidents immediately", "Follow established procedures", "Safety over blame"]
            }}
        }}
        """)
        
        parser = JsonOutputParser()
        chain = structure_prompt | self.llm | parser
        
        try:
            result = chain.invoke({"extracted_content": extracted_content})
            return {"success": True, "data": result}
        except Exception as e:
            return {"success": False, "error": f"Structure mapping failed: {str(e)}"}
    
    def enhance_knowledge_base(self, structured_data: Dict[str, Any]) -> Dict[str, Any]:
        """Step 3: Fill missing information and add conversation rules"""
        
        enhancement_prompt = ChatPromptTemplate.from_template("""
        You are a workplace safety expert. Enhance the knowledge base with missing critical information.
        
        Current Knowledge Base: {structured_data}
        
        Add these essential sections:
        
        1. CONVERSATION RULES - When characters appear and respond
        2. VISUAL GENERATION TRIGGERS - Key scenes for image generation
        3. ENFORCEMENT RULES - How strictly to follow protocols
        4. MISSING SAFETY DETAILS - Complete any gaps in protocols
        
        Return the complete enhanced knowledge base with all original content plus these additions:
        
        {{
            ... (keep all existing content) ...,
            "conversation_rules": {{
                "character_presence": {{
                    "vikram": {{
                        "appears_when": ["incident_occurs", "helping_colleague", "before_supervisor"],
                        "exits_when": ["supervisor_takes_over"]
                    }},
                    "mr_singh": {{
                        "appears_when": ["officially_reported", "serious_incident"],
                        "present_throughout": "once_involved"
                    }}
                }},
                "response_triggers": {{
                    "vikram_responds_if": ["raj_shows_panic", "safety_guidance_needed"],
                    "mr_singh_responds_if": ["officially_reported", "authorization_needed"]
                }},
                "enforcement_rules": {{
                    "must_enforce": ["safety_protocols", "incident_reporting"],
                    "enforcement_method": {{
                        "first_violation": "gentle_correction_with_explanation",
                        "repeated_violation": "firm_insistence_with_safety_reasons",
                        "dangerous_action": "immediate_intervention"
                    }}
                }}
            }},
            "visual_generation_triggers": {{
                "scenes_to_generate": [
                    {{
                        "trigger": "incident_start",
                        "description": "Chemical processing floor with spill area"
                    }},
                    {{
                        "trigger": "area_isolation",
                        "description": "Warning signs around spill area"
                    }},
                    {{
                        "trigger": "supervisor_office",
                        "description": "Industrial office with safety equipment"
                    }}
                ],
                "equipment_to_show": ["warning_signs", "spill_kit", "safety_gear"]
            }}
        }}
        """)
        
        parser = JsonOutputParser()
        chain = enhancement_prompt | self.llm | parser
        
        try:
            result = chain.invoke({"structured_data": structured_data})
            return {"success": True, "data": result}
        except Exception as e:
            return {"success": False, "error": f"Enhancement failed: {str(e)}"}
    
    def validate_knowledge_base(self, knowledge_base: Dict[str, Any]) -> Dict[str, Any]:
        """Step 4: Validate completeness and add metadata"""
        
        required_sections = ["scenario", "characters", "safety_protocols", "learning_objectives"]
        missing_sections = [section for section in required_sections if section not in knowledge_base]
        
        if missing_sections:
            return {
                "success": False, 
                "error": f"Missing required sections: {missing_sections}",
                "partial_data": knowledge_base
            }
        
        # Add metadata
        knowledge_base["metadata"] = {
            "created_from": "admin_input_conversion",
            "processing_status": "complete",
            "version": "1.0",
            "timestamp": "2024",
            "ready_for_roleplay": True
        }
        
        return {"success": True, "data": knowledge_base}
    
    def convert_knowledge_base(self, admin_input: str) -> Dict[str, Any]:
        """
        Main method to convert admin input to complete knowledge base
        """
        
        print("🚀 Starting Knowledge Base Conversion...")
        
        # Step 1: Extract content
        print("📝 Step 1: Extracting content...")
        extraction_result = self.extract_content(admin_input)
        if not extraction_result["success"]:
            return {
                "success": False,
                "error": extraction_result["error"],
                "stage": "extraction_failed"
            }
        
        # Step 2: Structure mapping
        print("🏗️ Step 2: Mapping to structure...")
        structure_result = self.structure_mapping(extraction_result["data"])
        if not structure_result["success"]:
            return {
                "success": False,
                "error": structure_result["error"],
                "stage": "structure_failed"
            }
        
        # Step 3: Enhancement
        print("✨ Step 3: Enhancing knowledge base...")
        enhancement_result = self.enhance_knowledge_base(structure_result["data"])
        if not enhancement_result["success"]:
            return {
                "success": False,
                "error": enhancement_result["error"],
                "stage": "enhancement_failed",
                "partial_data": structure_result["data"]
            }
        
        # Step 4: Validation
        print("✅ Step 4: Validating and finalizing...")
        validation_result = self.validate_knowledge_base(enhancement_result["data"])
        if not validation_result["success"]:
            return {
                "success": False,
                "error": validation_result["error"],
                "stage": "validation_failed",
                "partial_data": validation_result.get("partial_data")
            }
        
        print("🎉 Knowledge Base Conversion Complete!")
        return {
            "success": True,
            "knowledge_base": validation_result["data"],
            "stage": "complete",
            "metadata": {
                "conversion_successful": True,
                "ready_for_roleplay": True
            }
        }

# Example usage and testing
def test_converter():
    """Test the knowledge base converter with sample admin input"""
    
    # Check if OpenAI API key is set
 #   if not os.getenv("OPENAI_API_KEY"):
  #      print("❌ Error: Please set your OPENAI_API_KEY environment variable")
   #     print("Example: export OPENAI_API_KEY='your-api-key-here'")
    #    return
    
    print("🧪 Testing Knowledge Base Converter...")
    
    converter = KnowledgeBaseConverterAgent()
    
    sample_admin_input = """
    Create a chemical spill training scenario for a manufacturing plant. 
    
    Characters:
    - Raj: New employee who accidentally spills acetone-based solvent near mixing equipment
    - Vikram: Experienced worker, safety-focused, helps guide proper response
    - Mr. Singh: Shift supervisor, calm and educational, emphasizes safety over blame
    
    Safety Requirements:
    - Never touch spilled chemicals without PPE
    - Immediately isolate spill area with warning signs  
    - Report to supervisor without delay
    - Only trained HSE team performs cleanup
    - Document everything for learning
    
    Learning Goals:
    - Teach immediate incident reporting
    - Demonstrate proper emergency response sequence
    - Reinforce no-blame safety culture
    - Show importance of following protocols
    """
    
    result = converter.convert_knowledge_base(sample_admin_input)
    
    if result["success"]:
        print("\n✅ Knowledge Base Conversion Successful!")
        print(f"📊 Processing completed at stage: {result['stage']}")
        
        kb = result["knowledge_base"]
        print("\n📋 Generated Knowledge Base Summary:")
        print(f"- Scenario: {kb.get('scenario', {}).get('name', 'N/A')}")
        print(f"- Characters: {list(kb.get('characters', {}).keys())}")
        print(f"- Safety Protocol Steps: {len(kb.get('safety_protocols', {}).get('immediate_actions', []))}")
        print(f"- Learning Goals: {len(kb.get('learning_objectives', {}).get('primary_goals', []))}")
        
        # Save to file
        output_file = 'generated_knowledge_base.json'
        with open(output_file, 'w') as f:
            json.dump(kb, f, indent=2)
        print(f"\n💾 Complete knowledge base saved to '{output_file}'")
        
        # Display sample character info
        print("\n👥 Sample Character Info:")
        if 'vikram' in kb.get('characters', {}):
            vikram = kb['characters']['vikram']
            print(f"  Vikram - Role: {vikram.get('role', 'N/A')}")
            print(f"  Authority: {vikram.get('authority_level', 'N/A')}")
        
        print("\n🔒 Sample Safety Protocol:")
        immediate_actions = kb.get('safety_protocols', {}).get('immediate_actions', [])
        if immediate_actions:
            first_action = immediate_actions[0]
            print(f"  Step {first_action.get('step', 1)}: {first_action.get('action', 'N/A')}")
        
        print("\n🎯 Ready for roleplay agent integration!")
        
    else:
        print(f"\n❌ Conversion Failed at stage: {result['stage']}")
        print(f"❗ Error: {result['error']}")
        
        if result.get('partial_data'):
            print("📄 Partial data was generated - saving for debugging...")
            with open('partial_knowledge_base.json', 'w') as f:
                json.dump(result['partial_data'], f, indent=2)

if __name__ == "__main__":
    test_converter()