from backend.llm_local import get_llm

from rank_bm25 import BM25Okapi
import re

def _split_sentences(text: str):
    sents = re.split(r'(?<=[.!?])\s+', text)
    return [s.strip() for s in sents if s.strip()]

def extractive_sents(chunks: list[dict], query: str | None, max_sents=80):
    sents = []
    for ch in chunks:
        for s in _split_sentences(ch["text"]):
            sents.append((s, ch["meta"].get("page_num")))
    if not sents: return []
    corpus = [s for s,_ in sents]
    bm25 = BM25Okapi([c.split() for c in corpus])
    if query:
        scores = bm25.get_scores(query.split())
        ranked = [(corpus[i], sents[i][1]) for i,_ in sorted(enumerate(scores), key=lambda z:z[1], reverse=True)]
        return ranked[:max_sents]
    return sents[:max_sents]

def map_reduce(chunks: list[dict], focus: str, style: str, length: str):
    llm = get_llm()
    partials = []
    for ch in chunks:
        p = ch["meta"].get("page_num")
        prompt = f"Resuma em PT-BR, foco: {focus}, tamanho: {length}. Prefixe a página se houver.\n\n" \
                 f"{'[p.'+str(p)+'] ' if p else ''}{ch['text'][:2500]}"
        out = llm.create_chat_completion(messages=[{"role":"user","content":prompt}],
                                         temperature=0.2, max_tokens=450)
        partials.append(out["choices"][0]["message"]["content"])
    merged = "\n\n---\n\n".join(partials)
    final = llm.create_chat_completion(messages=[{"role":"user","content":
        f"Una os resumos abaixo num único resumo coerente, estilo: {style}, tamanho: {length}. "
        f"Conserve referências de páginas quando houver.\n\n{merged}"}],
        temperature=0.2, max_tokens=700)
    return final["choices"][0]["message"]["content"]