# backend/agent/chain_of_thought.py
import logging
try:
    from google import genai
    from google.genai import types
    from config import GEMINI_API_KEY, GEMINI_MODEL
except ImportError:
    pass

class AdversarialThoughtChain:
    def __init__(self, personality: str = "JARVIS"):
        self.personality = personality
        self._client = None
        try:
            self._client = genai.Client(api_key=GEMINI_API_KEY)
        except Exception as e:
            import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}")
        
    async def _ask_llm(self, prompt: str, system_instruction: str) -> str:
        if not self._client:
            return "Error: LLM client not initialized."
        try:
            resp = await self._client.aio.models.generate_content(
                model=GEMINI_MODEL,
                contents=prompt,
                config=types.GenerateContentConfig(
                    system_instruction=system_instruction,
                    temperature=0.7,
                ),
            )
            return resp.text.strip()
        except Exception as e:
            logging.error(f"CoT LLM Error: {e}")
            return f"Error: {e}"

    async def execute(self, user_input: str, context: str = "") -> dict:
        """Executes the Draft -> Critique -> Revise -> Adversarial loop."""
        
        try:
            from modules.assistant_identity import JARVIS_PERSONALITY_PROMPT, FRIDAY_PERSONALITY_PROMPT
            if self.personality.lower() == "jarvis":
                identity_prompt = JARVIS_PERSONALITY_PROMPT
            else:
                identity_prompt = FRIDAY_PERSONALITY_PROMPT
        except ImportError:
            identity_prompt = f"You are {self.personality}. Be highly intelligent and analytical."

        system_base = identity_prompt
        
        # 1. Draft
        draft_prompt = f"User: {user_input}\nContext: {context}\nProvide an initial comprehensive draft response."
        draft = await self._ask_llm(draft_prompt, system_base + " You are drafting an initial response. Do not self-censor or worry about perfection yet.")
        
        # 2. Critique
        critique_prompt = f"Draft Response:\n{draft}\n\nCritique this draft. What is unclear, logically flawed, or missing? Identify weaknesses."
        critique = await self._ask_llm(critique_prompt, system_base + " You are a harsh internal critic. Find flaws.")
        
        # 3. Revise
        revise_prompt = f"Original Draft:\n{draft}\n\nCritique:\n{critique}\n\nRevise the draft to address these criticisms."
        revision = await self._ask_llm(revise_prompt, system_base + " You are an editor. Fix the draft based on the critique.")
        
        # 4. Adversarial
        adversarial_prompt = f"You argued:\n{revision}\n\nNow argue the strongest possible counter-position. What assumptions are wrong? What did you miss?"
        adversarial = await self._ask_llm(adversarial_prompt, system_base + " You are an adversarial debater. Tear down the argument.")
        
        # 5. Final Resolution
        resolution_prompt = (
            f"User Query: {user_input}\n\n"
            f"Revised Response:\n{revision}\n\n"
            f"Adversarial Counter-Argument:\n{adversarial}\n\n"
            f"Synthesize the absolute best, most robust final answer. Acknowledge valid counter-points but deliver a definitive conclusion."
        )
        final_answer = await self._ask_llm(resolution_prompt, system_base + " Deliver the final, flawless response.")

        return {
            "draft": draft,
            "critique": critique,
            "revision": revision,
            "adversarial": adversarial,
            "final_answer": final_answer
        }