import os
from functools import lru_cache
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.tools import tool
from langchain_community.tools import DuckDuckGoSearchRun
from langgraph.prebuilt import create_react_agent
from dotenv import load_dotenv
from src.rag_engine import KnowledgeBase
from src.file_processor import FileProcessor

load_dotenv()

file_processor = FileProcessor()

_fallback_kb = KnowledgeBase(pdf_path=os.path.join("data", "policy.pdf"))
try:
    _fallback_kb.load_and_index()
except Exception as e:
    print(f"Fallback KB skipped: {e}")

_search_tool = DuckDuckGoSearchRun()

# gemini-2.5-flash reliably answers after tool calls; override per key via env.
MODEL_NAME = os.getenv("GEMINI_MODEL", "gemini-2.5-flash")

SYSTEM_PROMPT = (
    "You are a precise research assistant with two tools:\n"
    "- `lookup_documents`: searches the user's uploaded files and the internal knowledge base. "
    "Prefer this for questions about policies, documents, or any uploaded content.\n"
    "- `search_web`: searches the live web. Use this for current events, news, or general "
    "knowledge that is unlikely to be in the documents.\n\n"
    "Guidelines:\n"
    "1. Choose the tool that best fits the question; use both if needed.\n"
    "2. Ground your answer in the retrieved content and do not invent facts. If the documents "
    "do not contain the answer, say so and try the web.\n"
    "3. Use the conversation summary and recent turns to resolve follow-up questions "
    "(e.g. pronouns like 'it' or 'that').\n"
    "4. Be concise, accurate, and cite the source of your information when relevant."
)


@tool
def lookup_documents(query: str) -> str:
    """Search the user-uploaded documents for relevant information.
    Use this for questions about content in any uploaded files."""
    if file_processor.has_documents():
        result = file_processor.retrieve(query)
        if result:
            return result
    return _fallback_kb.retrieve(query)


@tool
def search_web(query: str) -> str:
    """Search the web for current events, news, or general knowledge not in uploaded documents."""
    try:
        return _search_tool.run(query)
    except Exception as e:
        return f"Search failed: {e}"


@lru_cache(maxsize=32)
def get_llm(api_key: str) -> ChatGoogleGenerativeAI:
    """Cached Gemini client per key, reused by the agent and the summarizer."""
    if not api_key or not api_key.strip():
        raise ValueError("A Google Gemini API key is required.")
    return ChatGoogleGenerativeAI(
        model=MODEL_NAME,
        temperature=0,
        google_api_key=api_key.strip(),
    )


@lru_cache(maxsize=32)
def get_agent_executor(api_key: str):
    """ReAct agent for a given key (BYOK). Tools and the index are shared; only the
    LLM is per-key, and the graph is cached so repeat calls don't rebuild it."""
    return create_react_agent(
        get_llm(api_key),
        [lookup_documents, search_web],
        prompt=SYSTEM_PROMPT,
    )