from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import requests
import subprocess
print("Start Start:")
print("Start Start:")
print("Start Start:")
print("Start Start:")
print("Start Start:")

# Install ollama using the provided script
install_command = "curl https://ollama.ai/install.sh | sh"
install_result = subprocess.run(install_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

# Print the installation output
print("Installation Output:")
print(install_result.stdout)

# Print any installation errors
print("Installation Errors:")
print(install_result.stderr)

# Check the return code of the installation
if install_result.returncode == 0:
    # Installation was successful, now run the 'ollama pull mistral' command
    pull_command = "ollama pull mistral"
    
    # Run the pull command
    pull_result = subprocess.run(pull_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

    # Print the pull output
    print("\nPull Output:")
    print(pull_result.stdout)

    # Print any pull errors
    print("Pull Errors:")
    print(pull_result.stderr)

    # Print the return code of the pull command
    print("Pull Return Code:", pull_result.returncode)

    # Check if the pull command was successful before starting the server
    if pull_result.returncode == 0:
        # Start the server using 'ollama serve'
        serve_command = "ollama serve"
        
        # Run the server start command
        serve_result = subprocess.run(serve_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

        # Print the server start output
        print("\nServer Start Output:")
        print(serve_result.stdout)

        # Print any server start errors
        print("Server Start Errors:")
        print(serve_result.stderr)

        # Print the return code of the server start command
        print("Server Start Return Code:", serve_result.returncode)
    else:
        print("Pull command failed. Aborting server start.")
else:
    print("Installation failed. Aborting pull and server start.")


# Command to run
command = "ollama pull mistral"

# Run the command
result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

# Print the output
print("Output:")
print(result.stdout)

# Print any errors
print("Errors:")
print(result.stderr)

# Print the return code
print("Return Code:", result.returncode)





def get_ollam_response(prompt):
    response = requests.post(
        "http://localhost:11434/api/generate",
        json={
            "model": "mistral",
            "prompt": prompt,
            "stream": False
        }
    )
    resp = response.json()
    return resp["response"], resp["eval_count"]


app = FastAPI(openapi_url="/api/v1/sparrow-data/openapi.json", docs_url="/api/v1/sparrow-data/docs")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
    allow_credentials=True,
)



@app.post("/ask_llm")
def ask_HFAPI_endpoint(prompt: str):
    result = get_ollam_response(prompt)
    return {"result": result}