import requests
from fastapi import FastAPI
from pydantic import BaseModel
from datetime import date
from ollama import Client
import os

app = FastAPI()
today_date = date.today()

JINA_PASS = os.environ.get("JINA_AUTH")
class PromptRequest(BaseModel):
    prompt: str
    temperature: float = 0.5

TOOL_PROMPT = f"""You are an autonomous AI Agent with real-time internet access.
Current Date: {today_date}

You have access to one tool:
- Search [query] : searches the internet and returns real-time results

RULES:
- If the user query needs current/real-time/recent information, respond with ONLY: Search [your search query]
- If you can answer from your own knowledge, answer directly
- Do NOT mix search command with other text
- After getting search results, answer the user query using those results

"""

def search_tool(query: str) -> str:
    try:
        r = requests.get(
            "https://s.jina.ai/",
            params={"q": query},
            headers={"Accept": "application/json", "Authorization":f"{JINA_PASS}"},
            timeout=120
        )
        data = r.json().get("data", [])[:5]
        results = []
        for item in data:
            title = item.get("title", "")
            desc = item.get("description", "") or item.get("content", "")[:300]
            url = item.get("url", "")
            results.append(f"- {title}\n  {desc}\n  Source: {url}")
        return "\n\n".join(results) if results else "No results found."
    except Exception as e:
        return f"Search failed: {str(e)}"

def run_llm(messages: list, temperature: float) -> str:
    client = Client(host="http://localhost:11434")
    response = client.chat(
        model="qwen2.5-coder:7b",
        messages=messages,
        options={"temperature": temperature}
    )
    return response["message"]["content"].strip()

@app.get("/")
def health():
    return {"ok": True}

@app.post("/qwen2.5-coder:7b")
async def generate_response(request: PromptRequest):
    # Step 1: ask LLM if it needs to search
    messages = [
        {"role": "system", "content": TOOL_PROMPT},
        {"role": "user", "content": request.prompt}
    ]
    response = run_llm(messages, request.temperature)
    print(f"[LLM RAW]: {response}")

    # Step 2: if LLM wants to search
    if response.strip().startswith("Search "):
        query = response.strip().removeprefix("Search ").strip("[]").strip()
        print(f"[SEARCH QUERY]: {query}")
        search_results = search_tool(query)
        print(f"[SEARCH RESULTS]: {search_results[:200]}...")

        # Step 3: feed results back to LLM
        messages.append({"role": "assistant", "content": response})
        messages.append({
            "role": "user",
            "content": f"Search Results:\n{search_results}\n\nNow answer the original question: {request.prompt}"
        })
        final_response = run_llm(messages, request.temperature)
        print(f"[FINAL]: {final_response}")
        return {"response": final_response, "searched": True, "query": query}

    return {"response": response, "searched": False}