from google import genai
from google.genai import types
import os
import time
from .note_server import NoteServer

class ObsidianAgent:
    def __init__(self, api_key: str, server: NoteServer):
        self.client = genai.Client(api_key=api_key)
        self.server = server
        self.model_id = "gemini-3-flash-preview"

    def process_vault(self, input_folder: str, philosophy: str, rewrite_tag: str) -> list[str]:
        """
        Scans the input folder and processes each note using the LLM.
        Returns a list of action logs.
        """
        logs = []
        notes = self.server.list_notes(input_folder)
        
        # Derive Vault Root (Assuming input_folder is inside the vault, e.g. .../Vault/Inbox)
        # We take the parent directory of the input folder.
        vault_root = os.path.dirname(input_folder.rstrip(os.sep))
        
        if not notes:
            return ["No notes found in input folder."]

        logs.append(f"Found {len(notes)} notes. Processing (Slow mode for Rate Limits)...")

        for i, note_path in enumerate(notes):
            # RATE LIMITING: Sleep to respect 5 RPM (1 request every 12s + buffer)
            if i > 0:
                time.sleep(15)
                
            content = self.server.read_note(note_path)
            if not content:
                continue
                
            log = self._process_single_note_with_retry(note_path, content, philosophy, rewrite_tag, vault_root)
            logs.append(log)
            
        return logs

    def _process_single_note_with_retry(self, note_path: str, content: str, philosophy: str, rewrite_tag: str, vault_root: str) -> str:
        """Wraps processing with retry logic for 429 errors."""
        max_retries = 3
        for attempt in range(max_retries):
            try:
                return self._process_single_note(note_path, content, philosophy, rewrite_tag, vault_root)
            except Exception as e:
                if "429" in str(e) or "RESOURCE_EXHAUSTED" in str(e):
                    if attempt < max_retries - 1:
                        wait_time = 60 * (attempt + 1)
                        print(f"Rate limit hit. Waiting {wait_time}s...")
                        time.sleep(wait_time)
                        continue
                return f"Error processing {os.path.basename(note_path)}: {str(e)}"
        return f"Failed after retries: {os.path.basename(note_path)}"

    def _process_single_note(self, note_path: str, content: str, philosophy: str, rewrite_tag: str, vault_root: str) -> str:
        
        # 1. Define Tools
        def move_note(target_folder: str):
            """Moves the note to a folder relative to Vault Root. E.g. 'Science/Biology'."""
            # Ensure we don't treat target_folder as absolute unless it starts with vault_root
            if not target_folder.startswith(vault_root):
                 full_target = os.path.join(vault_root, target_folder.strip(os.sep))
            else:
                 full_target = target_folder
                 
            return self.server.move_note(note_path, full_target)

        def flag_rewrite(reason: str):
            """Flags the current note for rewrite by appending a tag and reason."""
            return self.server.flag_rewrite(note_path, reason, rewrite_tag)

        # 2. Construct Prompt
        prompt = f"""
        You are an expert Knowledge Manager for an Obsidian Vault.
        
        PHILOSOPHY:
        {philosophy}
        
        CURRENT NOTE CONTENT:
        {content[:10000]}  # Truncate to avoid massive context if note is huge
        
        TASK:
        Analyze the note content against the Philosophy.
        1. If it fits well into a specific folder in the vault structure, move it there.
        2. If it is low quality, incomplete, or violates the philosophy, flag it for rewrite.
        
        You MUST call a tool.
        """

        # 3. Call Gemini
        response = self.client.models.generate_content(
            model=self.model_id,
            contents=prompt,
            config=types.GenerateContentConfig(
                tools=[move_note, flag_rewrite],
                tool_config=types.ToolConfig(
                    function_calling_config=types.FunctionCallingConfig(
                        mode="ANY"  # Force the model to use a tool
                    )
                )
            )
        )

        # 4. Execute Tool
        if not response.function_calls:
                return f"Skipped: {os.path.basename(note_path)} (Model failed to call tool)"
        
        # Execute the first function call found
        fc = response.function_calls[0]
        if fc.name == "move_note":
            result = move_note(**fc.args)
            return f"[AI] {result}"
        elif fc.name == "flag_rewrite":
            result = flag_rewrite(**fc.args)
            return f"[AI] {result}"
        else:
            return f"Error: Unknown tool called {fc.name}"