import json
import csv
import io
import requests
import html  # For escaping HTML characters
from bs4 import BeautifulSoup
from openai import OpenAI 

# Initialize OpenAI API with Nvidia's Mistral model
client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key="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"
)

def clean_test_case_output(text):
    """
    Cleans the output to handle HTML characters and unwanted tags.
    """
    # Unescape HTML entities (convert &lt; back to < and &gt; back to >)
    text = html.unescape(text)
    
    # Use BeautifulSoup to handle HTML tags more comprehensively
    soup = BeautifulSoup(text, 'html.parser')
    
    # Convert <br> tags to newlines and remove all other tags
    cleaned_text = soup.get_text(separator="\n").strip()
    
    return cleaned_text

def generate_testcases(user_story):
    """
    Generates advanced QA test cases based on a provided user story by interacting 
    with Nvidia's Mistral model API. The prompt is refined for clarity, 
    and the output is processed for better quality.
    
    :param user_story: A string representing the user story for which to generate test cases.
    :return: A list of test cases in the form of dictionaries.
    """
    try:
        # Example few-shot learning prompt to guide the model
        completion = client.chat.completions.create(
            model="nv-mistralai/mistral-nemo-12b-instruct",  # Using Mistral model
            messages=[
                {"role": "user", "content": f"Generate QA test cases for the following user story: {user_story}"}
            ],
            temperature=0.06,  # Further lowering temperature for precise and deterministic output
            top_p=0.5,  # Prioritize high-probability tokens even more
            max_tokens=4096,  # Increase max tokens to allow longer content
            stream=True  # Streaming the response for faster retrieval
        )

        # Initialize an empty string to accumulate the response
        test_cases_text = ""
        
        # Accumulate the response from the streaming chunks
        for chunk in completion:
            if chunk.choices[0].delta.content is not None:
                test_cases_text += chunk.choices[0].delta.content


        # Ensure the entire response is captured before cleaning
        if test_cases_text.strip() == "":
            return [{"test_case": "No test cases generated or output was empty."}]
        
        # Clean the output by unescaping HTML entities and replacing <br> tags
        test_cases_text = clean_test_case_output(test_cases_text)

        try:
            # Try to parse the output as JSON, assuming the model returns structured test cases
            test_cases = json.loads(test_cases_text)
            if isinstance(test_cases, list):
                return test_cases  # Return structured test cases

            else:
                return [{"test_case": test_cases_text}]  # Return as a list with the text wrapped in a dict

        except json.JSONDecodeError:
            # Fallback: return the raw text if JSON parsing fails
            return [{"test_case": test_cases_text}]
    
    except requests.exceptions.RequestException as e:
        print(f"API request failed: {str(e)}")
        return []
# Add options for multiple test case formats
def export_test_cases(test_cases, format='json'):
    if not test_cases:
        return "No test cases to export."

    # Convert test cases (which are currently strings) into a structured format for CSV
    structured_test_cases = [{'Test Case': case} for case in test_cases]

    if format == 'json':
        # Improve JSON export to be line-by-line formatted
        return json.dumps(test_cases, indent=4, separators=(',', ': '))  # More readable format
    elif format == 'csv':
        # Check if test_cases is a list of dicts
        if isinstance(test_cases, list) and isinstance(test_cases[0], dict):
            output = io.StringIO()
            csv_writer = csv.DictWriter(output, fieldnames=test_cases[0].keys(), quoting=csv.QUOTE_ALL)
            csv_writer.writeheader()
            csv_writer.writerows(test_cases)
            return output.getvalue()
        else:
            raise ValueError("Test cases must be a list of dictionaries for CSV export.")

# 2. Save test cases as a downloadable file
def save_test_cases_as_file(test_cases, format='json'):
    if not test_cases:
        return "No test cases to save."

    if format == 'json':
        with open('test_cases.json', 'w') as f:
            json.dump(test_cases, f)
    elif format == 'csv':
        with open('test_cases.csv', 'w', newline='') as file:
            dict_writer = csv.DictWriter(file, fieldnames=test_cases[0].keys())
            dict_writer.writeheader()
            dict_writer.writerows(test_cases)
    else:
        return f"Unsupported format: {format}"
    return f'{format} file saved'