brain.py

import asyncio
import os
from knowledge_base import KnowledgeBase
from scraper import search_and_extract
from errors import get_logger, GenerAIError, ErrorCode, fmt_exc

log = get_logger("brain")

LOCAL_THRESHOLD = 0.40

# Modello HuggingFace da usare (fine-tunato o pubblico)
# Imposta HF_MODEL=tuo_username/generai-model nel file .env o come variabile d'ambiente
HF_MODEL = os.environ.get("HF_MODEL", "")
HF_TOKEN = os.environ.get("HF_TOKEN", "")


def _rerank(results: list, query: str) -> list:
    words = set(query.lower().split())
    for r in results:
        stored_q = r["metadata"].get("query", "")
        overlap = len(words & set(stored_q.lower().split()))
        r["score"] = r["distance"] - overlap * 0.12
    ranked = sorted(results, key=lambda x: x["score"])
    if ranked:
        log.debug("Rerank: best score=%.3f (dist=%.3f)", ranked[0]["score"], ranked[0]["distance"])
    return ranked


def _extract_sentences(text: str, query: str, max_chars: int = 700) -> str:
    words = set(query.lower().split())
    sentences = [s.strip() for s in text.replace("\n", ". ").split(".") if len(s.strip()) > 20]
    scored = [(sum(1 for w in words if w in s.lower()), s) for s in sentences]
    scored.sort(key=lambda x: -x[0])
    result = ""
    for _, s in scored:
        if len(result) + len(s) + 2 > max_chars:
            break
        result += s + ". "
    return result.strip() or text[:max_chars]


# ── HuggingFace LLM client (opzionale) ────────────────────────────────────────

class _HFClient:
    """Wrapper leggero per HuggingFace Inference API."""

    SYSTEM = (
        "Sei GenerAI, un assistente AI specializzato in lingua italiana. "
        "Rispondi in modo chiaro, preciso e sempre in italiano. "
        "Se ti viene fornito del contesto, basati su quello per rispondere."
    )

    def __init__(self, model: str, token: str):
        try:
            from huggingface_hub import InferenceClient
            self._client = InferenceClient(model=model, token=token or None)
            self._model = model
            log.info("HuggingFace LLM pronto: %s", model)
        except ImportError:
            raise GenerAIError(
                ErrorCode.BRAIN_ASK_FAILED,
                "huggingface_hub non installato. Esegui: pip install huggingface-hub",
            )
        except Exception as e:
            raise GenerAIError(
                ErrorCode.BRAIN_ASK_FAILED,
                f"Impossibile connettersi a HuggingFace ({model}): {fmt_exc(e)}",
                cause=e,
            )

    def generate(self, question: str, context: str = "") -> str:
        user_msg = question
        if context:
            user_msg = f"Contesto:\n{context}\n\nDomanda: {question}"

        messages = [
            {"role": "system",  "content": self.SYSTEM},
            {"role": "user",    "content": user_msg},
        ]
        log.debug("HF generate — modello=%s, context_len=%d", self._model, len(context))
        try:
            response = self._client.chat_completion(
                messages=messages,
                max_tokens=512,
                temperature=0.3,
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            raise GenerAIError(
                ErrorCode.BRAIN_ASK_FAILED,
                f"Generazione HF fallita: {fmt_exc(e)}",
                cause=e,
            )


# ── Brain ──────────────────────────────────────────────────────────────────────

class Brain:
    def __init__(self):
        log.info("Avvio Brain...")
        try:
            self.kb = KnowledgeBase()
        except GenerAIError as e:
            e.log(log)
            raise

        self._last_doc_id: str | None = None

        # Carica client HF se configurato
        self._hf: _HFClient | None = None
        if HF_MODEL:
            try:
                self._hf = _HFClient(HF_MODEL, HF_TOKEN)
            except GenerAIError as e:
                e.log(log)
                log.warning("Fallback alla ricerca semantica (HF non disponibile).")
        else:
            log.info("HF_MODEL non impostato — uso ricerca semantica locale.")

        log.info("Brain pronto. Modalità: %s", "LLM+KB" if self._hf else "KB+Web")

    async def ask(self, question: str, on_status=None) -> tuple[str, str]:
        """
        Returns (answer, status).
        on_status: callable opzionale async(msg: str) per aggiornamenti in tempo reale.
        status: "local" | "searched" | "unknown" | "error" | "llm"
        """
        async def emit(msg: str):
            if on_status:
                await on_status(msg)

        self._last_doc_id = None
        log.info("Domanda: %r", question)

        # ── 1. Ricerca KB locale ───────────────────────────────────────────────
        await emit("🔍 Cerco nella memoria locale...")
        try:
            results = _rerank(self.kb.search(question, n_results=12), question)
        except Exception as e:
            log.warning("KB search fallita: %s", fmt_exc(e))
            results = []

        best_context = ""
        best_source = ""

        if results and results[0]["score"] < LOCAL_THRESHOLD:
            best = results[0]
            best_context = best["answer"]
            best_source  = best["metadata"].get("source", "")
            await emit(f"✅ Trovato in memoria locale (score: {results[0]['score']:.2f})")
            log.info("Match KB locale (score=%.3f)", results[0]["score"])

            if self._hf:
                await emit("🤖 Genero risposta con LLM...")
                return await self._llm_answer(question, best_context, best_source, status="llm")

            answer = _extract_sentences(best_context, question)
            if best_source and best_source != "grammatica_italiana":
                answer += f"\n\n*Fonte: {best_source}*"
            await emit("💬 Risposta pronta!")
            return answer, "local"

        # ── 2. Ricerca web ─────────────────────────────────────────────────────
        await emit("🌐 Non trovato in memoria — ricerca sul web...")
        try:
            web_results = await asyncio.to_thread(search_and_extract, question)
        except Exception as e:
            err = GenerAIError(ErrorCode.BRAIN_ASK_FAILED, f"Ricerca web fallita: {fmt_exc(e)}", cause=e)
            err.log(log)
            await emit(f"❌ Errore ricerca web: {fmt_exc(e)}")
            return err.user_message(), "error"

        if not web_results:
            await emit("⚠️ Nessun risultato trovato sul web")
            log.warning("Nessun risultato web per: %r", question)
            return (
                "Non ho trovato informazioni su questo argomento nel mio database né sul web. "
                "Prova a riformulare la domanda.",
                "unknown",
            )

        await emit(f"📄 Trovati {len(web_results)} risultati — estraggo il testo...")
        combined = "\n\n---\n\n".join(f"[{r['title']}]\n{r['text']}" for r in web_results)
        sources   = ", ".join(r["url"] for r in web_results if r.get("url"))

        # ── 3. Genera risposta ─────────────────────────────────────────────────
        if self._hf:
            await emit("🤖 Genero risposta con LLM...")
            answer, status = await self._llm_answer(question, combined, sources, status="llm")
        else:
            await emit("✂️ Estraggo le frasi più rilevanti...")
            answer = _extract_sentences(combined, question)
            if sources:
                answer += f"\n\n*Fonte: {sources}*"
            status = "searched"

        await emit("💬 Risposta pronta!")
        return answer, status

    async def _llm_answer(
        self, question: str, context: str, source: str, status: str
    ) -> tuple[str, str]:
        try:
            answer = await asyncio.to_thread(self._hf.generate, question, context)
            if source and source != "grammatica_italiana":
                answer += f"\n\n*Fonte: {source}*"
            log.info("Risposta generata da LLM (%d chars)", len(answer))
            return answer, status
        except GenerAIError as e:
            e.log(log)
            # Fallback all'estrazione testuale
            log.warning("Fallback a estrazione testuale.")
            answer = _extract_sentences(context, question)
            if source and source != "grammatica_italiana":
                answer += f"\n\n*Fonte: {source}*"
            return answer, "searched"

    def give_feedback(self, positive: bool):
        if self._last_doc_id:
            self.kb.reinforce(self._last_doc_id, positive)
        else:
            log.debug("give_feedback chiamato senza _last_doc_id.")

    @property
    def kb_size(self) -> int:
        return self.kb.count()