import os
import json
import http.client
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

class SearchRequest(BaseModel):
    query: str

@app.post("/search")
async def search_ollama(request: SearchRequest):
    api_key = os.getenv("OLLAMA_API_KEY")
    if not api_key:
        return {"error": "OLLAMA_API_KEY not found in environment variables"}

    conn = http.client.HTTPSConnection("ollama.com")
    payload = json.dumps({"query": request.query})
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    conn.request("POST", "https://ollama.com/api/web_search", body=payload, headers=headers)
    res = conn.getresponse()
    data = res.read()
    conn.close()

    try:
        return json.loads(data)
    except json.JSONDecodeError:
        return {"error": "Failed to parse response", "raw": data.decode()}

class FetchRequest(BaseModel):
    url: str

@app.post("/web-fetch")
async def web_fetch(request: FetchRequest):
    api_key = os.getenv("OLLAMA_API_KEY")
    if not api_key:
        return {"error": "OLLAMA_API_KEY not found in environment variables"}

    conn = http.client.HTTPSConnection("https://ollama.com/api/web_fetch")
    payload = json.dumps({"url": request.url})
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    conn.request("POST", "/api/web_fetch", body=payload, headers=headers)
    res = conn.getresponse()
    data = res.read()
    conn.close()

    try:
        return json.loads(data)
    except json.JSONDecodeError:
        return {"error": "Failed to parse response", "raw": data.decode("utf-8", errors="replace")}