src/fact_extractor.py

import json
import re
import time
from typing import Optional

import polars as pl

from src.config import Config
from src.llm import LocalLLM


class FactExtractor:
    def __init__(self, config: Config, llm: LocalLLM):
        self.config = config
        self.llm = llm

    def extract(self, chunks_df: pl.DataFrame) -> pl.DataFrame:
        if chunks_df.is_empty():
            return pl.DataFrame(
                schema={
                    "entity": pl.Utf8,
                    "attribute": pl.Utf8,
                    "value": pl.Utf8,
                    "unit": pl.Utf8,
                    "currency": pl.Utf8,
                    "min_quantity": pl.Float64,
                    "max_quantity": pl.Float64,
                    "condition": pl.Utf8,
                    "valid_from": pl.Utf8,
                    "valid_to": pl.Utf8,
                    "confidence": pl.Float64,
                    "source_url": pl.Utf8,
                    "source_title": pl.Utf8,
                    "chunk_id": pl.Int64,
                }
            )

        all_facts = []
        for row in chunks_df.to_dicts():
            text = row.get("text", "")
            url = row.get("url", "")
            title = row.get("title", "")
            chunk_id = row.get("chunk_id", 0)

            regex_facts = self._extract_regex(text, url, title, chunk_id)
            all_facts.extend(regex_facts)

        if not all_facts:
            return pl.DataFrame(
                schema={
                    "entity": pl.Utf8,
                    "attribute": pl.Utf8,
                    "value": pl.Utf8,
                    "unit": pl.Utf8,
                    "currency": pl.Utf8,
                    "min_quantity": pl.Float64,
                    "max_quantity": pl.Float64,
                    "condition": pl.Utf8,
                    "valid_from": pl.Utf8,
                    "valid_to": pl.Utf8,
                    "confidence": pl.Float64,
                    "source_url": pl.Utf8,
                    "source_title": pl.Utf8,
                    "chunk_id": pl.Int64,
                }
            )

        return pl.DataFrame(all_facts)

    def _extract_regex(
        self, text: str, source_url: str, source_title: str, chunk_id: int
    ) -> list[dict]:
        facts = []
        text_lower = text.lower()

        patterns = [
            (r"(?:price|cost|fee|rate|charge|payment|subscription|license)\s*(?:\#|no\.?|:)?\s*\$?\s*([\d,]+(?:\.\d{1,2})?)", "price", "USD", "per_item"),
            (r"(?:price|cost|fee|rate|charge|payment|subscription|license)\s*(?:\#|no\.?|:)?\s*£?\s*([\d,]+(?:\.\d{1,2})?)", "price", "GBP", "per_item"),
            (r"(?:price|cost|fee|rate|charge|payment|subscription|license)\s*(?:\#|no\.?|:)?\s*€?\s*([\d,]+(?:\.\d{1,2})?)", "price", "EUR", "per_item"),
            (r"\$?\s*([\d,]+(?:\.\d{1,2})?)\s*(?:per\s*(month|year|item|unit|day|week|hour|minute))", "price", "USD", "per_{}"),
            (r"£?\s*([\d,]+(?:\.\d{1,2})?)\s*(?:per\s*(month|year|item|unit|day|week|hour|minute))", "price", "GBP", "per_{}"),
            (r"€?\s*([\d,]+(?:\.\d{1,2})?)\s*(?:per\s*(month|year|item|unit|day|week|hour|minute))", "price", "EUR", "per_{}"),
            (r"(\d+)\s*(?:-\s*(\d+))?\s*(?:years?|months?)\s*(?:experience|required|needed|warranty)", "duration", None, "time"),
            (r"(?:speed|bandwidth|throughput|rate)\s*(?::|of|up to)?\s*(\d+)\s*(Mbps|Gbps|mbps|gbps)", "speed", None, "Mbps/Gbps"),
            (r"(?:capacity|storage|space|memory|ram)\s*(?::|of)?\s*(\d+)\s*(GB|TB|MB|gb|tb|mb)", "capacity", None, "GB/TB"),
            (r"(\d+)\s*(?:users?|seats?|licenses?|employees?)", "capacity", None, "per_seat"),
            (r"(\d+)\s*(?:-\s*(\d+))?\s*(?:days?)\s*(?:money.back|guarantee|refund|cancellation|notice)", "duration", None, "days"),
            (r"(?:minimum|max\.?|maximum)\s*(?:order|purchase|quantity|amount)\s*(?::|is)?\s*(\d+)", "quantity", None, "units"),
            (r"(?:free|included|trial)\s*(?:for\s*)?(\d+)\s*(?:days?|months?)", "duration", None, "trial_period"),
        ]

        for pattern, attr, currency, unit in patterns:
            for match in re.finditer(pattern, text, re.IGNORECASE):
                try:
                    value_raw = match.group(1).replace(",", "")
                    value = float(value_raw) if "." in value_raw else value_raw
                except (ValueError, IndexError):
                    continue

                actual_unit = unit.format(match.group(2)) if "{}" in unit else unit

                min_q = None
                max_q = None
                try:
                    min_q = float(match.group(1).replace(",", ""))
                    max_q = float(match.group(2).replace(",", "")) if match.group(2) else None
                except (IndexError, ValueError):
                    pass

                contextual_entity = self._guess_entity(text, text_lower)

                facts.append({
                    "entity": contextual_entity,
                    "attribute": attr,
                    "value": str(value),
                    "unit": actual_unit,
                    "currency": currency or "",
                    "min_quantity": min_q,
                    "max_quantity": max_q,
                    "condition": None,
                    "valid_from": None,
                    "valid_to": None,
                    "confidence": 0.7,
                    "source_url": source_url,
                    "source_title": source_title,
                    "chunk_id": chunk_id,
                })

        return facts

    def _guess_entity(self, text: str, text_lower: str) -> str:
        keywords_priority = [
            "openreach", "bt", "ee", "vodafone", "virgin media", "sky",
            "talktalk", "three", "o2", "plusnet", "shell energy",
        ]
        for kw in keywords_priority:
            if kw in text_lower:
                return kw.title()
        return "Unknown"

    def _extract_llm(self, text: str, source_url: str, source_title: str, chunk_id: int) -> list[dict]:
        if not self.llm:
            return []

        prompt = f"""Extract structured facts (prices, speeds, capacities, durations, quantities, ranges, conditions) from this text.

Text:
{text[:2000]}

Return a JSON array. Each fact has: entity, attribute, value, unit (or null), currency (or null), min_quantity (or null), max_quantity (or null), condition (or null), confidence (0-1)."""

        try:
            resp = self.llm.generate(
                messages=[
                    {"role": "system", "content": "You extract structured facts from text. Return only valid JSON arrays."},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=1024,
                temperature=0.01,
            )
            json_str = resp.strip()
            if json_str.startswith("```"):
                json_str = json_str.split("```")[1]
                if json_str.startswith("json"):
                    json_str = json_str[4:]
                json_str = json_str.strip()
            facts = json.loads(json_str)
            if isinstance(facts, list):
                for f in facts:
                    f["source_url"] = source_url
                    f["source_title"] = source_title
                    f["chunk_id"] = chunk_id
                return facts
        except Exception:
            pass

        return []