import asyncio
import os
from google.antigravity import types, Agent, LocalAgentConfig
from google.antigravity.hooks import hooks

class AgentWrapper:
    def __init__(self, project_path, thread_id, bot_client):
        self.project_path = project_path
        self.thread_id = thread_id
        self.bot = bot_client
        self.approval_future = None
        self.queue = asyncio.Queue()
        
        # Define hooks as closures inside __init__ so they have access to self
        # and are properly wrapped by the SDK decorators.
        
        @hooks.pre_tool_call_decide
        async def pre_tool_call_hook(tool_call: types.ToolCall) -> types.HookResult:
            """
            Intercepts any tool call (e.g., executing terminal commands or writing files)
            and sends a confirmation request to Discord.
            """
            # Create a Future to pause the async execution loop
            self.approval_future = asyncio.get_event_loop().create_future()
            
            # Format the message details
            args_str = str(tool_call.args)
            # Truncate if arguments are too long for Discord message limits
            if len(args_str) > 1000:
                args_str = args_str[:1000] + "..."
                
            message_content = (
                f"⚠️ **Antigravity wants to execute a tool:**\n"
                f"🔧 **Tool:** `{tool_call.name}`\n"
                f"📦 **Arguments:**\n```json\n{args_str}\n```\n"
                f"Please approve or deny this action."
            )
            
            # Send approval request with buttons
            await self.bot.send_approval_request(
                thread_id=self.thread_id,
                content=message_content
            )
            
            # Wait until the user clicks a button (resolves self.approval_future)
            approved = await self.approval_future
            self.approval_future = None
            
            return types.HookResult(allow=approved)

        @hooks.post_tool_call
        async def post_tool_call_hook(data):
            """
            Sends an update to Discord after a tool call completes.
            """
            await self.bot.send_to_thread(
                thread_id=self.thread_id,
                content=f"⚙️ **Tool execution completed.**"
            )

        # Configure the Google Antigravity Agent
        # Retrieve the API key, Base URL, and Model from environment variables
        api_key = os.getenv("GEMINI_API_KEY")
        base_url = os.getenv("GEMINI_BASE_URL")
        model_name = os.getenv("GEMINI_MODEL", "gemini-2.5-flash")
        
        if base_url:
            # Strip trailing '/v1' or '/v1beta' because the SDK appends the version automatically
            clean_base_url = base_url.rstrip("/")
            if clean_base_url.endswith("/v1"):
                clean_base_url = clean_base_url[:-3]
            elif clean_base_url.endswith("/v1beta"):
                clean_base_url = clean_base_url[:-7]
                
            # If a custom base URL (e.g. local proxy) is configured,
            # we explicitly construct the ModelTarget with a GeminiAPIEndpoint.
            endpoint = types.GeminiAPIEndpoint(
                api_key=api_key,
                base_url=clean_base_url
            )
            model_target = types.ModelTarget(
                name=model_name,
                types=[types.ModelType.TEXT],
                endpoint=endpoint
            )
            self.config = LocalAgentConfig(
                workspaces=[self.project_path],
                model=model_target,
                hooks=[
                    pre_tool_call_hook,
                    post_tool_call_hook
                ]
            )
        else:
            # Otherwise, let the SDK fall back to standard defaults (direct to Google)
            self.config = LocalAgentConfig(
                workspaces=[self.project_path],
                hooks=[
                    pre_tool_call_hook,
                    post_tool_call_hook
                ]
            )
        
        # Start the background agent session loop
        self.loop_task = asyncio.create_task(self.agent_loop())

    async def agent_loop(self):
        """
        Background loop that keeps the 'async with Agent' session alive
        and processes prompts from the queue sequentially.
        """
        try:
            async with Agent(self.config) as agent:
                while True:
                    prompt, future = await self.queue.get()
                    try:
                        # Call agent.chat within the active context
                        response = await agent.chat(prompt)
                        # Extract the text content from the ChatResponse object
                        text_response = await response.text()
                        future.set_result(text_response)
                    except Exception as e:
                        future.set_exception(e)
                    finally:
                        self.queue.task_done()
        except asyncio.CancelledError:
            print(f"Session loop for thread {self.thread_id} cancelled.")
        except Exception as e:
            print(f"Error in agent loop for thread {self.thread_id}: {e}")

    def resolve_approval(self, approved: bool):
        """
        Resolves the pending approval future with the user's decision.
        """
        if self.approval_future and not self.approval_future.done():
            self.approval_future.set_result(approved)

    async def run_prompt(self, prompt: str):
        """
        Runs a prompt against the Antigravity Agent by queuing it
        and waiting for the background session loop to process it.
        """
        # Send initial status
        status_msg = await self.bot.send_to_thread(
            thread_id=self.thread_id,
            content="🧠 *Antigravity is thinking...*"
        )
        
        # Create a future to receive the result from the background loop
        response_future = asyncio.get_event_loop().create_future()
        
        # Queue the prompt
        await self.queue.put((prompt, response_future))
        
        try:
            # Wait for the background loop to process it
            response = await response_future
            await status_msg.edit(content=response)
        except Exception as e:
            await self.bot.send_to_thread(
                thread_id=self.thread_id,
                content=f"💥 **An error occurred during execution:**\n`{str(e)}`"
            )