import streamlit as st
from groq import Groq
import json
from src.embeddings import embedding_model
from src.vector_store import VectorStore


# =====================================================
# GROQ CLIENT
# =====================================================

client = Groq(
    api_key=st.secrets["GROQ_API_KEY"]
)

vector_db = VectorStore()


# =====================================================
# CHUNKING
# =====================================================

def chunk_text(text, chunk_size=1000):

    if not text:
        return []

    return [
        text[i:i + chunk_size]
        for i in range(0, len(text), chunk_size)
    ]


# =====================================================
# VECTOR DATABASE
# =====================================================

def create_rag_database(papers):

    documents = []

    for paper in papers:

        chunks = chunk_text(
            paper.get("text", "")
        )

        documents.extend(chunks)

    if not documents:
        return

    embeddings = embedding_model.encode(
        documents,
        show_progress_bar=False
    )

    vector_db.build(
        embeddings,
        documents
    )


# =====================================================
# CHAT WITH PAPERS
# =====================================================

def ask_rag(question):

    query_embedding = embedding_model.encode(question)

    context = vector_db.search(query_embedding)

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

    prompt = f"""
You are an expert scientific research assistant.

Answer ONLY from the provided research paper.

If the answer is not available,
reply:

"I could not find this information in the uploaded paper."

Context:

{context_text}

Question:

{question}
"""

    response = client.chat.completions.create(

        model="llama-3.1-8b-instant",

        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ],

        temperature=0.2
    )

    return response.choices[0].message.content

# =====================================================
# COMPLETE PAPER ANALYSIS (Single AI Call)
# =====================================================

def analyze_paper(text, abstract=""):

    if not text:
        return {
            "summary": "No text available.",
            "abstract_summary": "No abstract available.",
            "limitations": "Not available.",
            "research_gaps": "Not available."
        }

    # Use only a limited amount of text to stay within token limits
    paper_text = text[:4000]

    if abstract:
        abstract = abstract[:1500]

    prompt = f"""
You are an expert scientific research assistant.

Analyze the following research paper.

Return ONLY valid JSON.

The JSON must have EXACTLY these keys:

{{
  "summary": "...",
  "abstract_summary": "...",
  "limitations": "...",
  "research_gaps": "..."
}}

Instructions:

- abstract_summary:
Summarize ONLY the abstract in 4-5 sentences.

- summary:
A concise summary (150-200 words).

- limitations:
List the main 2-3 limitations as bullet points.

- research_gaps:
List 3 future research directions.

ABSTRACT:

{abstract}

PAPER:

{paper_text}
"""

    response = client.chat.completions.create(

        model="llama-3.1-8b-instant",

        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ],

        temperature=0.2,

        response_format={
            "type": "json_object"
        }

    )

    try:

        result = json.loads(
            response.choices[0].message.content
        )

    except Exception:

        result = {


            "abstract_summary":
                "Abstract summary unavailable.",
            
            "summary":
                "Summary generation failed.",

            "limitations":
                "Limitations unavailable.",

            "research_gaps":
                "Research gaps unavailable."

        }

    return result