from .base import BaseRAGTechnique
from ..services.embed_service import get_embedding
from ..utils.json_utils import extract_json_block, repair_json
from typing import List, Dict, Any
import json

class AgenticRAG(BaseRAGTechnique):
    async def retrieve(self, query: str, document_id: str, top_k: int, **kwargs) -> List[Dict[str, Any]]:
        # This is the "agent loop"
        await self.emit("AGENT_INIT", "#22C55E", "Qwen3 agent ready with 4 tools")
        
        all_collected_chunks = []
        conversation_history = []
        
        system_prompt = f"""
You are an intelligent RAG agent. You have access to a document (ID: {document_id}).
Your goal is to answer the user query: "{query}"

Available tools:
1. search_docs(query: str, top_k: int) -> list of chunks
2. filter_search(filters: dict, query: str) -> list of chunks. filters can include "page" or "section".
3. get_page(page_num: int) -> text of that page
4. finish(answer: str) -> finish with final answer

Respond ONLY with a JSON object:
{{
    "thought": "your reasoning",
    "tool": "tool_name",
    "args": {{ ... }}
}}
"""
        
        for i in range(5): # Max 5 iterations
            await self.emit("PLAN", "#8B5CF6", f"Agent iteration {i+1}: Thinking...")
            
            agent_prompt = f"{system_prompt}\n\nHistory: {json.dumps(conversation_history)}\n\nAction:"
            response_text = self.llm.generate(agent_prompt)
            
            try:
                action_data = repair_json(extract_json_block(response_text))
                thought = action_data.get("thought", "")
                tool = action_data.get("tool", "")
                args = action_data.get("args", {})
                
                await self.emit("PLAN", "#8B5CF6", f"Thought: {thought[:100]}...")
                
                if tool == "finish":
                    self.final_agent_answer = args.get("answer", "")
                    break
                
                # Execute Tool
                await self.emit("TOOL", "#D97706", f"Tool call: {tool}({json.dumps(args)})")
                
                observation = ""
                if tool == "search_docs":
                    q = args.get("query", query)
                    tk = args.get("top_k", top_k)
                    q_vec = get_embedding(q)
                    results = await self.supabase.vector_search(q_vec, document_id, self.user_id, tk)
                    all_collected_chunks.extend(results)
                    observation = f"Found {len(results)} chunks."
                elif tool == "filter_search":
                    f = args.get("filters", {})
                    q = args.get("query", query)
                    matching_ids = await self.supabase.filter_chunk_ids(document_id, self.user_id, f)
                    if matching_ids:
                        q_vec = get_embedding(q)
                        results = await self.supabase.vector_search(q_vec, document_id, self.user_id, top_k, filter_chunk_ids=matching_ids)
                        all_collected_chunks.extend(results)
                        observation = f"Filtered search found {len(results)} chunks."
                    else:
                        observation = "No chunks matched the filters."
                elif tool == "get_page":
                    p = args.get("page_num")
                    results = await self.supabase.filter_chunk_ids(document_id, self.user_id, {"page": p})
                    if results:
                        chunks = await self.supabase.get_chunks_by_ids(results)
                        all_collected_chunks.extend(chunks)
                        observation = f"Retrieved page {p}."
                    else:
                        observation = f"Page {p} not found."
                
                await self.emit("OBSERVE", "#22C55E", observation)
                conversation_history.append({"action": action_data, "observation": observation})
                
            except Exception as e:
                logger.error(f"Agent error decoding JSON: {e}")
                conversation_history.append({"error": f"Invalid JSON response from your side. Use the required JSON format. error: {str(e)}"})

        await self.emit("FINAL", "#22C55E", "Answer generated after tool usage.")
        return all_collected_chunks

    async def generate(self, query: str, chunks: List[Dict[str, Any]]) -> str:
        # If the agent finished with a final answer, use it.
        if hasattr(self, "final_agent_answer") and self.final_agent_answer:
            return self.final_agent_answer
            
        await self.emit("GENERATE", "#7C3AED", "Qwen3 generating final summary...")
        prompt = self.build_prompt(query, chunks)
        return self.llm.generate(prompt)