from openai import OpenAI
from tenacity import (
    retry,
    stop_after_attempt,
    wait_random_exponential,
)


PQA_QUESTION2QUESTION_PROMPT = """\
Given a question, decompose it into multiple atomic questions.

### Instructions:
- The questions MUST be atomic, i.e., they MUST be answerable by only a single piece of information.
- The questions MUST be standalone, i.e., they MUST NOT reference any other question or the given question.
- The questions can be both open-ended or yes/no questions.
- The questions should be decomposed only from the main question.
- Each question should be on a new line and start with `**** `.

### Input:
{text}

### Output:
"""

qa_client = OpenAI(
    base_url="http://130.85.37.21:4774/v1",
    api_key="EMPTY"
)

@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(7))
def completion_with_backoff(client, **kwargs):
    return client.chat.completions.create(**kwargs)

def generation_to_questions(generated_text, header, numbered=True):
    try:
        lines = generated_text.split("\n")
        lines = [line.strip() for line in lines]
        lines = [line for line in lines if line.startswith(header)]
        lines = [line.replace(header, "").strip() for line in lines]
    except:
        lines = []
        print("Error in processing generated text")
    return lines

def get_questions(claim):
    prompt = PQA_QUESTION2QUESTION_PROMPT.format(text=claim)
    print(f"Question Generation Prompt: {prompt}")
    response = completion_with_backoff(
        client=qa_client,
        model="dipta007/Llama-3.1-8B-Instruct-finetuned-pqa",
        messages=[
            {"role": "user", "content": prompt},
        ],
        max_tokens=2048,
        top_p=1.0,
        temperature=0.0,
    )
    print(f"Questions: {response.choices[0].message}")
    generation = response.choices[0].message.content
    questions = generation_to_questions(generation, "****")
    questions = [f"- {q}" for q in questions]
    return "\n".join(questions)