import os
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent
import gradio as gr
import time

load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

llm = ChatGroq(
    model="llama-3.1-8b-instant",
    api_key=GROQ_API_KEY,
    temperature=0.3
)

search = DuckDuckGoSearchRun()

@tool
def web_search(query: str) -> str:
    """Search the web for current information about a topic."""
    max_retries = 3
    for attempt in range(max_retries):
        try:
            results = search.run(query)
            if results:
                return results
        except Exception as e:
            if attempt < max_retries - 1:
                time.sleep(2)  # tunggu 2 detik sebelum retry
            else:
                return f"Search tidak tersedia saat ini, menggunakan pengetahuan model: {query}"
    return "Search tidak tersedia."

@tool
def summarize_text(text: str) -> str:
    """Summarize a given text into key bullet points. Input should be the text to summarize."""
    response = llm.invoke(f"Summarize the following text concisely in 3-5 bullet points:\n\n{text}")
    return response.content

tools = [web_search, summarize_text]
agent_executor = create_react_agent(llm, tools)

def research_topic(topic: str, language: str = "Indonesia") -> str:
    if not topic.strip():
        return "Masukkan topik yang ingin diriset."
    try:
        query = f"""Research the following topic comprehensively and provide a detailed report in {language}:
Topic: {topic}
Please search for recent information at least 2-3 times with different queries, summarize findings, and structure a final report in {language}."""

        result = agent_executor.invoke({"messages": [{"role": "user", "content": query}]})
        final_answer = result["messages"][-1].content

        format_prompt = f"""Format the following research findings into a well-structured report in {language}:

{final_answer}

Structure:
# Laporan Riset: {topic}

## Temuan Utama
[key findings as bullet points]

## Tren & Perkembangan
[trends and developments]

## Insight & Analisis
[insights and analysis]

## Kesimpulan
[brief conclusion]

IMPORTANT: Do NOT add any appendix, references, or lampiran section. Only include the sections above. Write in {language}."""

        formatted = llm.invoke(format_prompt)
        return formatted.content
    except Exception as e:
        return f"Error: {str(e)}\n\nTip: Coba topik yang lebih spesifik."

css = """
.gradio-container {
    max-width: 900px !important;
    margin: 0 auto !important;
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
}
.header-section {
    text-align: center;
    padding: 2rem 0 1.5rem;
    border-bottom: 1px solid #e5e7eb;
    margin-bottom: 1.5rem;
}
.header-section h1 {
    font-size: 1.5rem;
    font-weight: 600;
    color: #111827;
    margin-bottom: 0.5rem;
}
.header-section p {
    font-size: 0.9rem;
    color: #6b7280;
    margin: 0;
}
.badge-row {
    display: flex;
    gap: 8px;
    justify-content: center;
    flex-wrap: wrap;
    margin-top: 0.75rem;
}
.badge {
    font-size: 0.75rem;
    padding: 3px 10px;
    border-radius: 999px;
    background: #f3f4f6;
    color: #374151;
    border: 1px solid #e5e7eb;
}
footer { display: none !important; }
"""

with gr.Blocks(css=css, title="Research Assistant") as demo:
    gr.HTML("""
    <div class="header-section">
        <h1>Research Assistant</h1>
        <p>Riset topik apapun secara otomatis menggunakan multi-agent AI</p>
        <div class="badge-row">
            <span class="badge">Search Agent</span>
            <span class="badge">Summarize Agent</span>
            <span class="badge">Writer Agent</span>
            <span class="badge">Groq · Llama 3.1</span>
            <span class="badge">DuckDuckGo</span>
        </div>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            topic_input = gr.Textbox(
                label="Topik Riset",
                placeholder="Contoh: Perkembangan AI di Indonesia 2024",
                lines=3,
                container=True
            )
            language_input = gr.Dropdown(
                choices=["Indonesia", "English"],
                value="Indonesia",
                label="Bahasa Output",
                container=True
            )
            research_btn = gr.Button(
                "Mulai Riset",
                variant="primary",
                size="lg"
            )
            gr.Examples(
                label="Contoh Topik",
                examples=[
                    ["Perkembangan AI Generatif di Indonesia 2024"],
                    ["Tren startup teknologi Asia Tenggara"],
                    ["Kebijakan regulasi AI di dunia"],
                    ["Perkembangan electric vehicle di Indonesia"],
                ],
                inputs=topic_input
            )

        with gr.Column(scale=2):
            output = gr.Markdown(
                value="Hasil riset akan muncul di sini...",
                label="Hasil Riset",
                container=True
            )

    research_btn.click(
        fn=research_topic,
        inputs=[topic_input, language_input],
        outputs=output,
        show_progress=True
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))