src/answer_generator.py

import json
import re

import polars as pl

from src.config import Config
from src.llm import LocalLLM
from src.query_router import QueryRouter
from src.retriever import Retriever


class AnswerGenerator:
    def __init__(
        self,
        config: Config,
        llm: LocalLLM,
        retriever: Retriever,
        facts_df: pl.DataFrame,
        pages_df: pl.DataFrame,
    ):
        self.config = config
        self.llm = llm
        self.router = QueryRouter()
        self.retriever = retriever
        self.facts_df = facts_df
        self.pages_df = pages_df

    def answer(self, query: str) -> str:
        route = self.router.route(query)
        retrieved = self.retriever.retrieve(query)

        numeric_results = None
        if route in ("numeric", "calculation", "range", "comparison"):
            numeric_results = self._search_facts(query)

        if route == "numeric":
            return self._answer_numeric(query, retrieved, numeric_results)
        elif route == "calculation":
            return self._answer_calculation(query, retrieved, numeric_results)
        elif route == "comparison":
            return self._answer_comparison(query, retrieved, numeric_results)
        elif route == "range":
            return self._answer_range(query, retrieved, numeric_results)
        elif route == "policy":
            return self._answer_policy(query, retrieved)
        elif route == "summary":
            return self._answer_summary(query, retrieved)
        else:
            return self._answer_general(query, retrieved)

    def _search_facts(self, query: str) -> list[dict]:
        if self.facts_df.is_empty():
            return []

        query_lower = query.lower()
        numbers = re.findall(r"\d+", query)
        query_numbers = [float(n) for n in numbers]

        relevant_facts = []
        for row in self.facts_df.to_dicts():
            entity = row.get("entity", "").lower()
            attribute = row.get("attribute", "").lower()

            entity_words = entity.split()
            if any(w in query_lower for w in entity_words if len(w) > 2):
                relevant_facts.append(row)
                continue

            if attribute in query_lower:
                relevant_facts.append(row)
                continue

            if query_numbers:
                mn = row.get("min_quantity")
                mx = row.get("max_quantity")
                if mn is not None and mx is not None:
                    for qn in query_numbers:
                        if mn <= qn <= mx:
                            relevant_facts.append(row)
                            break
                elif mn is not None and mx is None:
                    if any(qn >= mn for qn in query_numbers):
                        relevant_facts.append(row)

        return relevant_facts

    def _answer_numeric(self, query: str, retrieved: list[dict], facts: list[dict]) -> str:
        if facts:
            answer = self._build_numeric_answer(query, facts, retrieved)
            if answer:
                return answer

        context = self._format_context(retrieved[:5])
        return self._llm_answer(query, context)

    def _answer_calculation(self, query: str, retrieved: list[dict], facts: list[dict]) -> str:
        if facts:
            calc_result = self._try_calculate(query, facts)
            if calc_result:
                return calc_result

        context = self._format_context(retrieved[:5])
        return self._llm_answer(
            query,
            context,
            instruction="If the question requires calculation, perform the calculation step by step and show your reasoning.",
        )

    def _answer_comparison(self, query: str, retrieved: list[dict], facts: list[dict]) -> str:
        if facts:
            items = set(f["entity"] for f in facts if f.get("entity"))
            if len(items) >= 2:
                comparison = self._build_comparison(query, facts)
                if comparison:
                    return comparison

        context = self._format_context(retrieved[:7])
        return self._llm_answer(query, context)

    def _answer_range(self, query: str, retrieved: list[dict], facts: list[dict]) -> str:
        if facts:
            mn = min((f.get("min_quantity") or 0 for f in facts if f.get("min_quantity") is not None), default=None)
            mx = max((f.get("max_quantity") or 0 for f in facts if f.get("max_quantity") is not None), default=None)
            vals = sorted(set(f["value"] for f in facts if f.get("value")), key=float)

            if mn is not None and mx is not None:
                return (
                    f"The available range is {mn} to {mx}.\n"
                    + f"Prices in this range: {', '.join(vals)}\n"
                    + self._cite_sources(facts)
                )

        context = self._format_context(retrieved[:5])
        return self._llm_answer(query, context)

    def _answer_policy(self, query: str, retrieved: list[dict]) -> str:
        context = self._format_context(retrieved[:7])
        return self._llm_answer(
            query,
            context,
            instruction="Answer based on policy/terms information. If the answer depends on missing information (e.g., user's specific situation), ask a follow-up question to clarify.",
        )

    def _answer_summary(self, query: str, retrieved: list[dict]) -> str:
        context = self._format_context(retrieved[:7])
        return self._llm_answer(query, context)

    def _answer_general(self, query: str, retrieved: list[dict]) -> str:
        if not retrieved:
            return (
                "I don't have enough information from the crawled data to answer that question. "
                "Please try crawling a website first, or rephrase your question."
            )
        context = self._format_context(retrieved[:5])
        return self._llm_answer(query, context)

    def _try_calculate(self, query: str, facts: list[dict]) -> str | None:
        numbers = re.findall(r"\d+[.,]?\d*", query)
        if not numbers:
            return None

        query_lower = query.lower()
        has_price_per = any("price" in f.get("attribute", "").lower() or "per" in f.get("attribute", "").lower() for f in facts)
        has_unit_price = any(f.get("unit") == "per_item" for f in facts)

        if has_price_per or has_unit_price:
            for f in facts:
                attr = f.get("attribute", "").lower()
                if "price" in attr or "per" in attr:
                    try:
                        unit_price = float(f["value"])
                    except (ValueError, TypeError):
                        continue

                    qty = None
                    for n_str in numbers:
                        try:
                            candidate = float(n_str.replace(",", "."))
                            mn = f.get("min_quantity")
                            mx = f.get("max_quantity")
                            if mn is not None and mx is not None:
                                if mn <= candidate <= mx:
                                    qty = candidate
                                    break
                            else:
                                qty = candidate
                                break
                        except ValueError:
                            continue

                    if qty is not None:
                        total = qty * unit_price
                        currency = f.get("currency", "")
                        currency_symbol = {"USD": "$", "EUR": "€", "GBP": "£"}.get(currency, "")
                        return (
                            f"**Calculation:** {qty} × {unit_price} = {total}\n\n"
                            + f"For {qty} items at {currency_symbol}{unit_price} each, "
                            + f"the total is **{currency_symbol}{total:.2f}**.\n\n"
                            + self._cite_sources([f])
                        )

        return None

    def _build_numeric_answer(self, query: str, facts: list[dict], retrieved: list[dict]) -> str:
        for f in facts:
            attr = f.get("attribute", "").lower()
            if "price" in attr:
                value = f.get("value", "")
                currency = f.get("currency", "USD")
                unit = f.get("unit", "")
                mn = f.get("min_quantity")
                mx = f.get("max_quantity")
                entity = f.get("entity", "")

                parts = [f"The price for **{entity}** is **{value} {currency}**"]
                if unit:
                    parts.append(f"({unit})")
                if mn is not None and mx is not None:
                    if mn != mx:
                        parts.append(f"for quantities between {mn} and {mx}")
                    else:
                        parts.append(f"for quantity {mn}")
                parts.append(f"\n\n{self._cite_sources([f])}")
                return " ".join(parts)

        return ""

    def _build_comparison(self, query: str, facts: list[dict]) -> str:
        by_entity: dict[str, list[dict]] = {}
        for f in facts:
            ent = f.get("entity", "unknown")
            by_entity.setdefault(ent, []).append(f)

        lines = []
        for entity, efacts in by_entity.items():
            prices = [f for f in efacts if "price" in f.get("attribute", "").lower()]
            if prices:
                vals = [f"{p['value']} {p.get('currency', '')}" for p in prices]
                lines.append(f"- **{entity}**: {', '.join(vals)}")

        if lines:
            lines.append("")
            lines.append(self._cite_sources(facts))
            return "\n".join(lines)
        return ""

    def _llm_answer(self, query: str, context: str, instruction: str = "") -> str:
        if not context:
            return (
                "I don't have enough information from the crawled data to answer that question. "
                "Please try crawling a website first."
            )

        system_prompt = "You are a precise QA assistant that answers questions based ONLY on the provided crawled content."
        if instruction:
            system_prompt += f"\n\n{instruction}"

        try:
            text = self.llm.generate(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"},
                ],
                max_tokens=self.config.llm_max_tokens,
                temperature=self.config.llm_temperature,
            )
            if text:
                return text
        except Exception:
            pass

        return "I'm sorry, I couldn't generate an answer at this time. Please try again later."

    def _format_context(self, retrieved: list[dict]) -> str:
        if not retrieved:
            return ""

        parts = []
        seen_urls = set()
        for i, r in enumerate(retrieved, 1):
            text = r.get("text", "")[:500]
            url = r.get("url", "")
            title = r.get("title", "")

            source = f"[Source: {title}]({url})" if url else ""
            if url:
                seen_urls.add(url)

            parts.append(f"[{i}] {text}\n{source}")

        return "\n\n---\n\n".join(parts)

    def _cite_sources(self, facts: list[dict]) -> str:
        urls = set()
        for f in facts:
            url = f.get("source_url", "")
            if url:
                urls.add(url)
        if urls:
            return "**Sources:** " + ", ".join(sorted(urls))
        return ""