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

class IntelliMod:
    def __init__(self):
        load_dotenv()
        self.gemini_key = os.getenv("GEMINI_API_KEY")
        self.openrouter_key = os.getenv("OPENROUTER_API_KEY")
        
        # --- CLIENTS ---
        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"
            )

        # --- THE TIG REGISTRY (Layer 2) ---
        # Models use OpenRouter IDs for non-Gemini, raw names for Gemini direct
        self.tool_registry = {
            # TIER 1: HEAVY LIFTERS (High Cost / Complex Logic)
            "coding": "anthropic/claude-sonnet-4",
            "planning": "anthropic/claude-opus-4",
            "creative_text": "openai/gpt-5.1-chat",
            
            # TIER 2: SPECIALISTS (Medium Cost / Specific Tasks)
            "visual_image": "gemini-3-pro-preview",
            "research": "gemini-3-flash-preview",
            
            # TIER 3: SPEED & CHAT (Unlimited / Low Cost)
            "chat": "gemini-2.5-flash",
            "sorting": "gemini-2.5-flash"
        }
        
        self.active_model = self.tool_registry["chat"]

    def detect_intent(self, user_prompt):
        """
        Layer 1: Intent Classification
        """
        # 1. SHORT-CIRCUIT
        if len(user_prompt.split()) < 5:
             return "chat"

        if not self.gemini_client:
            return "chat"

        try:
            classifier_prompt = f"""
            ANALYZE this user prompt and output ONLY ONE word from this list:
            [coding, creative_text, research, visual_image, planning, chat]
            
            - coding: python, scripts, html, debugging, logic, "write code".
            - creative_text: stories, essays, poems, long-form writing.
            - research: facts, history, summarizing files, looking up info.
            - visual_image: descriptions of images, asking for image generation prompts.
            - planning: complex step-by-step plans, architecture, project management.
            - chat: casual conversation, greetings, simple questions, thank yous.
            
            PROMPT: "{user_prompt[:1000]}"
            """
            
            response = self.gemini_client.models.generate_content(
                model="gemini-2.5-flash",
                contents=classifier_prompt
            )
            intent = response.text.strip().lower()
            
            for valid in self.tool_registry.keys():
                if valid in intent:
                    return valid
            return "chat"
            
        except Exception as e:
            print(f"  [IntelliMod] Intent detection failed: {e}")
            return "chat"

    def run_tig_pipeline(self, prompt, force_model=None):
        # 1. Routing
        if force_model:
            target_model = force_model
            intent = "manual_override"
        else:
            intent = self.detect_intent(prompt)
            target_model = self.tool_registry.get(intent, "gemini-2.5-flash")
        
        # Update State
        self.active_model = target_model
        
        # 2. Execution
        return self._execute_model(target_model, prompt)

    def _execute_model(self, model, prompt):
        # FAST PATH: No system prompts, just raw speed.
        
        # PATH A: Google Direct
        if "gemini" in model and self.gemini_client:
            try:
                response = self.gemini_client.models.generate_content(
                    model=model,
                    contents=prompt
                )
                return response.text
            except Exception as e:
                print(f"  [IntelliMod] Google ({model}) failed.")

        # PATH B: OpenRouter (Claude, GPT, and other non-Gemini models)
        if self.openrouter_client:
            try:
                response = self.openrouter_client.chat.completions.create(
                    model=model,
                    messages=[{"role": "user", "content": prompt}],
                    temperature=0.7
                )
                return response.choices[0].message.content
            except Exception as e:
                return f"[System Critical] OpenRouter failed: {e}"

        return "[System Critical] No brains active."

    # --- MPA (Modular Prompt Auditor) ---
    MPA_PATH = os.path.join(os.path.dirname(__file__), "..", "intellimod_system", "content", "intellimod_knowledge", "audit-protocols", "modular-prompt-auditor.md")

    def run_mpa_pipeline(self, user_prompt, force_model=None):
        """Loads the MPA spec and runs a full 9-step prompt evaluation."""
        try:
            with open(self.MPA_PATH, "r") as f:
                mpa_spec = f.read()
        except Exception as e:
            return f"[MPA Error] Could not load auditor spec: {e}"

        # Trim the footer (everything from "Insert Prompt to Evaluate" onward)
        cutoff = mpa_spec.find("Insert Prompt to Evaluate")
        if cutoff != -1:
            mpa_spec = mpa_spec[:cutoff]

        # Build the evaluator prompt
        evaluator_prompt = f"""{mpa_spec.strip()}

---
## EVALUATION TARGET

Prompt to Evaluate:
```
{user_prompt}
```

---
## INSTRUCTIONS

Execute ALL 9 steps above against this prompt.
Be thorough and specific. Output each step clearly.
"""

        target_model = force_model or "anthropic/claude-sonnet-4"
        return self._execute_model(target_model, evaluator_prompt)