import os
import uuid
import gradio as gr

from huggingface_hub import InferenceClient
from duckduckgo_search import DDGS
import chromadb

# -----------------------------
# LLM CONFIGURATION
# -----------------------------

MODEL_NAME = "openai/gpt-oss-20b"

HF_TOKEN = os.getenv("HF_TOKEN")

# -----------------------------
# VECTOR MEMORY SETUP
# -----------------------------

client_db = chromadb.PersistentClient(path="./medini_memory")

collection = client_db.get_or_create_collection(
    name="conversation_memory"
)

# -----------------------------
# WEB SEARCH TOOL
# -----------------------------

def web_search(query):

    results = []

    try:

        with DDGS() as ddgs:

            for r in ddgs.text(query, max_results=3):

                body = r.get("body", "")

                if body:
                    results.append(body)

    except Exception as e:

        return f"Web search failed: {str(e)}"

    return "\n".join(results)

# -----------------------------
# MEMORY FUNCTIONS
# -----------------------------

def save_memory(user_message, assistant_response):

    memory_text = f"""
USER: {user_message}

ASSISTANT: {assistant_response}
"""

    collection.add(
        documents=[memory_text],
        ids=[str(uuid.uuid4())]
    )

def retrieve_memory(query):

    try:

        results = collection.query(
            query_texts=[query],
            n_results=3
        )

        docs = results.get("documents", [[]])[0]

        return "\n".join(docs)

    except Exception:
        return ""

# -----------------------------
# PLANNER AGENT
# -----------------------------

def planner(user_input):

    user_input = user_input.lower()

    if any(
        keyword in user_input
        for keyword in [
            "search",
            "latest",
            "news",
            "find",
            "lookup",
            "web",
        ]
    ):
        return "web_search"

    return "chat"

# -----------------------------
# MAIN AI AGENT
# -----------------------------

def agent_respond(
    message,
    history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):

    # -----------------------------
    # CONNECT MODEL
    # -----------------------------

    client = InferenceClient(
        model=MODEL_NAME,
        token=HF_TOKEN
    )

    # -----------------------------
    # PLANNING
    # -----------------------------

    action = planner(message)

    tool_context = ""

    # -----------------------------
    # TOOL EXECUTION
    # -----------------------------

    if action == "web_search":

        tool_context = web_search(message)

    # -----------------------------
    # MEMORY RETRIEVAL
    # -----------------------------

    memory_context = retrieve_memory(message)

    # -----------------------------
    # SYSTEM PROMPT
    # -----------------------------

    enhanced_system_prompt = f"""
{system_message}

You are Medini Intelligence AI Agent.

You have:
- long-term memory
- web search capability
- contextual reasoning

MEMORY:
{memory_context}

TOOL RESULTS:
{tool_context}

Use the information intelligently.
"""

    # -----------------------------
    # BUILD MESSAGE HISTORY
    # -----------------------------

    messages = [
        {
            "role": "system",
            "content": enhanced_system_prompt,
        }
    ]

    if history:
        messages.extend(history)

    messages.append(
        {
            "role": "user",
            "content": message
        }
    )

    # -----------------------------
    # STREAM RESPONSE
    # -----------------------------

    response = ""

    try:

        stream = client.chat.completions.create(
            messages=messages,
            model=MODEL_NAME,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        )

        for chunk in stream:

            token = ""

            if (
                chunk.choices
                and hasattr(chunk.choices[0].delta, "content")
                and chunk.choices[0].delta.content
            ):
                token = chunk.choices[0].delta.content

            response += token

            yield response

    except Exception as e:

        yield f"Error: {str(e)}"
        return

    # -----------------------------
    # SAVE MEMORY
    # -----------------------------

    save_memory(message, response)

# -----------------------------
# GRADIO UI
# -----------------------------

chatbot = gr.ChatInterface(
    fn=agent_respond,
    type="messages",
    additional_inputs=[

        gr.Textbox(
            value="You are Medini Intelligence AI Agent.",
            label="System Message",
        ),

        gr.Slider(
            minimum=1,
            maximum=4096,
            value=1024,
            step=1,
            label="Max Tokens",
        ),

        gr.Slider(
            minimum=0.1,
            maximum=2.0,
            value=0.7,
            step=0.1,
            label="Temperature",
        ),

        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p",
        ),
    ],
)

# -----------------------------
# APP LAYOUT
# -----------------------------

with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown(
        """
# Medini Intelligence AI Agent

### Features
- Conversational AI
- Memory
- Web Search
- Tool Use
- Autonomous Reasoning
"""
    )

    chatbot.render()

# -----------------------------
# RUN APP
# -----------------------------

if __name__ == "__main__":

    demo. launch(
        server_name="0.0.0.0",
        server_port=7860
    )