from __future__ import annotations

import json
import logging
import re
from typing import List, Tuple, Dict


logger = logging.getLogger(__name__)


def _extract_json_from_markdown(response: str) -> str:
    if not isinstance(response, str):
        return str(response)
    pattern = r"```(?:json)?\s*\n?(.*?)\n?```"
    match = re.search(pattern, response, re.DOTALL)
    if match:
        return match.group(1).strip()
    return response.strip()


class RelevanceFilter:
    """RelevanceFilter batches chunk texts to an LLM and selects those with
    relevance scores above a threshold.
    """

    def __init__(self, llm_client):
        self.llm_client = llm_client

    def _build_prompt(self, question: str, chunks: List[Tuple[str, Dict]]) -> str:
        texts = [chunk_text for (chunk_text, _meta) in chunks]
        lines = []
        for idx, t in enumerate(texts, start=1):
            lines.append(f"Chunk {idx}: {t}")
        chunks_formatted = "\n".join(lines)
        prompt = (
            f"Given question '{question}' and these document chunks, rate each 0-10 for relevance. "
            f"Return JSON array of scores.\n{chunks_formatted}\n"
        )
        return prompt

    async def filter(
        self, question: str, chunks: List[Tuple[str, Dict]], threshold: float = 7.0
    ) -> List[Tuple[str, Dict]]:
        if not chunks:
            return []

        prompt = self._build_prompt(question, chunks)
        try:
            response = await self.llm_client.complete(prompt, temperature=0.0, step_name="RelevanceFilter")
        except Exception as exc:
            logger.error("RelevanceFilter LLM call failed: %s", exc)
            return []

        scores: List[float] = []
        try:
            response = _extract_json_from_markdown(response)
            parsed = json.loads(response)
            if not isinstance(parsed, list):
                return []
            for v in parsed:
                if isinstance(v, (int, float)):
                    scores.append(float(v))
                else:
                    return []
        except Exception as exc:
            logger.error("RelevanceFilter JSON parse failed: %s", exc)
            return []

        if len(scores) != len(chunks):
            return []

        result: List[Tuple[str, Dict]] = []
        for (chunk, meta), score in zip(chunks, scores):
            if score > threshold:
                result.append((chunk, meta))

        return result
