import os
import json
from google import genai
from openai import OpenAI

class Brain:
    def __init__(self):
        # --- CONNECTIONS ---
        self.gemini_key = os.getenv("GEMINI_API_KEY")
        self.openrouter_key = os.getenv("OPENROUTER_API_KEY")
        
        self.gemini_client = None
        if self.gemini_key:
            self.gemini_client = genai.Client(api_key=self.gemini_key)
            
        self.openrouter_client = None
        if self.openrouter_key:
            self.openrouter_client = OpenAI(
                api_key=self.openrouter_key,
                base_url="https://openrouter.ai/api/v1"
            )

        # --- TIG: THE TOOL REGISTRY ---
        self.tool_registry = {
            "coding": "anthropic/claude-sonnet-4",
            "creative_text": "openai/gpt-5.1-chat",
            "research": "gemini-3-flash-preview",
            "chat": "gemini-2.5-flash"
        }
        
        # State
        self.active_model = self.tool_registry["chat"]
        self.last_intent = "startup"

    def detect_intent(self, user_prompt):
        """
        Layer 1: Uses a fast, cheap model to classify the user's intent.
        Returns: 'coding', 'creative_text', 'research', or 'chat'
        """
        if not self.gemini_client:
            return "chat" # Fallback if Google is dead

        try:
            # Quick classification prompt
            classifier_prompt = f"""
            ANALYZE this user prompt and output ONLY ONE word from this list:
            [coding, creative_text, research, chat]
            
            - coding: asking for python, scripts, html, debugging, logic.
            - creative_text: writing stories, poems, complex essays.
            - research: asking for facts, summaries of files, history.
            - chat: casual conversation, greetings, simple questions.
            
            PROMPT: "{user_prompt[:500]}"
            """
            
            response = self.gemini_client.models.generate_content(
                model="gemini-2.5-flash",
                contents=classifier_prompt
            )
            intent = response.text.strip().lower()
            
            # Cleaning result just in case
            for valid in self.tool_registry.keys():
                if valid in intent:
                    return valid
            return "chat"
            
        except Exception as e:
            print(f"  [TIG Router] Intent detection failed: {e}")
            return "chat"

    def think(self, prompt, force_model=None):
        # 1. Did user force a model?
        target_model = force_model
        
        # 2. If not, Run TIG Intent Detection
        if not target_model:
            intent = self.detect_intent(prompt)
            
            # --- THE FIX: ALWAYS RESET TO REGISTRY DEFAULT ---
            # Previously, we kept 'self.active_model' stuck on the last choice.
            # Now, we look up the intent fresh every time.
            target_model = self.tool_registry.get(intent, "gemini-2.5-flash")
            
            self.last_intent = intent
            self.active_model = target_model # Update state
            
            # Only announce if it's NOT the default chat
            if intent != "chat":
                print(f"  [TIG] Intent: {intent.upper()} -> Routing to {target_model}")

        # 3. EXECUTION (The "Do It" Phase)
        
        # PATH A: Use Google Direct (If the chosen model is Gemini)
        if "gemini" in target_model and self.gemini_client:
            try:
                response = self.gemini_client.models.generate_content(
                    model="gemini-2.5-flash",
                    contents=prompt
                )
                return response.text
            except Exception as e:
                print(f"  [Brain] Google failed ({e}). Failover to OpenRouter.")
                # Fall through to OpenRouter

        # PATH B: Use OpenRouter (For Claude, GPT, or Google Failover)
        if self.openrouter_client:
            try:
                # print(f"  [Brain] Calling {target_model} via OpenRouter...")
                response = self.openrouter_client.chat.completions.create(
                    model=target_model,
                    messages=[{"role": "user", "content": prompt}],
                    temperature=0.7
                )
                return response.choices[0].message.content
            except Exception as e:
                return f"[System Critical] OpenRouter call failed: {e}"

        return "[System Critical] No API keys active. Check .env."