"""
Trend Longevity Analyser
========================
A RAG-powered tool that classifies whether a topic/trend is Early, Rising,
at Peak, Declining, or Fading — based on live news data.

RAG Pipeline (6 steps):
  1. FETCH    → NewsAPI pulls up to 50 recent articles on the topic
  2. CHUNK    → LangChain splits each article into overlapping 400-char chunks
  3. EMBED    → Sentence Transformers encodes every chunk as a dense vector
  4. STORE    → ChromaDB holds all vectors in an in-memory collection
  5. RETRIEVE → Semantic search returns the 6 most relevant chunks
  6. GENERATE → Gemini Flash analyses the retrieved context and returns a structured signal

Built by: Sammie Wong
"""

import os
import json
import uuid
import time
import gradio as gr
import google.generativeai as genai
import chromadb
from chromadb.utils import embedding_functions
from newsapi import NewsApiClient
from langchain_text_splitters import RecursiveCharacterTextSplitter
from datetime import datetime, timedelta


# ─── STEP 1: FETCH ────────────────────────────────────────────────────────────

def fetch_articles(topic: str, days_back: int, newsapi_key: str,
                   boolean_query: str = "") -> list[dict]:
    """
    Pull recent news articles from NewsAPI.

    Two modes:
      Boolean mode — when boolean_query is provided, passed directly to NewsAPI
                     q parameter. Supports AND, OR, NOT, "exact phrase", (grouping).
                     No automatic quoting or fallback — you control the query precisely.
      Simple mode  — two-pass automatic:
                     Pass 1: exact phrase match e.g. "green hydrogen"
                     Pass 2: broad keyword fallback if < 5 results returned

    Free tier: 100 requests/day. Boolean = 1 request; simple = 1–2 requests.
    """
    client = NewsApiClient(api_key=newsapi_key)
    from_date = (datetime.now() - timedelta(days=days_back)).strftime('%Y-%m-%d')

    common_params = dict(
        language='en',
        sort_by='publishedAt',
        from_param=from_date,
        page_size=50,
    )

    if boolean_query.strip():
        # Boolean mode: pass query directly, no modification
        response = client.get_everything(q=boolean_query.strip(), **common_params)
        return response.get('articles', [])

    # Simple mode: two-pass
    response = client.get_everything(q=f'"{topic}"', **common_params)
    articles = response.get('articles', [])
    if len(articles) < 5:
        response = client.get_everything(q=topic, **common_params)
        articles = response.get('articles', [])
    return articles


# ─── STEP 2: CHUNK ────────────────────────────────────────────────────────────

def chunk_articles(articles: list[dict]) -> tuple[list[str], list[dict]]:
    """
    Split each article into overlapping chunks.

    Why chunk? LLMs have context limits and RAG works better with focused
    pieces of text than entire articles. Overlap (40 chars) prevents
    losing meaning at chunk boundaries.
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=400,
        chunk_overlap=40,
    )

    all_chunks, all_metadata = [], []

    for article in articles:
        full_text = " ".join(filter(None, [
            article.get('title', ''),
            article.get('description', ''),
            article.get('content', ''),
        ]))

        for chunk in splitter.split_text(full_text):
            all_chunks.append(chunk)
            all_metadata.append({
                'source':       article['source']['name'],
                'published_at': article['publishedAt'][:10],
                'title':        (article.get('title') or '')[:100],
            })

    return all_chunks, all_metadata


# ─── STEPS 3 & 4: EMBED + STORE ───────────────────────────────────────────────

def build_vector_store(chunks: list[str], metadatas: list[dict]):
    """
    Encode every chunk into a 384-dimensional embedding vector and store
    in an in-memory ChromaDB collection.

    Model: all-MiniLM-L6-v2 — fast, accurate, no GPU needed.
    EphemeralClient = fully in-memory, stateless per request.
    """
    embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name="all-MiniLM-L6-v2"
    )

    try:
        db = chromadb.EphemeralClient()
    except AttributeError:
        db = chromadb.Client()

    collection = db.create_collection(
        name=f"trend_{uuid.uuid4().hex[:8]}",
        embedding_function=embedding_fn,
    )

    batch_size = 100
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i:i + batch_size]
        collection.add(
            documents=batch,
            metadatas=metadatas[i:i + batch_size],
            ids=[f"c{i + j}" for j in range(len(batch))],
        )

    return collection


# ─── STEP 5: RETRIEVE ─────────────────────────────────────────────────────────

def retrieve_context(collection, topic: str, n_results: int = 6) -> tuple[list[str], list[dict]]:
    """
    Retrieve the most relevant chunks via cosine similarity.

    n_results=6 (down from 10) keeps the Gemini prompt comfortably
    under the free-tier token-per-minute limit while preserving signal quality.
    """
    query = f"trend momentum growth decline viral popularity media coverage {topic}"

    results = collection.query(
        query_texts=[query],
        n_results=min(n_results, collection.count()),
    )
    return results['documents'][0], results['metadatas'][0]


# ─── STEP 6: GENERATE ─────────────────────────────────────────────────────────

# Hard cap on context chars sent to Gemini.
# 6 chunks × 400 chars = ~2400 chars raw. We cap at 3000 to stay well
# under the free-tier 32K tokens-per-minute limit.
MAX_CONTEXT_CHARS = 3000

def analyse_with_gemini(
    topic: str,
    retrieved_docs: list[str],
    retrieved_meta: list[dict],
    total_articles: int,
    days_back: int,
    gemini_key: str,
) -> dict:
    """
    Pass retrieved chunks to Gemini Flash for trend analysis.

    Reliability features:
    - Context is hard-capped at MAX_CONTEXT_CHARS to prevent 429 token-rate errors.
    - Auto-retry: on a 429 response, waits 15 s then retries up to 2 more times.
    - response_mime_type="application/json" forces clean JSON — no parsing issues.
    """
    # Build context, truncating to stay inside token-rate limit
    context_parts = []
    total_chars = 0
    for doc, m in zip(retrieved_docs, retrieved_meta):
        entry = f"[{m['source']} · {m['published_at']}]\n{doc}"
        if total_chars + len(entry) > MAX_CONTEXT_CHARS:
            # Include a truncated version rather than skipping entirely
            remaining = MAX_CONTEXT_CHARS - total_chars
            if remaining > 80:
                context_parts.append(entry[:remaining] + "…")
            break
        context_parts.append(entry)
        total_chars += len(entry)

    context = "\n\n".join(context_parts)

    prompt = f"""You are a senior trend analyst at a social listening company (like Pulsar, Brandwatch, or Sprinklr).

You have been given the {len(context_parts)} most semantically relevant news excerpts about "{topic}",
retrieved via RAG from a corpus of {total_articles} articles published in the last {days_back} days.

RETRIEVED CONTEXT:
{context}

Based on this evidence, classify the trend and return a JSON object with exactly these fields:
{{
    "phase": "Early|Rising|Peak|Declining|Fading",
    "confidence": <integer 0-100>,
    "momentum_score": <integer 0-100>,
    "summary": "<2-3 sentence narrative of the trend's current state and why>",
    "key_signals": [
        "<specific evidence signal from the articles>",
        "<specific evidence signal from the articles>",
        "<specific evidence signal from the articles>",
        "<specific evidence signal from the articles>"
    ],
    "prediction_30d": "<one concrete sentence: what will likely happen to this trend in 30 days>",
    "media_spread": "Niche|Specialist|Mainstream|Viral",
    "top_themes": ["<theme 1>", "<theme 2>", "<theme 3>"]
}}

Phase definitions:
- Early:     Emerging signal, limited coverage, mostly specialist/niche sources
- Rising:    Volume growing, broadening coverage, sentiment intensifying
- Peak:      Maximum velocity, mainstream saturation, brand/celebrity involvement
- Declining: Volume dropping, novelty fading, sentiment normalising
- Fading:    Minimal coverage, topic becoming dated, resolved, or replaced"""

    genai.configure(api_key=gemini_key)
    model = genai.GenerativeModel(
        model_name="gemini-2.0-flash",
        generation_config=genai.GenerationConfig(
            response_mime_type="application/json",
            max_output_tokens=600,
        ),
    )

    # Retry loop: up to 3 attempts, 15-second pause on 429
    last_error = None
    for attempt in range(3):
        try:
            response = model.generate_content(prompt)
            return json.loads(response.text)
        except Exception as e:
            last_error = e
            if "429" in str(e) and attempt < 2:
                wait = 15 * (attempt + 1)   # 15 s, then 30 s
                time.sleep(wait)
                continue
            break

    raise last_error


# ─── PIPELINE ORCHESTRATOR ────────────────────────────────────────────────────

def run_analysis(topic: str, days_back: int, newsapi_key_input: str, gemini_key_input: str,
                 boolean_query: str = ""):
    """
    Orchestrates the full 6-step RAG pipeline.
    Returns 6 Markdown strings for the Gradio output components.
    """
    EMPTY = ("", "", "", "", "", "")

    if not topic.strip():
        return ("⚠️ Please enter a topic label (used for the Gemini prompt even in Boolean mode).", *EMPTY[1:])

    news_key   = newsapi_key_input.strip()  or os.environ.get("NEWSAPI_KEY", "")
    gemini_key = gemini_key_input.strip()   or os.environ.get("GEMINI_API_KEY", "")

    if not news_key:
        return ("⚠️ NewsAPI key required. Free key at newsapi.org", *EMPTY[1:])
    if not gemini_key:
        return ("⚠️ Gemini key required. Free key at aistudio.google.com", *EMPTY[1:])

    using_boolean = bool(boolean_query.strip())
    search_label  = boolean_query.strip() if using_boolean else topic

    try:
        # Step 1
        articles = fetch_articles(topic, int(days_back), news_key, boolean_query)
        if not articles:
            hint = (
                f"Boolean query returned no results: `{boolean_query}`\n\nCheck your operators and try broader terms."
                if using_boolean else
                f"No articles found for **'{topic}'** in the last {days_back} days.\n\nTry a shorter or broader topic name."
            )
            return (f"⚠️ {hint}", *EMPTY[1:])

        # Steps 2–4
        chunks, metadatas = chunk_articles(articles)
        collection = build_vector_store(chunks, metadatas)

        # Step 5
        retrieved_docs, retrieved_meta = retrieve_context(collection, topic)

        # Step 6
        result = analyse_with_gemini(
            topic, retrieved_docs, retrieved_meta,
            len(articles), int(days_back), gemini_key,
        )

        # ── Format outputs ────────────────────────────────────────────────
        phase_icons  = {"Early": "🌱", "Rising": "📈", "Peak": "🔥", "Declining": "📉", "Fading": "🌫️"}
        phase        = result.get("phase", "Unknown")
        icon         = phase_icons.get(phase, "❓")
        phase_md     = f"## {icon} {phase.upper()}"

        mode_label = f"Boolean: `{boolean_query.strip()}`" if using_boolean else "Simple (auto)"
        scores_md = (
            f"### 📊 Metrics\n| | |\n|---|---|\n"
            f"| **Confidence** | `{result.get('confidence', '—')}%` |\n"
            f"| **Momentum Score** | `{result.get('momentum_score', '—')} / 100` |\n"
            f"| **Media Spread** | `{result.get('media_spread', '—')}` |\n"
            f"| **Search Mode** | {mode_label} |\n"
            f"| **Articles Fetched** | `{len(articles)}` |\n"
            f"| **Chunks in Vector Store** | `{len(chunks)}` |\n"
            f"| **Chunks Sent to Gemini** | `{len(retrieved_docs)}` |"
        )

        summary_md    = f"### 📝 Narrative\n{result.get('summary', '')}"
        signals_md    = (
            "### 🔍 Key Signals\n"
            + "\n".join(f"- {s}" for s in result.get('key_signals', []))
            + "\n\n**Top Themes:** "
            + "  ".join(f"`{t}`" for t in result.get('top_themes', []))
        )
        prediction_md = f"### 🔮 30-Day Prediction\n_{result.get('prediction_30d', '')}_"
        sources_md    = "### 📰 Sources Retrieved via RAG\n" + "\n".join(
            f"- **{m['source']}** ({m['published_at']}): _{m['title']}_"
            for m in retrieved_meta[:6]
        )

        return phase_md, scores_md, summary_md, signals_md, prediction_md, sources_md

    except json.JSONDecodeError:
        return ("⚠️ Could not parse Gemini response. Please try again.", *EMPTY[1:])
    except Exception as e:
        err = str(e)
        if "429" in err:
            return (
                "⚠️ **Rate limit hit (429).** Gemini free tier allows ~15 requests/minute.\n\n"
                "Wait 60 seconds and try again — the retry logic above handled 2 attempts already.",
                *EMPTY[1:]
            )
        return (f"⚠️ Error: {err}", *EMPTY[1:])


# ─── GRADIO UI ────────────────────────────────────────────────────────────────

CSS = """
#title-block { text-align: center; padding: 8px 0 4px 0; }
#title-block h1 { font-size: 2rem; font-weight: 800; margin-bottom: 2px; }
#title-block p  { color: #64748b; font-size: 0.95rem; }
.phase-box      { text-align: center; font-size: 1.8rem; font-weight: 800;
                  padding: 18px; background: #f8fafc; border-radius: 12px;
                  border: 2px solid #e2e8f0; min-height: 80px; }
.metric-card    { background: #f8fafc; border-radius: 10px; padding: 4px 8px;
                  border-left: 3px solid #3b82f6; }
footer { display: none !important; }
"""

HOWTO = """
**Step 1 — Fetch:** NewsAPI retrieves up to 50 articles. Two modes: *Simple* (auto exact-phrase then broad fallback) or *Boolean* — where your query is passed directly to NewsAPI supporting `AND`, `OR`, `NOT`, `"exact phrase"`, and `(grouping)`. Boolean mode gives you precise control over what enters the RAG corpus.

**Step 2 — Chunk:** LangChain's `RecursiveCharacterTextSplitter` breaks each article into overlapping 400-character chunks. Overlap prevents information loss at chunk boundaries.

**Step 3 — Embed:** Each chunk is encoded into a 384-dimensional dense vector using `all-MiniLM-L6-v2` (Sentence Transformers) — no GPU needed.

**Step 4 — Store:** All vectors are loaded into a ChromaDB `EphemeralClient` collection (fully in-memory, stateless per request).

**Step 5 — Retrieve:** A trend-signal query runs against the vector store. ChromaDB returns the 6 most semantically relevant chunks via cosine similarity.

**Step 6 — Generate:** Gemini Flash receives only the retrieved chunks (capped at 3,000 chars) — not all 50 articles. This keeps the prompt inside the free-tier token-rate limit and is the core RAG advantage. Auto-retries on 429 with a 15-second wait.
"""

with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="slate",
        font=[gr.themes.GoogleFont("DM Sans"), "ui-sans-serif", "sans-serif"],
    ),
    css=CSS,
    title="Trend Longevity Analyser",
) as demo:

    with gr.Column(elem_id="title-block"):
        gr.Markdown("""
# 📊 Trend Longevity Analyser
**RAG-powered trend intelligence** · NewsAPI + ChromaDB + Sentence Transformers + Gemini Flash
""")

    gr.Markdown("---")

    with gr.Accordion("🔑 API Keys", open=False):
        gr.Markdown(
            "_Keys are used only for your request and never stored. "
            "Set `NEWSAPI_KEY` and `GEMINI_API_KEY` as Space Secrets to avoid entering them each time._"
        )
        with gr.Row():
            newsapi_input = gr.Textbox(label="NewsAPI Key",          placeholder="newsapi.org — free",         type="password")
            gemini_input  = gr.Textbox(label="Google Gemini API Key", placeholder="aistudio.google.com — free", type="password")

    with gr.Row():
        with gr.Column(scale=4):
            topic_input = gr.Textbox(
                label="Topic Label",
                placeholder="e.g.  Gen Z workplace burnout  ·  green hydrogen  ·  AI companions",
                info="Used as the display name and passed to Gemini. Always required.",
                lines=1,
            )
        with gr.Column(scale=1):
            days_input = gr.Slider(
                label="Days Back", minimum=7, maximum=30, value=30, step=7,
                info="Free NewsAPI tier = 30 days max",
            )

    boolean_input = gr.Textbox(
        label="⚡ Boolean Search Query (optional — overrides simple search)",
        placeholder='e.g.  "Gen Z" AND (burnout OR "quiet quitting") NOT productivity',
        info='Operators: AND  OR  NOT  "exact phrase"  (grouping)  — passed directly to NewsAPI',
        lines=1,
    )

    analyse_btn = gr.Button("🔍  Analyse Trend", variant="primary", size="lg")

    gr.Markdown("---")

    phase_output = gr.Markdown(elem_classes=["phase-box"])

    with gr.Row():
        with gr.Column():
            scores_output  = gr.Markdown(elem_classes=["metric-card"])
            signals_output = gr.Markdown()
        with gr.Column():
            summary_output    = gr.Markdown()
            prediction_output = gr.Markdown()

    sources_output = gr.Markdown()

    with gr.Accordion("🧠  How the RAG Pipeline Works", open=False):
        gr.Markdown(HOWTO)

    gr.Markdown(
        "_Built by [Sammie Wong](https://linkedin.com/in/) · "
        "Powered by [NewsAPI](https://newsapi.org) + [ChromaDB](https://www.trychroma.com) + "
        "[Gemini Flash](https://aistudio.google.com)_"
    )

    analyse_btn.click(
        fn=run_analysis,
        inputs=[topic_input, days_input, newsapi_input, gemini_input, boolean_input],
        outputs=[phase_output, scores_output, summary_output,
                 signals_output, prediction_output, sources_output],
    )

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