"""
Massa QA - Question Answering sur la documentation Massa blockchain
Utilise l'API Anthropic (Claude Sonnet) avec approche hybride :
- Context stuffing pour les questions simples (docs pré-chargées)
- RAG avec FAISS pour les corpus plus larges
"""

import os
import gradio as gr
import anthropic
import requests
from bs4 import BeautifulSoup
import faiss
import numpy as np
import pickle
import hashlib
import time

# --- Configuration ---
ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY")
MODEL = "claude-sonnet-4-20250514"
MAX_CONTEXT_TOKENS = 8000  # seuil pour basculer en RAG
EMBEDDING_DIM = 384  # dimension pour les embeddings légers

# --- Pages de documentation Massa à scraper ---
MASSA_DOC_URLS = [
    "https://docs.massa.net/docs/learn/introduction",
    "https://docs.massa.net/docs/learn/architecture/basic-concepts",
    "https://docs.massa.net/docs/learn/decentralized-web",
    "https://docs.massa.net/docs/build/smart-contract/intro",
    "https://docs.massa.net/docs/build/standards",
    "https://docs.massa.net/docs/build/networks-faucets/public-networks",
    "https://docs.massa.net/docs/build/massa-web3/intro",
]

# --- Scraping de la documentation ---
def fetch_page(url: str) -> dict:
    """Récupère le contenu texte d'une page de documentation."""
    try:
        resp = requests.get(url, timeout=15)
        resp.raise_for_status()
        soup = BeautifulSoup(resp.text, "html.parser")
        # Docusaurus : le contenu principal est dans <article> ou <main>
        article = soup.find("article") or soup.find("main") or soup.find("body")
        if article:
            # Nettoyer les scripts et styles
            for tag in article.find_all(["script", "style", "nav", "footer"]):
                tag.decompose()
            text = article.get_text(separator="\n", strip=True)
        else:
            text = soup.get_text(separator="\n", strip=True)
        return {"url": url, "text": text[:6000]}  # limiter par page
    except Exception as e:
        return {"url": url, "text": f"[Erreur de chargement: {e}]"}


def load_corpus() -> list[dict]:
    """Charge toutes les pages de documentation."""
    corpus = []
    for url in MASSA_DOC_URLS:
        page = fetch_page(url)
        if page["text"] and not page["text"].startswith("[Erreur"):
            corpus.append(page)
    return corpus


# --- Embeddings légers avec TF-IDF simplifié ---
def simple_tokenize(text: str) -> list[str]:
    """Tokenisation basique."""
    import re
    return re.findall(r'\b\w+\b', text.lower())


def build_vocab(corpus_texts: list[str], max_features: int = EMBEDDING_DIM) -> dict:
    """Construit un vocabulaire à partir du corpus."""
    from collections import Counter
    all_tokens = Counter()
    for text in corpus_texts:
        all_tokens.update(set(simple_tokenize(text)))
    # Garder les plus fréquents
    most_common = all_tokens.most_common(max_features)
    return {word: idx for idx, (word, _) in enumerate(most_common)}


def text_to_vector(text: str, vocab: dict) -> np.ndarray:
    """Convertit un texte en vecteur TF simplifié."""
    vec = np.zeros(len(vocab), dtype=np.float32)
    tokens = simple_tokenize(text)
    for token in tokens:
        if token in vocab:
            vec[vocab[token]] += 1
    # Normaliser
    norm = np.linalg.norm(vec)
    if norm > 0:
        vec /= norm
    return vec


# --- Classe principale ---
class MassaQA:
    def __init__(self):
        self.client = None
        self.corpus = []
        self.chunks = []
        self.vocab = {}
        self.index = None
        self.use_rag = False
        self.loaded = False

    def initialize(self):
        """Charge le corpus et prépare le système."""
        if self.loaded:
            return

        # Client Anthropic
        if ANTHROPIC_API_KEY:
            self.client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)

        # Charger la doc
        self.corpus = load_corpus()

        # Découper en chunks
        self.chunks = []
        for doc in self.corpus:
            text = doc["text"]
            url = doc["url"]
            # Chunks de ~800 caractères avec overlap
            chunk_size = 800
            overlap = 100
            for i in range(0, len(text), chunk_size - overlap):
                chunk = text[i:i + chunk_size]
                if len(chunk.strip()) > 50:
                    self.chunks.append({"text": chunk, "url": url})

        # Décider de l'approche
        total_chars = sum(len(c["text"]) for c in self.chunks)
        total_tokens_approx = total_chars // 4

        if total_tokens_approx > MAX_CONTEXT_TOKENS:
            # RAG mode : construire l'index FAISS
            self.use_rag = True
            corpus_texts = [c["text"] for c in self.chunks]
            self.vocab = build_vocab(corpus_texts)
            vectors = np.array([text_to_vector(t, self.vocab) for t in corpus_texts])

            dim = vectors.shape[1]
            self.index = faiss.IndexFlatIP(dim)  # Inner Product (cosine sur vecteurs normalisés)
            self.index.add(vectors)
        else:
            self.use_rag = False

        self.loaded = True

    def get_context(self, question: str, top_k: int = 5) -> tuple[str, list[dict]]:
        """Récupère le contexte pertinent pour une question.
        Retourne (context_text, sources) où sources est une liste de {url, snippet}."""
        sources = []
        if not self.use_rag:
            # Context stuffing : tout envoyer
            context_parts = []
            for doc in self.corpus:
                context_parts.append(f"--- Source: {doc['url']} ---\n{doc['text']}")
                sources.append({
                    "url": doc["url"],
                    "snippet": doc["text"][:150].replace("\n", " ").strip()
                })
            return "\n\n".join(context_parts), sources
        else:
            # RAG : recherche par similarité
            q_vec = text_to_vector(question, self.vocab).reshape(1, -1)
            scores, indices = self.index.search(q_vec, top_k)
            context_parts = []
            seen_urls = set()
            for i, idx in enumerate(indices[0]):
                if idx < len(self.chunks):
                    chunk = self.chunks[idx]
                    score = float(scores[0][i])
                    source_tag = f"[Source: {chunk['url']}]" if chunk['url'] not in seen_urls else ""
                    seen_urls.add(chunk['url'])
                    context_parts.append(f"{source_tag}\n{chunk['text']}")
                    sources.append({
                        "url": chunk["url"],
                        "snippet": chunk["text"][:150].replace("\n", " ").strip(),
                        "score": score
                    })
            return "\n\n".join(context_parts), sources

    def ask(self, question: str, history: list) -> str:
        """Répond à une question sur Massa."""
        self.initialize()

        if not self.client:
            return "⚠️ Clé API Anthropic non configurée. Ajoutez ANTHROPIC_API_KEY dans les secrets du Space."

        if not question.strip():
            return "Posez une question sur la blockchain Massa !"

        context, sources = self.get_context(question)
        mode = "RAG" if self.use_rag else "Context"

        system_prompt = f"""Tu es un assistant expert sur la blockchain Massa. 
Tu réponds aux questions en te basant UNIQUEMENT sur la documentation fournie ci-dessous.
Si la réponse n'est pas dans la documentation, dis-le clairement.
Réponds en français sauf si la question est en anglais.
Cite les sources quand c'est pertinent.
Sois concis et précis.

Mode de recherche : {mode} ({len(self.chunks)} chunks indexés)

--- DOCUMENTATION MASSA ---
{context}
--- FIN DE LA DOCUMENTATION ---"""

        try:
            message = self.client.messages.create(
                model=MODEL,
                max_tokens=1024,
                system=system_prompt,
                messages=[{"role": "user", "content": question}]
            )
            answer = message.content[0].text

            # Compteur de tokens depuis l'API
            usage = message.usage
            input_tokens = usage.input_tokens
            output_tokens = usage.output_tokens
            total_tokens = input_tokens + output_tokens

            # Ajouter les top 3 sources utilisées
            top_sources = sources[:3]
            if top_sources:
                answer += "\n\n---\n📚 **Sources utilisées :**\n"
                for i, src in enumerate(top_sources, 1):
                    score_str = f" (score: {src['score']:.2f})" if "score" in src else ""
                    answer += f"\n{i}. [{src['url'].split('/')[-1] or src['url'].split('/')[-2]}]({src['url']}){score_str}\n"
                    answer += f"   > *{src['snippet']}...*\n"

            # Ajouter le compteur de tokens
            answer += f"\n---\n🔢 **Tokens** : {input_tokens} in + {output_tokens} out = **{total_tokens}** total | Mode: {mode} | Modèle: {MODEL}\n"

            return answer
        except anthropic.APIError as e:
            return f"⚠️ Erreur API Anthropic : {e}"
        except Exception as e:
            return f"⚠️ Erreur inattendue : {e}"


# --- Instance globale ---
qa = MassaQA()


# --- Interface Gradio ---
def respond(message, history):
    response = qa.ask(message, history)
    return response


DESCRIPTION = """
# 🔗 Massa QA — Assistant Documentation Massa

Posez vos questions sur la **blockchain Massa** : architecture, smart contracts, DeWeb, tokens, réseau...

Propulsé par **Claude Sonnet** (Anthropic) avec recherche sur la [documentation officielle](https://docs.massa.net).

> 💡 *Approche hybride : context stuffing si le corpus tient dans le prompt, RAG/FAISS sinon.*
"""

demo = gr.ChatInterface(
    fn=respond,
    title="Massa QA",
    description=DESCRIPTION,
    examples=[
        "Qu'est-ce que Massa ?",
        "Comment fonctionnent les smart contracts sur Massa ?",
        "Qu'est-ce que le DeWeb ?",
        "Quels sont les réseaux disponibles (mainnet, buildnet) ?",
        "Comment créer un token sur Massa ?",
    ],
    cache_examples=False,
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="indigo",
    ),
)

if __name__ == "__main__":
    demo.launch()