"""
Google Grounding Search tool for Armorer agents.
Replaces DuckDuckGo with Google's official grounding API.
"""

import os
from typing import Optional
from armorer.agent_sdk import tool
from armorer.config import ConfigManager
from armorer.profiles import NetworkPolicy
from armorer.security.policy import get_policy_guard

# ── Lazy-initialised client ──────────────────────────────────────────────────

_client = None

def _get_api_key() -> str | None:
    """Resolve the Gemini API key from env vars or Armorer config."""
    key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
    if not key:
        try:
            key = ConfigManager().get_key("GEMINI_API_KEY")
        except Exception:
            pass
    return key or None

def is_available() -> bool:
    """Return True if a Gemini API key is configured."""
    return _get_api_key() is not None

def _get_client():
    """Create the genai client on first use."""
    global _client
    if _client is None:
        from google import genai
        api_key = _get_api_key()
        if not api_key:
            raise RuntimeError("No GEMINI_API_KEY or GOOGLE_API_KEY configured. Please run 'armorer config' to set it.")
        _client = genai.Client(api_key=api_key)
    return _client

@tool(category="knowledge", verification_level="external_grounded_search", safe_for_retrieval=True)
def google_search(query: str) -> str:
    """
    Search the web using Google Grounding Search for authoritative answers.
    Returns a grounded answer with inline source citations.
    
    Args:
        query: The search query or question to research.
    """
    try:
        guard = get_policy_guard()
        if guard.network_policy == NetworkPolicy.NONE:
            return "Google Search blocked by offline network policy."
        if guard.network_policy == NetworkPolicy.WHITELIST:
            return "Google Search blocked by scoped network policy (direct external search disabled)."

        from google.genai import types

        client = _get_client()
        grounding_tool = types.Tool(google_search=types.GoogleSearch())
        config = types.GenerateContentConfig(
            tools=[grounding_tool],
            temperature=0.0 # Grounding works best with low temp
        )

        # We ask the model to answer the query using search
        response = client.models.generate_content(
            model="gemini-2.5-flash-lite",
            contents=query,
            config=config,
        )

        text = response.text or ""

        # ── Build inline citations from grounding metadata ───────────
        try:
            metadata = response.candidates[0].grounding_metadata
            if metadata and metadata.grounding_supports and metadata.grounding_chunks:
                supports = sorted(
                    metadata.grounding_supports,
                    key=lambda s: s.segment.end_index,
                    reverse=True,
                )
                for support in supports:
                    end_idx = support.segment.end_index
                    if support.grounding_chunk_indices:
                        links = []
                        for i in support.grounding_chunk_indices:
                            if i < len(metadata.grounding_chunks):
                                chunk = metadata.grounding_chunks[i]
                                uri = chunk.web.uri
                                title = chunk.web.title or f"Source {i+1}"
                                links.append(f"[{title}]({uri})")
                        if links:
                            citation = " " + ", ".join(links)
                            text = text[:end_idx] + citation + text[end_idx:]
        except Exception:
            pass  # Citations are best-effort

        return text or "No results found."

    except Exception as e:
        return f"Google Search failed: {e}"
