from fastapi import FastAPI
from pydantic import BaseModel
import requests
import trafilatura
from ddgs import DDGS
import uvicorn
import re

from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

# --------------------------------------------------
# APP
# --------------------------------------------------

app = FastAPI(title="Llama AI Deep Research Backend")

# --------------------------------------------------
# EMBEDDING MODEL
# --------------------------------------------------

embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

# --------------------------------------------------
# REQUEST MODEL
# --------------------------------------------------

class ResearchRequest(BaseModel):
    query: str

# --------------------------------------------------
# SETTINGS
# --------------------------------------------------

MAX_RESULTS = 5
REQUEST_TIMEOUT = 10
CHUNK_SIZE = 700
CHUNK_OVERLAP = 120
TOP_RELEVANT_CHUNKS = 4

# FIX: Reduced from 2500 → 600 chars per source.
# Mobile context is 512 tokens (~2000 chars total).
# With 3–5 sources at 600 chars each we would already overflow.
# 600 chars per source gives enough signal without flooding the client.
MAX_SOURCE_CONTENT_CHARS = 600

# FIX: Reduced from 5000 → 900 chars for the global summary.
# 900 chars ≈ 225 tokens — fits comfortably in the mobile prompt budget
# while still conveying the key findings across all sources.
MAX_FINAL_SUMMARY_CHARS = 900

# --------------------------------------------------
# BLOCKED DOMAINS
# --------------------------------------------------

BLOCKED_DOMAINS = [
    "sourceforge.net",
    "topbusinesssoftware.com",
    "g2.com",
    "capterra.com",
    "softwareadvice.com",
    "slashdot.org",
]

# --------------------------------------------------
# JUNK SENTENCE FILTER
# FIX: The semantic ranker scores promotional sentences highly when they
# contain query words (e.g. "AI", "latest").  Filter them out before
# ranking so they never appear in the summary or compressed content.
# --------------------------------------------------

JUNK_PHRASES = [
    "newsletter", "subscribe", "sign up", "sign-up",
    "delivered to your inbox", "your inbox", "our inbox",
    "trending now", "stay ahead", "chatai", "what is chatai",
    "advertisement", "sponsored", "promoted",
    "click here", "read more", "learn more",
    "follow us", "more headlines", "by relevance", "by date",
    "top stories", "breaking news",
    "terms of service", "privacy policy", "cookie",
    "all rights reserved", "copyright",
    "please enable javascript", "javascript is required",
    "log in", "sign in", "register now",
]

def is_junk_sentence(sentence: str) -> bool:
    lower = sentence.lower()
    return any(phrase in lower for phrase in JUNK_PHRASES)

# --------------------------------------------------
# CLEAN TEXT
# --------------------------------------------------

def clean_text(text):
    if not text:
        return ""
    text = re.sub(r'\s+', ' ', text)
    return text.strip()

# --------------------------------------------------
# SENTENCE SPLITTER
# Now filters junk sentences at the source.
# --------------------------------------------------

def split_sentences(text):
    sentences = re.split(r'(?<=[.!?])\s+', text)
    return [
        s.strip()
        for s in sentences
        if len(s.strip()) > 40 and not is_junk_sentence(s)
    ]

# --------------------------------------------------
# DOMAIN FILTER
# --------------------------------------------------

def is_blocked_domain(url):
    for domain in BLOCKED_DOMAINS:
        if domain in url:
            return True
    return False

# --------------------------------------------------
# CHUNKING
# --------------------------------------------------

def chunk_text(text, chunk_size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
    words = text.split()
    chunks = []
    start = 0
    while start < len(words):
        end = start + chunk_size
        chunk = words[start:end]
        chunks.append(" ".join(chunk))
        start += chunk_size - overlap
    return chunks

# --------------------------------------------------
# EXTRACT CONTENT
# --------------------------------------------------

def extract_content(url):
    try:
        response = requests.get(
            url,
            timeout=REQUEST_TIMEOUT,
            headers={"User-Agent": "Mozilla/5.0"},
        )
        if response.status_code != 200:
            return ""
        extracted = trafilatura.extract(
            response.text,
            include_links=False,
            include_images=False,
            favor_precision=True,
        )
        if not extracted:
            return ""
        return clean_text(extracted)
    except Exception:
        return ""

# --------------------------------------------------
# SEMANTIC RANKING
# --------------------------------------------------

def rank_chunks(query, chunks):
    if not chunks:
        return []
    try:
        query_embedding = embedding_model.encode([query])
        chunk_embeddings = embedding_model.encode(chunks)
        similarities = cosine_similarity(query_embedding, chunk_embeddings)[0]
        ranked = sorted(
            zip(chunks, similarities),
            key=lambda x: x[1],
            reverse=True,
        )
        return [chunk for chunk, score in ranked[:TOP_RELEVANT_CHUNKS]]
    except Exception:
        return chunks[:TOP_RELEVANT_CHUNKS]

# --------------------------------------------------
# SOURCE COMPRESSION
# --------------------------------------------------

def compress_chunks(query, chunks):
    try:
        all_sentences = []
        for chunk in chunks:
            sentences = split_sentences(chunk)   # junk already filtered here
            all_sentences.extend(sentences)
        if not all_sentences:
            return ""
        query_embedding = embedding_model.encode([query])
        sentence_embeddings = embedding_model.encode(all_sentences)
        similarities = cosine_similarity(query_embedding, sentence_embeddings)[0]
        ranked = sorted(
            zip(all_sentences, similarities),
            key=lambda x: x[1],
            reverse=True,
        )
        final_sentences = []
        used = set()
        total_chars = 0
        for sentence, score in ranked:
            normalized = sentence.lower().strip()
            if normalized in used:
                continue
            used.add(normalized)
            if total_chars + len(sentence) > MAX_SOURCE_CONTENT_CHARS:
                break
            final_sentences.append(sentence)
            total_chars += len(sentence)
        return " ".join(final_sentences)   # FIX: single paragraph, not \n\n — saves chars
    except Exception:
        return ""

# --------------------------------------------------
# GLOBAL SYNTHESIS
# --------------------------------------------------

def synthesize_research(query, results):
    try:
        all_sentences = []
        for result in results:
            content = result.get("content", "")
            if not content:
                continue
            sentences = split_sentences(content)   # junk already filtered
            all_sentences.extend(sentences)
        if not all_sentences:
            return ""
        query_embedding = embedding_model.encode([query])
        sentence_embeddings = embedding_model.encode(all_sentences)
        similarities = cosine_similarity(query_embedding, sentence_embeddings)[0]
        ranked = sorted(
            zip(all_sentences, similarities),
            key=lambda x: x[1],
            reverse=True,
        )
        final_sentences = []
        used = set()
        total_chars = 0
        for sentence, score in ranked:
            normalized = sentence.lower().strip()
            if normalized in used:
                continue
            used.add(normalized)
            if total_chars + len(sentence) > MAX_FINAL_SUMMARY_CHARS:
                break
            final_sentences.append(sentence)
            total_chars += len(sentence)
        return " ".join(final_sentences)   # FIX: single paragraph, not \n\n
    except Exception:
        return ""

# --------------------------------------------------
# RESEARCH PIPELINE
# --------------------------------------------------

def perform_research(query):
    source_results = []
    try:
        with DDGS() as ddgs:
            search_results = list(ddgs.text(query, max_results=MAX_RESULTS))
        if not search_results:
            return None
        for item in search_results:
            try:
                title = item.get("title", "")
                url = item.get("href", "")
                snippet = item.get("body", "")
                if not url:
                    continue
                if is_blocked_domain(url):
                    continue
                content = extract_content(url)
                if not content:
                    content = snippet
                if not content:
                    continue
                chunks = chunk_text(content)
                ranked_chunks = rank_chunks(query, chunks)
                compressed_content = compress_chunks(query, ranked_chunks)
                if not compressed_content:
                    continue
                source_results.append({
                    "title": title,
                    "url": url,
                    "snippet": snippet,
                    "content": compressed_content,   # semantically ranked extraction
                })
            except Exception:
                continue
        if not source_results:
            return None
        research_summary = synthesize_research(query, source_results)
        return {
            "sources": source_results,
            "research_summary": research_summary,
        }
    except Exception:
        return None

# --------------------------------------------------
# ROOT
# --------------------------------------------------

@app.get("/")
async def root():
    return {"status": "running"}

# --------------------------------------------------
# DEEP RESEARCH ENDPOINT
# --------------------------------------------------

@app.post("/deep-research")
async def deep_research(req: ResearchRequest):
    query = req.query.strip()
    if not query:
        return {
            "success": False,
            "message": "Query is empty.",
            "results": [],
        }
    research_data = perform_research(query)
    if not research_data:
        return {
            "success": False,
            "message": "No useful research results found.",
            "results": [],
        }
    return {
        "success": True,
        "query": query,
        "research_summary": research_data["research_summary"],
        "results": research_data["sources"],
    }

# --------------------------------------------------
# START
# --------------------------------------------------

if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860)