import re
import requests
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from bs4 import BeautifulSoup
import time
import yaml
import os

# Load prompt template
_PROMPT_TPL = None
try:
    with open("prompts.yaml", "r") as f:
        _PROMPT_TPL = yaml.safe_load(f)["answer_prompt"]
except Exception:
    _PROMPT_TPL = "Question: {question}\nContext: {context}\nAnswer:"

class SimpleAgent:
    """
    Lightweight agent:
    - uses a small seq2seq LLM (Flan-T5 small) to generate concise answers
    - uses a quick web retrieval (Wikipedia API and DuckDuckGo snippets) to build context
    - returns stripped answers ready for EXACT-MATCH scoring
    """

    def __init__(self, model_name="google/flan-t5-small"):
        self.model_name = model_name
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
            self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        except Exception as e:
            raise RuntimeError(f"Failed to load model {model_name}: {e}")

    def _wiki_search(self, query, sentences=2):
        """Simple Wikipedia summary fetch"""
        try:
            url = "https://en.wikipedia.org/api/rest_v1/page/summary/" + requests.utils.quote(query)
            r = requests.get(url, timeout=6)
            if r.status_code == 200:
                data = r.json()
                return data.get("extract", "")
        except Exception:
            pass
        return ""

    def _duckduckgo_snippets(self, query, max_chunks=2):
        """
        Lightweight scraping of DuckDuckGo HTML results for short snippets.
        (If blocked, it will fail gracefully — agent still works with model-only.)
        """
        try:
            headers = {"User-Agent": "Mozilla/5.0"}
            url = f"https://duckduckgo.com/html/?q={requests.utils.quote(query)}"
            r = requests.get(url, headers=headers, timeout=6)
            soup = BeautifulSoup(r.text, "html.parser")
            snippets = []
            for res in soup.select(".result__snippet")[:max_chunks]:
                snippets.append(res.get_text(separator=" "))
            return " ".join(snippets)
        except Exception:
            return ""

    def _clean_answer(self, text):
        # Keep it compact: strip whitespace and newlines, remove leading/trailing punctuation
        if not text:
            return ""
        a = text.strip()
        # remove multiple newlines/spaces
        a = re.sub(r"\s+", " ", a)
        # remove leading and trailing quotes or dashes
        a = a.strip(" \n\"'`-:;")
        return a

    def _build_context(self, question):
        # Use quick heuristics to extract keywords and fetch context
        # Try wikipedia, then duckduckgo
        kw = question.split("?")[0][:120]  # shorthand
        wiki = self._wiki_search(kw)
        ddg = self._duckduckgo_snippets(kw)
        context_parts = [p for p in [wiki, ddg] if p]
        return " ".join(context_parts)[:3000]  # limit context length

    def answer(self, question):
        """
        Return a single string: the final answer ONLY (no commentary).
        """
        context = self._build_context(question)
        prompt = _PROMPT_TPL.format(question=question, context=context)

        # generate
        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
        out = self.model.generate(**inputs, max_new_tokens=128, do_sample=False)
        text = self.tokenizer.decode(out[0], skip_special_tokens=True)
        return self._clean_answer(text)