import logging
from typing import Any, Dict

import httpx
from inspect_ai.log import transcript
from inspect_ai.model import ModelUsage
from inspect_ai.tool import Tool, ToolDef, tool

from astabench.util.model import record_model_usage_with_inspect

DEFAULT_ASYNC_TIMEOUT = 600

logger = logging.getLogger(__name__)

QUERY_ATTEMPTS = 3
pf_model_name_2_litellm_model_name = {
    "models/gemini-2.0-flash": "gemini/gemini-2.0-flash",
    "gemini-2.0-flash": "gemini/gemini-2.0-flash",
    "gpt-4o": "gpt-4o-2024-08-06",
}


async def find_papers(
    query: str,
    base_url: str,
    timeout: int = DEFAULT_ASYNC_TIMEOUT,
    inserted_before: str | None = None,
    read_results_from_cache: bool = True,
) -> Dict[str, Any]:
    """
    Search via the AI2I paper-finder API and return the response.

    Args:
        query: The search query (paper description)
        base_url: Base API URL
        timeout: Maximum time to wait for the search to complete (in seconds)
        inserted_before: cut off papers from this date onwards

    Returns:
        Raw response from paper-finder (dictionary containing the search results).  The format varies based on whether the `enriched` flag is set.

        See the API documentation for more info.  The output schema is as
        follows:

        ```json
        {
            "response_text": "I found 253 papers matching your query.\n\n---\n\nThis is what I searched for: ...",
            "doc_collection": {
                "documents": [
                    {
                        "corpus_id": "string",
                        "markdown": "string",
                        "relevance_judgement": Optional["integer"]
                        # might have many other non-relevant fields
                    }
                ]
            },
            "query": "string",
            "error": "string",
            "session_id": "string",
            "token_breakdown_by_model": dict["string", dict["string", "integer"]]
            }
        ```
        token_breakdown_by_model is a dict from model name to model usage.
        model usage is a dict from "prompt"/"completion"/"total" to the integer count
    ```

    Raises:
        httpx.HTTPStatusError: For API errors
        ValueError: For invalid responses
    """
    base_url = base_url.rstrip("/")

    async with httpx.AsyncClient(http2=True) as client:
        try:
            resp = await client.post(
                f"{base_url}/api/2/rounds",
                json={
                    "paper_description": query,
                    "operation_mode": "diligent",
                    "inserted_before": inserted_before,
                    "read_results_from_cache": read_results_from_cache,
                },
                timeout=timeout,
            )
            logger.info("Status code from rounds api: %s", resp.status_code)
            resp.raise_for_status()
        except httpx.HTTPStatusError as e:
            raise httpx.HTTPStatusError(
                f"request failed: {e.response.status_code}",
                request=e.request,
                response=e.response,
            )

        response = resp.json()

        # this will store it into the events list in the sample's log file
        transcript().info({"session_id": response.get("session_id")})

        # record cost
        for pf_model_name, token_breakdown_by_model in response.get(
            "token_breakdown_by_model", dict()
        ).items():
            usage = ModelUsage(
                input_tokens=token_breakdown_by_model.get("prompt"),
                output_tokens=token_breakdown_by_model.get("completion"),
                total_tokens=token_breakdown_by_model.get("total"),
            )
            record_model_usage_with_inspect(
                pf_model_name_2_litellm_model_name.get(pf_model_name, pf_model_name),
                usage,
            )
        return response


@tool
def paper_finder_tool(
    base_url: str,
    async_timeout: int = DEFAULT_ASYNC_TIMEOUT,
    inserted_before: str | None = None,
    read_results_from_cache: bool = True,
) -> Tool:
    async def run_paper_finder(paper_description: str) -> dict[str, Any]:
        response = dict()
        attempts = QUERY_ATTEMPTS
        while attempts > 0:
            try:
                response = await find_papers(
                    paper_description,
                    base_url=base_url,
                    timeout=async_timeout,
                    inserted_before=inserted_before,
                    read_results_from_cache=read_results_from_cache,
                )
                break
            except Exception as e:
                logger.error("Exception occurred during paper search: %s", e)
                attempts -= 1
                if attempts == 0:
                    raise e
        return response

    tdef = ToolDef(
        tool=run_paper_finder,
        name="paper_finder",
        description="""Paper Finding Agent: Searches academic papers based on user-specified criteria.

        SOME EXAMPLES OF SEARCH CAPABILITIES:
        - Description/keywords
        - Author name(s)
        - Paper title
        - Similarity to anchor papers (optional)

        QUERY MODIFIERS:
        - Temporal scope (e.g., "last 5 years")
        - Impact metrics (e.g., "highly-cited")
        - Search scope (single paper vs. collection)

        OPERATION:
        - Asynchronous processing
        - May ask clarifying questions
            * Direct clarification questions to users verbatim
            * User responses should populate the `answer` field, NOT `paper_description`

        For best results, always pass the FULL search requirements and preferences in the `paper_description` field.
        """,
        parameters={
            "paper_description": "A natural language query to find a paper or a collection of papers. This query should clearly express all of the user's search requirements and preferences.",
        },
    )

    return tdef.as_tool()
