from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import requests
import os
from dotenv import load_dotenv

load_dotenv()

app = FastAPI(title="Ollama Workflow Generator API")

OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "phi3:mini")


class PromptRequest(BaseModel):
    prompt: str


@app.post("/api/prompt")
def generate_workflow(request: PromptRequest):
    payload = {
        "model": OLLAMA_MODEL,
        "prompt": request.prompt
    }

    try:
        response = requests.post(f"{OLLAMA_BASE_URL}/api/generate", json=payload, stream=True)
        response.raise_for_status()
        
        full_response = ""
        for line in response.iter_lines():
            if line:
                data = line.decode("utf-8")
                if '"response"' in data:
                    full_response += data.split('"response":"')[1].split('"')[0]
        
        return {"model": OLLAMA_MODEL, "response": full_response.strip()}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/")
def root():
    return {"message": "Ollama Workflow Generator API is running!"}