import os
import asyncio
from datetime import datetime
from langchain_groq import ChatGroq
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_community.utilities import GoogleSerperAPIWrapper
from src.config import SystemConfig
from src.memory import MemoryJournal

class KitchenBrain:
    def __init__(self):
        self.cfg = SystemConfig()
        
        # PRIMARY BRAIN (Fast, Llama 3)
        self.primary_llm = ChatGroq(
            model="llama-3.3-70b-versatile", 
            api_key=self.cfg.groq_key,
            streaming=True
        )
        
        # BACKUP BRAIN (Reliable, OpenAI)
        if self.cfg.openai_api_key:
            self.backup_llm = ChatOpenAI(
                model="gpt-4o-mini",
                api_key=self.cfg.openai_api_key,
                streaming=True
            )
        else:
            self.backup_llm = None

        self.memory = MemoryJournal()
        
        if self.cfg.serper_key:
            self.search = GoogleSerperAPIWrapper(serper_api_key=self.cfg.serper_key)
        else:
            self.search = None

    async def route_and_process(self, user_input):
        self.memory.save_interaction("user", user_input, "👤")
        text = user_input.lower()
        
        # --- STRICT ROUTING LOGIC ---
        
        # Brie is now INVITE ONLY. She only appears if you explicitly ask for output.
        # We removed vague words like 'prepare', 'make', 'cook', 'food'.
        brie_triggers = [
            'recipe', 
            'ingredients for', 
            'instructions for', 
            'how do i cook', 
            'how do i make', 
            'how to cook', 
            'how to make',
            'shopping list'
        ]
        
        # Check if ANY of the strict triggers are in the text
        is_requesting_chef = any(trigger in text for trigger in brie_triggers)

        persona = "Olivia"
        handoff_msg = ""

        if is_requesting_chef:
            persona = "Brie"
            handoff_msg = self.get_handoff_message(user_input)
            generator = self.stream_brie(user_input)
        else:
            # Default to Olivia for EVERYTHING else.
            # Even "I am preparing dinner" stays with Olivia now.
            generator = self.stream_olivia(user_input)

        return persona, handoff_msg, generator

    def get_handoff_message(self, text):
        return "That sounds delicious. I'll ask Brie to handle the culinary details."

    async def _safe_stream(self, messages):
        try:
            async for chunk in self.primary_llm.astream(messages):
                yield chunk.content
            return
        except Exception as e:
            print(f"⚠️ Primary Brain Failed: {e}")
            
        if self.backup_llm:
            try:
                print("🔄 Switching to Backup Brain (OpenAI)...")
                async for chunk in self.backup_llm.astream(messages):
                    yield chunk.content
                return
            except Exception as e:
                print(f"⚠️ Backup Brain Failed: {e}")
        
        yield "I'm having trouble connecting to my networks right now. Please try again in a moment."

    async def stream_olivia(self, text):
        now = datetime.now().strftime("%A, %B %d, %Y at %I:%M %p")
        past_context = self.memory.get_context_string(limit=10)
        
        search_data = ""
        triggers = ['weather', 'news', 'score', 'price', 'who is', 'what is', 'when is', 'location', 'find', 'near me']
        skip = ['sad', 'happy', 'tired', 'love', 'hate', 'joke']
        
        if self.search and any(t in text.lower() for t in triggers) and not any(s in text.lower() for s in skip):
            try:
                query = f"{text} in {self.cfg.location} ({now})"
                res = self.search.run(query)
                search_data = f"\n[REAL-TIME INFO]: {res}"
            except: pass

        # Tuned System Prompt to handle food chat better
        sys_prompt = f"""You are Olivia, a sophisticated Household Companion.
        Time: {now}. Location: {self.cfg.location}.
        User Name: {self.cfg.user_name}.
        
        MEMORY: {past_context}
        CONTEXT: {search_data}
        
        GUIDANCE:
        - You are the Manager. You handle chat, scheduling, and life updates.
        - If the user talks about food (e.g., "I'm making dinner"), be supportive and conversational.
        - DO NOT generate full recipes yourself.
        - If the user explicitly asks for a recipe, you can suggest asking Brie.
        - Be warm, professional, and concise."""
        
        msgs = [SystemMessage(content=sys_prompt), HumanMessage(content=text)]
        
        async for chunk in self._safe_stream(msgs):
            yield chunk

    async def stream_brie(self, text):
        prompt = """You are Brie, an elite private chef and cooking companion. You are warm, encouraging, and love helping people cook!

STRICT OUTPUT FORMAT - Follow this exactly:

**[Recipe Name]**

**Ingredients:**
- [ingredient 1]
- [ingredient 2]
- [ingredient 3]
(list all ingredients as bullet points)

**Instructions:**
1. [First step]
2. [Second step]
3. [Third step]
(number all steps clearly)

**Chef's Note:**
[One helpful tip or variation suggestion]

IMPORTANT RULES:
- Always use bullet points (-) for ingredients
- Always use numbers (1. 2. 3.) for instructions
- Keep instructions clear and concise
- Be encouraging and friendly in your Chef's Note
- Do NOT add extra sections or commentary outside this format"""
        
        msgs = [SystemMessage(content=prompt), HumanMessage(content=text)]
        
        async for chunk in self._safe_stream(msgs):
            yield chunk