"""
Summary generation utilities for research papers.
"""
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from src.config import GROQ_API_KEY, LLM_MODEL, LLM_TEMPERATURE


def ensure_groq_key():
    if not GROQ_API_KEY:
        raise RuntimeError("GROQ_API_KEY is not configured. Set it in .env or Hugging Face secrets.")


def generate_paper_summary(text: str, max_length: int = 500) -> str:
    """Generate a concise summary of the paper."""
    ensure_groq_key()
    llm = ChatGroq(
        model=LLM_MODEL,
        api_key=GROQ_API_KEY,
        temperature=LLM_TEMPERATURE
    )
    
    prompt = ChatPromptTemplate.from_template("""You are a research paper analyst. 
Generate a clear, concise summary of the following research paper in {max_length} words or less.
Focus on: problem statement, methodology, key findings, and implications.

Paper content:
{text}

Summary:""")
    
    chain = prompt | llm | StrOutputParser()
    summary = chain.invoke({"text": text[:3000], "max_length": max_length})
    return summary


def generate_section_summaries(chunks: list[str], num_sections: int = 5) -> list[str]:
    """Generate summaries for different sections of the paper."""
    ensure_groq_key()
    if len(chunks) < num_sections:
        num_sections = len(chunks)
    
    section_size = len(chunks) // num_sections
    summaries = []
    
    llm = ChatGroq(
        model=LLM_MODEL,
        api_key=GROQ_API_KEY,
        temperature=LLM_TEMPERATURE
    )
    
    prompt = ChatPromptTemplate.from_template("""Summarize this section of a research paper in 1-2 sentences:
{section_text}""")
    
    chain = prompt | llm | StrOutputParser()
    
    for i in range(num_sections):
        start_idx = i * section_size
        end_idx = (i + 1) * section_size if i < num_sections - 1 else len(chunks)
        section_text = "\n".join(chunks[start_idx:end_idx])
        summary = chain.invoke({"section_text": section_text[:2000]})
        summaries.append(f"Section {i+1}: {summary}")
    
    return summaries