"""Kodit retrieval integration for mini-swe-agent benchmarks."""

import structlog

from benchmark.swebench.instance import SWEBenchInstance
from benchmark.swebench.retriever import KoditRetriever, RetrievedSnippet


class KoditContextProvider:
    """Pre-fetches Kodit snippets and formats them for mini-swe-agent."""

    def __init__(
        self,
        kodit_base_url: str,
        top_k: int = 10,
        timeout: float = 30.0,
    ) -> None:
        """Initialize with Kodit server URL."""
        self._retriever = KoditRetriever(
            kodit_base_url=kodit_base_url,
            timeout=timeout,
        )
        self._top_k = top_k
        self._log = structlog.get_logger(__name__)

    def get_context(self, instance: SWEBenchInstance) -> str:
        """Retrieve and format Kodit snippets for an instance."""
        self._log.info(
            "Fetching Kodit context",
            instance_id=instance.instance_id,
            top_k=self._top_k,
        )

        snippets = self._retriever.retrieve(instance, top_k=self._top_k)

        if not snippets:
            self._log.warning(
                "No snippets retrieved",
                instance_id=instance.instance_id,
            )
            return ""

        self._log.info(
            "Retrieved snippets",
            instance_id=instance.instance_id,
            count=len(snippets),
        )

        return self._format_snippets(snippets)

    def _format_snippets(self, snippets: list[RetrievedSnippet]) -> str:
        """Format snippets as markdown code blocks."""
        parts = []
        for snippet in snippets:
            parts.append(f"### {snippet.file_path}")
            parts.append("```")
            parts.append(snippet.content)
            parts.append("```")
            parts.append("")

        return "\n".join(parts)

    def augment_problem_statement(
        self,
        instance: SWEBenchInstance,
    ) -> str:
        """Augment the problem statement with Kodit-retrieved context."""
        context = self.get_context(instance)

        if not context:
            return instance.problem_statement

        return f"""## Relevant Code Context
The following code snippets from the repository may be relevant to this issue:

{context}

## Issue Description
{instance.problem_statement}"""
