"""SID-2 preview on Baseten.

Before using this template: implement the vector_search and keyword_search backend
functions, set ID_FIELD and SNIPPET_FIELD to your document's fields, and
customize the INSTRUCTIONS, QUESTION, and TOOL_EXAMPLES to fit your data. You can
also define additional tools, or extend existing tools with more parameters (e.g. a
filter). Just make sure to always update the tool's description and usage example.
The report_helpful_ids tool is fixed and cannot be changed or modified.

Run:
    BASETEN_API_KEY=your-api-key uv run --with sid-sdk --with openai reference_implementation_baseten.py
"""

import json
import os
import uuid
from concurrent.futures import ThreadPoolExecutor

from openai import OpenAI

from sid import DocumentCache, InvalidCharacterRange, parse_rendered_model_facing_id
from sid.document_cache import SNIPPET_SIZE_DEFAULT


class ToolError(Exception):
    """A mistake the *model* made (bad id, malformed ref, invalid arguments),
    these errors are recoverable and passed back to the model as a message."""

# ---------------------------------------------------------------------------
# The episode. TODO: your question. The model is served by Baseten,
# under one BASETEN_API_KEY.
# ---------------------------------------------------------------------------

QUESTION = "What government position was held by the woman who portrayed Corliss Archer in the film Kiss and Tell?"

BASE_URL = os.environ.get("BASETEN_BASE_URL", "https://inference.baseten.co/v1")
MODEL = "sid/sid-2-preview"

MAX_TURNS = 10
REWARD = "f1"  # the trained scoring prompt: "f1" (few, precise ids) | "ndcg" (many, best first)
REASONING_EFFORT = "low"  # "low" | "medium" | "high"

DEFAULT_LIMIT = 5  # results per search when the model does not ask for a number
MAX_PARALLEL_FETCHES = 8  # backend queries in flight at once, across a turn's search calls

# ---------------------------------------------------------------------------
# TODO: your document backend. The tools below call these functions;
# implement them against your database. A document is a dict with:
#   ID_FIELD     : your document id (any stable string; never shown to the model)
#   SNIPPET_FIELD: the document's canonical text. Return the identical string
#                   for the same id everywhere: the seen-ledger's character
#                   offsets index it.
#   extras       : (optional) any other fields, such as a "title", rendered
#                   as display attributes. Every cached field is shown to the
#                   model (ID_FIELD is popped before caching), so leave
#                   anything private out of your documents.
# ---------------------------------------------------------------------------

ID_FIELD = "id"
SNIPPET_FIELD = "content"


def vector_search(query: str, limit: int) -> list[dict]:
    """TODO: your semantic (embedding) search: the ``limit`` most relevant
    documents for ``query``, best first."""
    raise NotImplementedError("implement vector_search against your document store")


def keyword_search(query: str, limit: int) -> list[dict]:
    """TODO: your full-text keyword (BM25) search: the ``limit`` best
    matching documents for ``query``."""
    raise NotImplementedError("implement keyword_search against your document store")


# ---------------------------------------------------------------------------
# The system prompt: one sentence naming your subject matter, then the
# search -> repeat -> report scaffolding. TODO: edit the first paragraph to
# describe your data and goal; keep the steps.
# ---------------------------------------------------------------------------

INSTRUCTIONS = """You are an expert research assistant that is given a question and must use the provided search tools to find all documents needed to answer the question.

Steps:
1. Reflect on what information is needed to answer the question and use the search tools to find documents. Each document has an id.
2. Repeat step 1 until all documents necessary and sufficient to answer the question have been found. Take as many turns and searches as needed – you can make multiple searches per turn! Most questions will require multiple turns. Many will need more.
3. Use the report_helpful_ids tool to report the most helpful document ids. List the most helpful document ids first (important!).

The episode ends once report_helpful_ids is called."""


# ---------------------------------------------------------------------------
# The tool schemas. Each tool should be demonstrated in TOOL_EXAMPLES.
# ---------------------------------------------------------------------------

SNIPPET_SIZE_DESCRIPTION = (
    f"Words shown per result snippet (default {SNIPPET_SIZE_DEFAULT}). Must be a positive integer. "
    f"Pass a larger value to see more of each document."
)

SEARCH_TOOL = {
    "type": "function",
    "name": "search",
    "description": (
        "Semantic (vector) search over the dataset. Best for conceptual or paraphrased "
        "queries: a long, descriptive query of what you want beats a few keywords; use "
        "text_search when you know the distinctive words a match would contain. Result "
        "ids may be ranged references 'id#start:end' naming the exact character span of "
        "the document shown; stretches you have already been shown are replaced by "
        "'[seen: \"#start:end\"]' markers instead of repeating, and doc_length gives the "
        "document's total character count."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "query": {"type": "string", "description": "The query to search the dataset with"},
            "limit": {"type": "integer", "description": "The number of results to return"},
            "snippet_size": {"type": "integer", "description": SNIPPET_SIZE_DESCRIPTION},
        },
        "required": ["query"],
    },
}

TEXT_SEARCH_TOOL = {
    "type": "function",
    "name": "text_search",
    "description": (
        "Full-text keyword search over the dataset, ranked by BM25. Provide the terms you "
        "want matched as plain words separated by spaces: the distinctive content words a "
        "matching document would contain. Result ids may be ranged references 'id#start:end' naming the exact "
        "character span of the document shown; stretches you have already been shown are "
        "replaced by '[seen: \"#start:end\"]' markers instead of repeating, and doc_length "
        "gives the document's total character count."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": (
                    "Plain space-separated content words to match (BM25)."
                ),
            },
            "limit": {"type": "integer", "description": "The number of results to return"},
            "snippet_size": {"type": "integer", "description": SNIPPET_SIZE_DESCRIPTION},
        },
        "required": ["query"],
    },
}

READ_TOOL = {
    "type": "function",
    "name": "read",
    "description": "Search results show only short snippets. Read up to 10,000 characters of a document per call. A bare ID starts at character 0; a ranged reference starts at its requested offset. Use further ranged reads to continue.",
    "parameters": {
        "type": "object",
        "properties": {
            "id": {
                "type": "string",
                "description": (
                    "The id of the document to read. Also accepts a ranged reference "
                    "'id#start:end' (character offsets, e.g. '123456#800:1600') to read only "
                    "that part of the document's text; a range running past the end of the "
                    "document is trimmed to the part that overlaps it."
                ),
            },
        },
        "required": ["id"],
    },
}

# The fixed terminal tool: do not change its name, description or parameters.
REPORT_TOOL = {
    "type": "function",
    "name": "report_helpful_ids",
    "description": (
        "Report the most helpful document ids from the search results. This is required in "
        "the last step. Each id may be reported at most once; repeated ids are penalized. "
        "Ids may be whole documents ('123456') or ranged references ('123456#800:1600'); a "
        "ranged reference counts as its whole document, so there is no need to strip the "
        "'#start:end' part."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "ids": {
                "type": "array",
                "description": "The ids of the most helpful documents from the search results.",
                "items": {"type": "string"},
            },
        },
        "required": ["ids"],
    },
}

TOOLS = [SEARCH_TOOL, TEXT_SEARCH_TOOL, READ_TOOL, REPORT_TOOL]
TOOLS_BY_NAME = {tool["name"]: tool for tool in TOOLS}

# Few-shot calls, sent on every request via the sid extension. TODO: write search
# examples that fit your data; keep the read and report examples as they are
TOOL_EXAMPLES = [
    {
        "caption": "Example: a semantic search with a long, descriptive query:",
        "name": "search",
        "arguments": {
            "query": "the actress who played Corliss Archer in the 1945 film, and the government post she later held",
            "limit": 5,
        },
    },
    {
        "caption": "Example: a keyword search:",
        "name": "text_search",
        "arguments": {"query": "Kiss and Tell 1945 Shirley Temple Corliss Archer", "limit": 5},
    },
    {
        "caption": "Example: reading one character range:",
        "name": "read",
        "arguments": {"id": "placeholder_1#800:1600"},
    },
    {
        "caption": "Example: ending the episode with a report:",
        "name": "report_helpful_ids",
        "arguments": {"ids": ["placeholder_1", "placeholder_2", "placeholder_3"]},
    },
]

# ---------------------------------------------------------------------------
# The tools, split in two: fetch_results is the slow backend call and runs in
# parallel; everything below it touches the cache and runs serially, in call
# order.
# ---------------------------------------------------------------------------


def fetch_results(name: str, arguments: dict) -> list[dict] | None:
    """The backend call behind one search tool call. `read` and the report
    have nothing to fetch (reads are served from the document cache), so they
    get ``None``."""
    if name not in ("search", "text_search"):
        return None

    # Catch mistyped arguments as a ToolError that the model can recover from.
    properties = TOOLS_BY_NAME[name]["parameters"]["properties"]
    for key, value in arguments.items():
        expected = str if properties[key]["type"] == "string" else int
        if not isinstance(value, expected) or isinstance(value, bool):
            raise ToolError(
                f"Invalid {name} arguments: {key} must be of type {properties[key]['type']}, "
                f"got {type(value).__name__}"
            )

    query = arguments.get("query", "")  # tool_arguments strips ""; treat missing as empty
    if not query:
        raise ToolError(
            "search requires a non-empty `query` to rank documents against. "
            "To fetch a specific document by id, use the `read` tool. "
            "To rank candidates against the question, call `search` with a meaningful query."
        )

    snippet_size = arguments.get("snippet_size", SNIPPET_SIZE_DEFAULT)
    if snippet_size <= 0:
        raise ToolError(
            f"Invalid {name} arguments: snippet_size must be a positive integer, got {snippet_size}"
        )

    limit = arguments.get("limit", DEFAULT_LIMIT)
    return vector_search(query, limit) if name == "search" else keyword_search(query, limit)


def render_results(cache: DocumentCache, results: list[dict], query: str, snippet_size: int) -> str:
    """One <doc> per result, in backend order. Shared by both search
    tools; only their fetch differs."""
    rendered = []
    for result in results:
        if not isinstance(result, dict) or not isinstance(result.get(ID_FIELD), str) or not result[ID_FIELD]:
            raise ValueError(f"backend document must carry a non-empty string {ID_FIELD!r}: {result!r}")

        # Keep backend results intact when searches reuse the same dictionaries.
        document = dict(result)

        # Popped, not read: every cached field renders, and your database ids must not.
        data_id = document.pop(ID_FIELD)
        content = document.get(SNIPPET_FIELD)
        if not isinstance(content, str) or not content:
            # Not a model mistake: the seen-ledger's offsets index this string.
            raise ValueError(f"backend document {data_id!r} must carry non-empty {SNIPPET_FIELD!r} text")

        # Searches routinely return the same document, and the cache is add-once.
        if data_id not in cache:
            cache.add_document(data_id, document)

        # BM25 snippet, with already-shown stretches collapsed to [seen: ...].
        view = cache.apply_snippet(
            data_id,
            SNIPPET_FIELD,
            query,
            snippet_size=snippet_size,
        )
        rendered.append(view.render_xml())
        cache.update_seen(view)  # later views of this span collapse to [seen: ...]

    return "\n".join(rendered) if rendered else "No results found. Try a different query."


def parse_reference(reference) -> tuple[str, tuple[int, int] | None]:
    """The model's ref in (``'R'`` or ``'R#start:end'``), split into id and
    optional char range. The SDK parser raises ``ValueError`` on malformed
    shapes with a descriptive message already written."""
    try:
        return parse_rendered_model_facing_id(reference)
    except ValueError as error:
        raise ToolError(str(error)) from None


def read(cache: DocumentCache, reference: str | None) -> str:
    """Fetch one document (or a character range of it) by the id a search
    returned. Explicit reads are not masked: the model asked for this span, so
    it renders up to 10,000 characters. Model mistakes (bad ref grammar,
    an unknown or hallucinated id, an unusable range) raise ToolError."""
    model_facing_id, char_range = parse_reference(reference)  # ToolError on bad grammar

    if not cache.contains_model_facing_id(model_facing_id):
        raise ToolError(
            f"Invalid document id '{model_facing_id}'. "
            f"You can only read documents returned by search. "
            f"Use 'search' or 'text_search' first."
        )

    data_id = cache.to_data_id(model_facing_id)  # the one conversion back into your ids

    if char_range is not None:
        # Lenient mode trims partially overlapping ranges to the document.
        # Invalid ranges become recoverable tool errors.
        try:
            char_range = cache.resolve_char_range(data_id, SNIPPET_FIELD, char_range)
        except InvalidCharacterRange as error:
            raise ToolError(
                f"{error}. Re-read '{model_facing_id}' without a range to read from the "
                f"start, or pass a range inside 0:doc_length."
            ) from None

    # Cap document text per read; further ranged reads can continue from here.
    if char_range is not None:
        start, end = char_range
        char_range = (start, min(end, start + 10_000))
    elif len(cache.get_document(data_id)[SNIPPET_FIELD]) > 10_000:
        char_range = (0, 10_000)

    view = cache.get_single_span_document_view(
        data_id,
        SNIPPET_FIELD,
        snippet_display_span=char_range,
    )
    rendered = view.render_xml()
    cache.update_seen(view)
    return rendered


def report(cache: DocumentCache, ids) -> list[str]:
    """The terminal report: ids are doc-granular (a ranged ref denotes its
    document), deduplicated preserving order, and translated back to your
    data_ids. IDs outside the model-facing ID map are dropped."""
    if not isinstance(ids, list) or not all(isinstance(reference, str) for reference in ids):
        raise ToolError("ids must be a list of strings")
    if not ids:
        raise ToolError("ids list cannot be empty")

    reported: list[str] = []
    for reference in ids:
        model_facing_id, _ = parse_reference(reference)  # ToolError on bad grammar
        if model_facing_id not in reported:
            reported.append(model_facing_id)

    return [
        cache.to_data_id(model_facing_id)
        for model_facing_id in reported
        if cache.contains_model_facing_id(model_facing_id)
    ]


# ---------------------------------------------------------------------------
# The agent loop: one Responses API call per turn, tool outputs paired back
# by call_id. The episode ends when the model reports.
#
# Requests are stateless (store=False): nothing is retained server-side, so
# the full conversation is resent on every call. The model's reasoning comes
# back as an opaque ``encrypted_content`` blob on each reasoning item, and
# carries over by echoing the output items back verbatim in the next input.
#
# The prompt (instructions, tools, reasoning effort, tool_examples, reward) is
# fixed for the life of the conversation. Build the request dict once and send
# it unchanged on every turn; only `input` grows.
# ---------------------------------------------------------------------------


def main() -> None:
    client = OpenAI(base_url=BASE_URL, api_key=os.environ["BASETEN_API_KEY"])
    cache = DocumentCache(range_mode="lenient")  # this episode's id map, document cache and seen-ledger

    session_id = uuid.uuid4().hex # KV cache routing information: one session per episode

    # Caps how many queries your backend sees at once (1 makes the turn serial).
    pool = ThreadPoolExecutor(max_workers=MAX_PARALLEL_FETCHES)

    # One request shape, only update the history at each turn.
    request = dict(
        model=MODEL,
        max_output_tokens=4096,
        store=False,
        include=["reasoning.encrypted_content"],
        instructions=INSTRUCTIONS,
        tools=TOOLS,
        # Reasoning effort, tool_examples, and the target reward are managed server-side
        reasoning={"effort": REASONING_EFFORT},
        extra_body={"sid": {"tool_examples": TOOL_EXAMPLES, "reward": REWARD}},
        extra_headers={"x-session-affinity": session_id},
    )

    # The whole conversation as Responses API input items, grown every turn.
    history: list[dict] = [{"role": "user", "content": QUESTION}]

    print(f"Q: {QUESTION}")
    response = client.responses.create(input=history, **request)

    for turn in range(1, MAX_TURNS + 1):
        print_model_output(turn, response)
        calls = [item for item in response.output if item.type == "function_call"]
        arguments = [tool_arguments(call) for call in calls]

        # Execute the model's tool calls in parallel
        fetches = [pool.submit(fetch_results, call.name, args) for call, args in zip(calls, arguments, strict=True)]

        # Compose observations serially, in call order: what earlier calls show, later calls mask.
        observations: list[str] = []
        reported: list[str] | None = None
        for call, args, fetch in zip(calls, arguments, fetches, strict=True):
            try:
                if call.name in ("search", "text_search"):
                    observation = render_results(
                        cache,
                        fetch.result(),
                        args["query"],
                        args.get("snippet_size", SNIPPET_SIZE_DEFAULT),
                    )

                elif call.name == "read":
                    observation = read(cache, args.get("id"))

                elif call.name == "report_helpful_ids":
                    if len(calls) > 1:
                        raise ToolError("Combining report tool call with other tool calls is not allowed.")
                    if turn == 1:
                        raise ToolError("Cannot report in the first turn. Must make a search first.")
                    reported = report(cache, args.get("ids"))
                    observation = "reported ids submitted: " + ",".join(reported)

                else:
                    observation = f"Tool '{call.name}' is not a valid tool. Available tools: {list(TOOLS_BY_NAME)}"

            except ToolError as e:
                # A recoverable mistake made by the model.
                observation = f"Error executing {call.name}: {e}"

            observations.append(observation)
            print(f"  [{call.name}] -> {observation}")  # verbatim: what you see is what the model receives

        # Termination conditions
        if reported is not None:
            print_report(cache, reported)
            return

        if turn == MAX_TURNS:
            print(f"\n[max turns reached: {MAX_TURNS}: episode ends without a report]")
            return

        if response.status == "incomplete":
            print("\n[generation hit max_output_tokens; ending the episode]")
            return

        # Non-terminal turn: append the model's output verbatim (the reasoning
        # items carry their encrypted_content), then feed back the tool outputs
        # (paired by call_id) plus the turn budget indicator as a user message.
        history += [item.model_dump(mode="json", exclude_none=True) for item in response.output]
        history += [
            {"type": "function_call_output", "call_id": call.call_id, "output": observation}
            for call, observation in zip(calls, observations, strict=True)
        ]

        remaining = MAX_TURNS - turn
        messages = [f"You have {remaining} {'turn' if remaining == 1 else 'turns'} remaining out of {MAX_TURNS} total."]

        if turn == MAX_TURNS - 1:
            messages.append(
                f"This is your last turn: you must call {REPORT_TOOL['name']} now. "
                f"Anything you have found but not reported is lost when this turn ends."
            )

        history += [{"role": "user", "content": message} for message in messages]
        for message in messages:
            print(f"[user] {message}")

        response = client.responses.create(input=history, **request)


def tool_arguments(call) -> dict:
    """Parse a call's JSON arguments: keep only the parameters the tool's
    schema declares, and drop the ""/null "defaults" LLMs like to send.
    Anything else the model got wrong comes back as a ToolError from the
    fetch or the tool."""
    try:
        arguments = json.loads(call.arguments)
    except json.JSONDecodeError:
        return {}
    schema = TOOLS_BY_NAME.get(call.name)
    if not isinstance(arguments, dict) or schema is None:
        return {}
    properties = schema["parameters"]["properties"]
    return {key: value for key, value in arguments.items() if key in properties and value not in ("", None)}


def print_model_output(turn: int, response) -> None:
    """The model's side of one turn: its thinking, any prose, and its calls."""
    print(f"\n--- turn {turn} (status={response.status}) " + "-" * 40)
    for item in response.output:
        if item.type == "reasoning":
            for part in getattr(item, "content", None) or []:
                print(f"[thinking] {part.text}")
        elif item.type == "message":
            print(f"[assistant] {response.output_text}")
        elif item.type == "function_call":
            print(f">>> {item.name}({item.arguments})")


def print_report(cache: DocumentCache, reported_ids: list[str]) -> None:
    """The episode's answer: your database ids, most helpful first."""
    print(f"\nReported for: {QUESTION}")
    for data_id in reported_ids:
        document = cache.get_document(data_id)
        print(f"   {data_id}  {document.get('title', '')}")


if __name__ == "__main__":
    main()