from __future__ import annotations

import asyncio
import json
import logging
from pathlib import Path
from typing import Any

from google import genai

from epochler.config import get_settings

logger = logging.getLogger(__name__)

_SYSTEM_PROMPT = (Path(__file__).parent / "prompts" / "searcher.md").read_text(encoding="utf-8")


def _sync_generate(query: str) -> str:
    """Run the synchronous Gemini API call (intended for use with asyncio.to_thread)."""
    settings = get_settings()
    client = genai.Client(api_key=settings.gemini_api_key)
    tool = genai.types.Tool(google_search=genai.types.GoogleSearch())
    response = client.models.generate_content(
        model=settings.gemini_model,
        contents=f"{_SYSTEM_PROMPT}\n\n---\n\nSearch query: {query}\n\nReturn your findings as JSON.",
        config=genai.types.GenerateContentConfig(
            tools=[tool],
            temperature=0.2,
        ),
    )
    return response.text.strip()


async def search(query: str) -> dict[str, Any]:
    """Use Gemini with grounding (Google Search) to answer an ML research query."""
    settings = get_settings()
    if not settings.gemini_api_key:
        return {
            "query": query,
            "summary": "Search unavailable — no Gemini API key configured.",
            "findings": [],
            "recommendation": "",
        }

    try:
        raw_text = await asyncio.wait_for(
            asyncio.to_thread(_sync_generate, query),
            timeout=60,
        )
    except Exception as e:
        logger.exception("Gemini API call failed for query: %s", query)
        return {
            "query": query,
            "summary": f"Search failed: {e}",
            "findings": [],
            "recommendation": "",
        }

    try:
        if "```json" in raw_text:
            raw_text = raw_text.split("```json", 1)[1].split("```", 1)[0].strip()
        elif "```" in raw_text:
            raw_text = raw_text.split("```", 1)[1].split("```", 1)[0].strip()
        result = json.loads(raw_text)
    except (json.JSONDecodeError, IndexError):
        result = {
            "query": query,
            "summary": raw_text[:500],
            "findings": [],
            "recommendation": "",
        }

    result.setdefault("query", query)
    return result
