#!/usr/bin/env python3
"""Get chunk alignment from ChatGPT.

The two prompts live in ``prompts/`` rather than in this file: they are 2 kB
of instructions each, on lines far past any line-length limit, and wrapping
them here would put newlines into the middle of the prompt itself.
"""

import os
from pathlib import Path
from typing import TYPE_CHECKING

from openai import OpenAI

if TYPE_CHECKING:
    import pandas as pd

PROMPTS = Path(__file__).parent / "prompts"
THE_PROMPT = (PROMPTS / "alignment.txt").read_text(encoding="utf-8")
THE_PROMPT_WITH_COMMENT = (PROMPTS / "alignment_with_comment.txt").read_text(
    encoding="utf-8"
)


def create_gpt() -> OpenAI:
    """Build an OpenAI client from OPENAI_API_KEY.

    The key used to be a placeholder string literal here, which meant the
    script could not actually run as committed.
    """
    key = os.environ.get("OPENAI_API_KEY")
    if not key:
        msg = "set OPENAI_API_KEY to run the alignment scripts"
        raise RuntimeError(msg)
    return OpenAI(api_key=key)


def call_api(client: OpenAI, chunks: str) -> str:
    """Ask the model to align one pair of chunked sentences."""
    response = client.chat.completions.create(
        model="gpt-3.5-turbo-1106",
        messages=[
            {"role": "user", "content": THE_PROMPT},
            {"role": "user", "content": chunks},
        ],
    )
    return response.choices[0].message.content.strip()


def call_api_examples(client: OpenAI, examples: pd.DataFrame, alignment: str) -> str:
    """Ask the model to align, few-shot prompted with worked examples."""
    user_input = []
    assistant_output = []
    for _index, row in examples.iterrows():
        chunks = ""
        for chunk in row["chunked_sentance1"]:
            chunks = chunks + "[ " + chunk + " ] "
        chunks = chunks + "\n"
        for chunk in row["chunked_sentance2"]:
            chunks = chunks + "[ " + chunk + " ] "
        user_input.append(chunks)
        assistant_output.append(row["alignment_text"])

    messages = []
    messages.append({"role": "user", "content": THE_PROMPT})
    for u, a in zip(user_input, assistant_output, strict=True):
        messages.append({"role": "user", "content": u})
        messages.append({"role": "assistant", "content": a})

    messages.append({"role": "user", "content": alignment})

    response = client.chat.completions.create(
        model="gpt-3.5-turbo-1106", messages=messages
    )
    return response.choices[0].message.content.strip()