import modal
import openai
from fastapi import Depends, FastAPI, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from modal import Image, Stub
from openai import OpenAI
from pydantic import BaseModel, Field
from pydantic_settings import BaseSettings, SettingsConfigDict

stub = Stub()
auth_scheme = HTTPBearer()
image = Image.debian_slim().pip_install_from_pyproject("pyproject.toml")


class Settings(BaseSettings):
    openai_api_key: str
    erato_token: str

    model_config = SettingsConfigDict(env_file=".env")


settings = Settings()

openai.organization = "org-xJUtqr0WhHZ39GSAdXBeH8mt"
openai.api_key = settings.openai_api_key

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


openai_client = OpenAI(organization=openai.organization, api_key=openai.api_key)

DEFAULT_TEMP = 1.0


class CompletionParams(BaseModel):
    context: str = ""
    cursor: int | None = None
    first_word: bool = Field(False, alias="firstWord")
    max_tokens: int = Field(500, alias="maxTokens")
    temperature: float = DEFAULT_TEMP
    tax: float | None = None


stub = Stub(name="erato")


@app.post("/api/completions")
def make_completions(
    params: CompletionParams, token: HTTPAuthorizationCredentials = Depends(auth_scheme)
):
    if token.credentials != settings.erato_token:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect token",
            headers={"WWW-Authenticate": "Bearer"},
        )
    context = params.context
    cursor = params.cursor if params.cursor is not None else len(context)
    first_word = params.first_word
    max_tokens = params.max_tokens
    temperature = params.temperature
    print("max tokens:", max_tokens)
    SYSTEM_PROMPT = """You're a critically acclaimed songwriter writing a song in collaboration with a musician. You'll be given part of a song and asked to complete a line. The line you'll be asked to complete will be marked with '[LYRICS GO HERE]'. Respond with the text you'd like to replace '[LYRICS GO HERE]'. Only respond with the line, do not copy the whole song or add any preamble. If you wish to start a new word, begin your response with a space."""

    num_responses = 10 if first_word else 3
    user_prompt = context[:cursor] + "[LYRICS GO HERE]" + context[cursor:]
    completion = openai_client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ],
        n=num_responses,
        max_tokens=max_tokens,
        temperature=temperature,
    )
    raw_choices = [(c.message.content) for c in completion.choices]
    print("raw choices:")
    for rc in raw_choices:
        print(rc)
    choices = [_prep_choice(rc) for rc in raw_choices]
    if first_word:
        choices = [_prep_choice(c.split()[0]) for c in choices]

    return {"choices": choices}


def _prep_choice(choice: str):
    if choice.startswith('"'):
        choice = choice[1:]
    if choice.endswith('"'):
        choice = choice[:-1]
    if not choice.startswith(" "):
        choice = " " + choice
    return choice


@stub.function(
    image=image,
    secrets=[
        modal.Secret.from_name("openai-secret"),
        modal.Secret.from_name("erato-token"),
    ],
)
@modal.asgi_app()
def fastapi_app():
    return app