# Koyeb Agent.py
import streamlit as st
import os
from smolagents import (
    CodeAgent,
    ToolCallingAgent,
    HfApiModel,
    DuckDuckGoSearchTool,
    LiteLLMModel,
    tool,
    VisitWebpageTool,
    GradioUI,
    PythonInterpreterTool
)
import requests
import json
from typing import List, Dict, Any

# Configuration
PAGE_CONFIG = {
    "page_title": "Koyeb Documentation Search",
    "page_icon": "📘",
    "layout": "wide"
}

AUTHORIZED_IMPORTS = [
    "requests", "os", "pandas", "markdownify", "numpy", "sympy", "json", "bs4",
    "pubchempy", "xml", "yahoo_finance", "Bio", "sklearn", "scipy", "pydub",
    "io", "PIL", "chess", "PyPDF2", "pptx", "torch", "datetime", "fractions",
    "csv", "huggingface_hub"
]

# Model configuration
MODEL_CONFIGS = {
    "nano": {
        "model_id": "openrouter/openai/gpt-4.1-nano",
        "temperature": 0.6
    },
    "dolphin": {
        "model_id": "openrouter/cognitivecomputations/dolphin3.0-r1-mistral-24b:free",
        "temperature": 0.5
    }
}

# Agent Initialization
def initialize_koyeb_agent(model=None):
    """Initialize the Koyeb Documentation Agent"""
    tools = [
        DuckDuckGoSearchTool(),
        VisitWebpageTool()
    ]
    if model is None:
        from smolagents import HfApiModel
        model = LiteLLMModel(
            model_id=MODEL_CONFIGS["nano"]["model_id"],
            api_key=os.environ["OPENROUTER_API_KEY"],
            temperature=MODEL_CONFIGS["nano"]["temperature"]
        )
    
    return CodeAgent(
        name="koyeb_agent",
        description="Koyeb Documentation AI Agent",
        tools=tools,
        model=model,
        max_steps=12,
        additional_authorized_imports=AUTHORIZED_IMPORTS
    )

# UI Components
def setup_ui():
    """Setup Streamlit UI components"""
    st.set_page_config(**PAGE_CONFIG)
    
    st.title("💻 Koyeb Documentation Search")
    st.markdown("""
    This tool uses an AI agent to help you explore Koyeb documentation.
    Ask questions about Koyeb features, usage, or troubleshooting.
    """)

def display_results(results: str):
    """Display search results in a formatted way"""
    st.markdown("### 📝 Results")
    st.markdown(results, unsafe_allow_html=True)

def setup_sidebar():
    """Setup sidebar content"""
    with st.sidebar:
        st.markdown("### About This Tool")
        st.markdown("""
        This agent specializes in Koyeb documentation, providing:
        - 📚 Quick answers from web searches
        - 🌐 Webpage content retrieval
        - 🤖 AI-powered insights
        """)
        st.markdown("### Tips")
        st.markdown("""
        - Use specific queries like "How to install Koyeb cli" or "Koyeb deepseek model deployment".
        - Results may include raw data from searches or webpages.
        """)

# Main App
@st.cache_resource
def initialize_app():
    """Initialize the application"""
    return initialize_koyeb_agent()

def main():
    setup_ui()
    setup_sidebar()
    
    if 'agent' not in st.session_state:
        with st.spinner("Initializing Koyeb Documentation Agent..."):
            st.session_state.agent = initialize_app()

    search_query = st.text_input(
        "🔍 Search Koyeb Documentation",
        placeholder="E.g., How to use Terraform with the koyeb provider ?"
    )

    # Single button to handle both cases
    if st.button("Search", type="primary", key="search_button"):
        if search_query:
            with st.spinner("Searching Koyeb documentation..."):
                try:
                    results = st.session_state.agent.run(search_query, additional_args={"reference_websites": ['https://www.koyeb.com/docs', 'https://www.koyeb.com/tutorials', 'https://www.koyeb.com/deploy'],"instructions": "Response to the user answer as a structured markdown documentation format."})
                    display_results(results)
                    
                    st.markdown("---")
                    st.info("""
                    💡 **How to read the results:**
                    - Results may include JSON from searches or webpage snippets.
                    - Check the sidebar for tips on refining your query.
                    """)
                except Exception as e:
                    st.error(f"An error occurred: {e}")
        else:
            st.warning("Please enter a search query.")

    st.markdown("---")
    st.caption("Powered by SmolAgents and Koyeb")

if __name__ == "__main__":
    main()