import PyPDF2
import gradio as gr
import openai
import os

openai.api_key = os.environ.get("OPENAI_API_KEY")

# Set the default PDF path here
DEFAULT_PDF_PATH = "spruiv7b.pdf"

def extract_text_from_block(pdf_path, block_pages):
    text = ""
    with open(pdf_path, 'rb') as file:
        reader = PyPDF2.PdfReader(file)
        for page_num in range(block_pages[0] - 1, block_pages[1]):
            page = reader.pages[page_num]
            text += page.extract_text()
    return text

def generate_test_cases(text, prompt_template):
    prompt = prompt_template.format(text=text)
    messages = [
        {"role": "system", "content": "You are a helpful assistant that generates test cases based on given text."},
        {"role": "user", "content": prompt}
    ]
    
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4o-mini", 
            messages=messages,
            max_tokens=3000,
            n=1,
            temperature=0.7,
            timeout=30  # Timeout after 30 seconds
        )
    except Exception as e:
        return f"[Error during API call: {str(e)}]"
    
    content = response.choices[0].message['content'].strip()
    finish_reason = response.choices[0].get("finish_reason", None)
    
    # Check if the response appears truncated; if so, warn the user.
    if finish_reason == "length" or content.endswith("..."):
        content += "\n[WARNING: The output may be truncated due to token limits.]"
    
    return content.strip()




def process_pdf(pdf_file, selected_block, prompt_template):
    block_config = {
        "5.4.1 OSPI-xSPI-QSPI-SPI Boot": (482, 488),
        "5.4.2 I2C Boot": (489, 489),
        "5.4.3 SD Card Boot": (490, 491),
        "5.4.4 eMMC Boot": (492, 493),
        "5.4.5 Ethernet Boot": (494, 497),
        "5.4.6 USB Boot": (498, 498),
        "5.4.7 UART Boot": (499, 500),
        "5.4.8 GPMC NOR Boot": (501, 502),
        "5.4.9 GPMC NAND Boot": (503, 503),
        "5.4.10 Serial NAND Boot": (504, 505),
        "5.4.11 No boot/Development boot": (506, 506)
    }

    try:
        if pdf_file is None:
            # Use the default PDF if no file is uploaded
            pdf_path = DEFAULT_PDF_PATH
        else:
            # Use the uploaded PDF
            pdf_path = pdf_file.name

        selected_pages = block_config[selected_block]
        
        text = extract_text_from_block(pdf_path, selected_pages)
        test_cases = generate_test_cases(text, prompt_template)

        return f"Generated Test Cases for {selected_block}:\n\n{test_cases}"
    except Exception as e:
        return f"An error occurred: {str(e)}"

block_names = [
    "5.4.1 OSPI-xSPI-QSPI-SPI Boot",
    "5.4.2 I2C Boot",
    "5.4.3 SD Card Boot",
    "5.4.4 eMMC Boot",
    "5.4.5 Ethernet Boot",
    "5.4.6 USB Boot",
    "5.4.7 UART Boot",
    "5.4.8 GPMC NOR Boot",
    "5.4.9 GPMC NAND Boot",
    "5.4.10 Serial NAND Boot",
    "5.4.11 No boot/Development boot"
]

default_prompt = """You are an expert test engineer with extensive experience in creating detailed, comprehensive test cases. Given the following technical context, generate a complete set of test cases covering every aspect of the system described. Your output must include all test cases—if your response appears to end before all test cases are generated, please continue with "Continuing with remaining test cases..." and then list the rest.

For each test case, include the following sections:

1. **Test Objective:** A clear statement of the purpose of the test.
2. **Pre-requisites:** A detailed list of required hardware, software, configurations, and environmental conditions.
3. **Test Design:** A description of how the functionality will be verified. This should cover:
   - Initialization and configuration checks.
   - Normal operational behavior.
   - Error handling and fallback mechanisms.
   - Performance aspects (if applicable).
4. **Test Procedure:** Step-by-step instructions to execute the test case. Include optional steps (for example, using a logic analyzer) where applicable.
5. **Expected Result:** A detailed explanation of the expected outcome if the system behaves correctly, covering both positive (successful boot/operation) and negative (failure, fallback, or error handling) scenarios.

Important: Ensure that your output includes all relevant test cases. Do not skip or omit any parts—even if the response becomes lengthy. If you reach the token limit, split your answer into multiple parts by indicating "Continuing with remaining test cases..." and then provide the remainder.

Context:
{text}

Generate a complete and thorough set of test cases for the provided context.
"""

app = gr.Interface(
    fn=process_pdf,
    inputs=[
        gr.File(label="Upload PDF (optional)"),
        gr.Dropdown(choices=block_names, label="Select Block"),
        gr.Textbox(label="Prompt Template", value=default_prompt, lines=10)
       
    ],
    outputs="text",
    title="PDF Block Test Case Generator",
    description="Upload a PDF (optional), select a block, edit the prompt if needed, and generate test cases for that block. If no PDF is uploaded, a default PDF will be used."
)

app.launch()