"""
chatbot.py — Answers naturally from what the AI has learned.
Filters out junk, synthesizes real paragraph answers.
"""
import json, re, math, os, random
from collections import Counter, defaultdict
from html import unescape

KNOWLEDGE_FILE = 'knowledge.jsonl'

STOPWORDS = {
    'a','an','the','is','it','in','on','at','to','for','of','and','or','but',
    'was','are','were','be','been','have','has','had','do','does','did','will',
    'would','could','should','may','might','that','this','these','those','with',
    'from','by','as','not','also','than','then','so','if','when','what','how',
    'who','which','its','their','our','your','my','his','her','we','they','he',
    'she','you','i','me','him','us','them','said','says',
}

# Sentences that are website UI / taglines — not real knowledge
JUNK = re.compile(
    r'(sign up|subscribe|click here|breaking news:|stay informed|'
    r'all rights reserved|privacy policy|cookie policy|advertisement|'
    r'follow us|share this|read more|learn more|'
    r'has everything you need|stay up to date|get the latest|'
    r'news today ap|ap news|everything you need to know for|'
    r'top breaking|live science\.|sciencedaily|newscientist|'
    r'latest \w+ news articles today)',
    re.IGNORECASE
)

def tokenize(text):
    text = re.sub(r'[^\w\s]', ' ', text.lower())
    return [w for w in text.split() if w not in STOPWORDS and len(w) > 2]

def scrub(text):
    """Clean HTML artifacts from stored text."""
    text = unescape(str(text))
    text = re.sub(r'&#?[a-zA-Z0-9]+;', ' ', text)
    text = re.sub(r'<[^>]+>', ' ', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def is_junk(sentence):
    if len(sentence.split()) < 7: return True
    if JUNK.search(sentence): return True
    return False


class KnowledgeBase:
    def __init__(self):
        self.docs=[]; self.doc_tokens=[]; self.idf={}; self.last_count=0
        self._load()

    def _load(self):
        if not os.path.exists(KNOWLEDGE_FILE): return
        new_docs = []
        try:
            with open(KNOWLEDGE_FILE,'r',encoding='utf-8') as f:
                for line in f:
                    l=line.strip()
                    if l:
                        try: new_docs.append(json.loads(l))
                        except: pass
        except: return
        if len(new_docs)==self.last_count: return
        self.docs=new_docs; self.last_count=len(new_docs)
        self.doc_tokens=[tokenize(scrub(d.get('text',''))) for d in self.docs]
        N=len(self.docs)
        df=defaultdict(int)
        for tok in self.doc_tokens:
            for w in set(tok): df[w]+=1
        self.idf={w:math.log((N+1)/(c+1))+1 for w,c in df.items()}

    def search(self, query, top_k=10):
        self._load()
        if not self.docs: return []
        q=tokenize(query)
        if not q: return []
        def sc(dt):
            tf=Counter(dt); n=max(len(dt),1)
            return sum((tf[w]/n)*self.idf.get(w,0) for w in q if w in tf)
        ranked=sorted(range(len(self.docs)),key=lambda i:-sc(self.doc_tokens[i]))
        return [(self.docs[i], sc(self.doc_tokens[i]))
                for i in ranked[:top_k] if sc(self.doc_tokens[i])>0]

    def get_stats(self):
        self._load()
        return {
            'total':      len(self.docs),
            'categories': dict(Counter(d.get('category','other') for d in self.docs)),
            'sources':    dict(Counter(d.get('source','unknown') for d in self.docs)),
        }


def build_answer(query, results):
    """
    Pull the best real sentences from retrieved docs and write a natural answer.
    Skips junk/tagline sentences. Returns None if nothing good found.
    """
    q_tokens = set(tokenize(query))
    candidates = []

    for doc, relevance in results:
        text = scrub(doc.get('text',''))
        # Split into sentences
        sents = re.split(r'(?<=[.!?])\s+', text)
        for s in sents:
            s = s.strip()
            if is_junk(s): continue
            overlap = len(q_tokens & set(tokenize(s)))
            # Score = topic overlap + doc relevance
            candidates.append((overlap * 2 + relevance, s))

    if not candidates:
        return None

    # Deduplicate, pick top unique sentences
    seen = set(); good = []
    for score, sent in sorted(candidates, reverse=True):
        key = ' '.join(sent.lower().split()[:6])
        if key not in seen:
            seen.add(key)
            good.append(sent)
        if len(good) >= 5: break

    if not good:
        return None

    # Write as flowing paragraphs (not bullets)
    if len(good) == 1:
        return good[0]
    elif len(good) <= 3:
        return ' '.join(good)
    else:
        # Split into two paragraphs
        return ' '.join(good[:2]) + '\n\n' + ' '.join(good[2:])


class RAGChatbot:
    def __init__(self):
        self.kb = KnowledgeBase()

    def _get_reply(self, query, results, stats):
        total = stats['total']

        if total == 0:
            return ("My brain is empty right now 🧠\n\n"
                    "The app is already auto-training — just wait a minute "
                    "while I read some articles, then ask again!")

        answer = build_answer(query, results) if results else None

        if not answer:
            topics = ', '.join(list(stats['categories'].keys())[:5])
            return (f"I've read **{total} articles** but haven't learned enough "
                    f"about *\"{query}\"* yet.\n\n"
                    f"I know most about: **{topics}**.\n\n"
                    f"Keep me running and I'll learn more!")

        return answer

    def chat(self, user_message, history):
        history = list(history or [])
        msg = user_message.strip()
        if not msg: return history

        lower = msg.lower()
        stats = self.kb.get_stats()

        if lower in ('hi','hello','hey','sup','yo','hiya','howdy','helo','hai'):
            total = stats['total']
            if total == 0:
                bot_reply = "Hey! 👋 I'm just starting up — give me a minute to read some articles and I'll be ready to answer questions!"
            else:
                topics = ', '.join(list(stats['categories'].keys())[:4])
                bot_reply = f"Hey! 👋 I've read **{total} articles** so far. I know quite a bit about **{topics}**. What do you want to know?"

        elif any(p in lower for p in ('how are you',"what's up",'whats up')):
            bot_reply = "I'm doing great, always learning! 🤖 What do you want to know?"

        elif any(p in lower for p in ('who are you','what are you','what can you do')):
            bot_reply = (
                "I'm a **Living Neural Network** — an AI that only knows what it has actually read. 🧠\n\n"
                f"So far I've read **{stats['total']} articles** from the internet. "
                "I don't have any pre-built knowledge — everything I know came from real articles I've been fed. "
                "The more I train, the smarter I get. Ask me anything!"
            )

        elif lower in ('stats','status','knowledge'):
            lines = '\n'.join(f"• **{k}**: {v} articles"
                             for k,v in sorted(stats['categories'].items(), key=lambda x:-x[1]))
            srcs  = ', '.join(f"{k} ({v})" for k,v in list(stats['sources'].items())[:6])
            bot_reply = (f"**What I've learned:**\n\n{lines or '• nothing yet'}\n\n"
                         f"**Sources:** {srcs or 'none yet'}\n"
                         f"**Total:** {stats['total']} articles")

        elif lower in ('help','?'):
            bot_reply = ("Just ask me anything — I'll answer from what I've actually read!\n\n"
                         "Try: *What is science?* · *Who is Elon Musk?* · *What's happening in AI?*\n\n"
                         "Type `stats` to see what I know.")

        elif any(p in lower for p in ('thanks','thank you','thx','ty')):
            bot_reply = "You're welcome! 😊"

        else:
            results   = self.kb.search(msg, top_k=10)
            bot_reply = self._get_reply(msg, results, stats)

        history.append({"role": "user",      "content": user_message})
        history.append({"role": "assistant", "content": bot_reply})
        return history