# import os
# import re
# import streamlit as st
# import ast
# import json
# import openai
# from llama_index.llms.openai import OpenAI
# from llama_index.core.llms import ChatMessage
# from llama_index.llms.anthropic import Anthropic
# from llama_index.llms.mistralai import MistralAI

# import nest_asyncio
# nest_asyncio.apply()

# # import ollama
# # from llama_index.llms.ollama import Ollama
# # from llama_index.core.llms import ChatMessage


# # OpenAI credentials
# # key = os.getenv('OPENAI_API_KEY')
# # openai.api_key = key 
# # os.environ["OPENAI_API_KEY"] = key

# # Anthropic credentials
# # key = os.getenv('CLAUDE_API_KEY')
# # os.environ["ANTHROPIC_API_KEY"] = key

# # Mistral
# key = os.getenv('MISTRAL_API_KEY')
# os.environ["MISTRAL_API_KEY"] = key



# # Streamlit UI
# st.title("Auto Test Case Generation using LLM")

# uploaded_files = st.file_uploader("Upload a python or Java file", type=[".py","java"], accept_multiple_files=True)

# if uploaded_files:
#     for uploaded_file in uploaded_files:
#         with open(f"./data/{uploaded_file.name}", 'wb') as f: 
#             f.write(uploaded_file.getbuffer())
#         st.success("File uploaded...")

#         # Check file type
#         _, file_extension = os.path.splitext(uploaded_file.name)
#         print(file_extension)

#         st.success("Fetching list of functions...")
#         file_path = f"./data/{uploaded_file.name}"
        
#         def extract_functions_from_file(file_path, file_extension):

#             if file_extension == '.py':
#                 with open(file_path, "r") as file:
#                     file_content = file.read()
                
#                 parsed_content = ast.parse(file_content)
#                 methods = {}
                
#                 for node in ast.walk(parsed_content):
#                     if isinstance(node, ast.FunctionDef):
#                         func_name = node.name
#                         func_body = ast.get_source_segment(file_content, node)
#                         methods[func_name] = func_body
                        
#             elif file_extension == '.java':
#                 with open(file_path, 'r') as file:
#                     lines = file.readlines()
                
#                 methods = {}
#                 inside_method = False
#                 method_name = None
#                 method_body = []
#                 brace_count = 0
                
#                 method_signature_pattern = re.compile(r'((?:public|protected|private|static|\s)*)\s+[\w<>\[\]]+\s+(\w+)\s*\([^)]*\)\s*\{')
                
#                 for line in lines:
#                     if not inside_method:
#                         match = method_signature_pattern.search(line)
#                         if match:
#                             modifiers, method_name = match.groups()
#                             inside_method = True
#                             method_body.append(line)
#                             brace_count = line.count('{') - line.count('}')
#                     else:
#                         method_body.append(line)
#                         brace_count += line.count('{') - line.count('}')
#                         if brace_count == 0:
#                             inside_method = False
#                             methods[method_name] = ''.join(method_body)
#                             method_body = []
                            
#                 if 'main' in methods.keys():
#                     del(methods['main'])

#             return methods
        
#         functions = extract_functions_from_file(file_path, file_extension)
#         list_of_functions = list(functions.keys())
#         st.write(list_of_functions)

#         def res(prompt, model=None):

#             # response = openai.chat.completions.create(
#             #     model=model,
#             #     messages=[
#             #         {"role": "user",
#             #          "content": prompt,
#             #         }
#             #     ]
#             # )

#             # return response.choices[0].message.content

#             response = [
#                 ChatMessage(role="system", content="You are a sincere and helpful coding assistant"),
#                 ChatMessage(role="user", content=prompt),
#             ]
#             # resp = Anthropic(model=model).chat(response)
#             resp = MistralAI(model).chat(response)
#             return resp

#         # Initialize session state for chat messages
#         if "messages" not in st.session_state:
#             st.session_state.messages = []

#          # Display chat messages from history on app rerun
#         for message in st.session_state.messages:
#             with st.chat_message(message["role"]):
#                 st.markdown(message["content"])

#         # Accept user input
#         if func := st.chat_input("Enter the function name for generating test cases:"):
#             st.session_state.messages.append({"role": "assistant", "content": f"Generating test cases for {func}"})
#             st.success(f"Generating test cases for {func}")

#             func = ''.join(func.split())

#             if func not in list_of_functions:
#                 st.write("Incorrect function name")
            
#             else:
#                 snippet = functions[func]

#                 # Generation
#                 # model = "gpt-3.5-turbo"
#                 # model = "claude-3-haiku-20240307"
#                 # model = "claude-3-sonnet-20240229"
#                 # model = "claude-3-opus-20240229"
#                 model = "codestral-latest"
        
#                 # Generation
#                 # resp = ollama.generate(model='codellama', 
#                 #                     prompt=f""" Your task is to generate unit test cases for this function : {snippet}\
#                 #                     \n\n Politely refuse if the function is not suitable for generating test cases.
#                 #                     \n\n Generate atleast 5 unit test case. Include couple of edge cases as well.
#                 #                     \n\n There should be no duplicate test cases.
#                 #                   \n\n Avoid generating repeated statements.
#                 #                     """)

#                 prompt=f""" Your task is to generate unit test cases for this function : \n\n{snippet}\
#                         \n\n Generate between 3 to 8 unique unit test cases. Include couple of edge cases as well.
#                         \n\n All the test cases should have the mandatory assert statement. 
#                         \n\n Every test case should be defined as a method inside the class.
#                         \n\n All the test cases should have textual description.
#                         \n\n Politely refuse if the function is not suitable for generating test cases.
#                         \n\n There should be no duplicate and incomplete test case. 
#                         \n\n Avoid generating repeated statements.
#                         \n\n Recheck your response before generating.
#                         \n\n Do not share the last Test Case. 
#                     """

#                 # print(prompt)
              
#                 resp = res(prompt = prompt, model = model)

#                 # Post Processing
#                 post_prompt = f"""Except the last test case, display everything that is present in this end to end: \n\n{resp}\
#                         \n\n Do not add anything extra. Just copy and paste everything except the last test case.
#                         \n\n Do not mention the count of total number of test cases in the response.
#                         \n\n Do not mention this sentence - "I have excluded the last test case as per your request"
#                     """
#                 post_resp = res(prompt = post_prompt, model = model)
#                 st.session_state.messages.append({"role": "assistant", "content": f"{post_resp}"})
#                 st.markdown(post_resp)
#                 # st.session_state.messages.append({"role": "assistant", "content": f"{resp['response']}"})
#                 # st.markdown(resp['response'])




import os
import re
import ast
import streamlit as st
from llama_index.llms.openai import OpenAI
from llama_index.core.llms import ChatMessage
from llama_index.llms.anthropic import Anthropic
from llama_index.llms.mistralai import MistralAI
import nest_asyncio

class TestCaseGenerator:
    def __init__(self):
        nest_asyncio.apply()
        self.key = os.getenv('MISTRAL_API_KEY')
        os.environ["MISTRAL_API_KEY"] = self.key
        self.model = "codestral-latest"
        self.functions = {}
        self.list_of_functions = []

    def setup_streamlit_ui(self):
        st.title("Auto Test Case Generation using LLM")
        uploaded_files = st.file_uploader("Upload a python or Java file", type=[".py", "java"], accept_multiple_files=True)
        if uploaded_files:
            for uploaded_file in uploaded_files:
                self.process_uploaded_file(uploaded_file)

    def process_uploaded_file(self, uploaded_file):
        with open(f"./data/{uploaded_file.name}", 'wb') as f:
            f.write(uploaded_file.getbuffer())
        st.success("File uploaded...")
        _, file_extension = os.path.splitext(uploaded_file.name)
        print(file_extension)
        st.success("Fetching list of functions...")
        file_path = f"./data/{uploaded_file.name}"
        self.extract_functions_from_file(file_path, file_extension)
        st.write(self.list_of_functions)

    def extract_functions_from_file(self, file_path, file_extension):
        if file_extension == '.py':
            self.extract_python_functions(file_path)
        elif file_extension == '.java':
            self.extract_java_functions(file_path)
        if 'main' in self.functions.keys():
            del(self.functions['main'])
        self.list_of_functions = list(self.functions.keys())

    def extract_python_functions(self, file_path):
        with open(file_path, "r") as file:
            file_content = file.read()
        parsed_content = ast.parse(file_content)
        for node in ast.walk(parsed_content):
            if isinstance(node, ast.FunctionDef):
                func_name = node.name
                func_body = ast.get_source_segment(file_content, node)
                self.functions[func_name] = func_body

    def extract_java_functions(self, file_path):
        with open(file_path, 'r') as file:
            lines = file.readlines()
        inside_method = False
        method_name = None
        method_body = []
        brace_count = 0
        method_signature_pattern = re.compile(r'((?:public|protected|private|static|\s)*)\s+[\w<>\[\]]+\s+(\w+)\s*\([^)]*\)\s*\{')
        for line in lines:
            if not inside_method:
                match = method_signature_pattern.search(line)
                if match:
                    modifiers, method_name = match.groups()
                    inside_method = True
                    method_body.append(line)
                    brace_count = line.count('{') - line.count('}')
            else:
                method_body.append(line)
                brace_count += line.count('{') - line.count('}')
                if brace_count == 0:
                    inside_method = False
                    self.functions[method_name] = ''.join(method_body)
                    method_body = []

    # def generate_response(self, prompt):
    #     response = [
    #         ChatMessage(role="system", content="You are a sincere and helpful coding assistant"),
    #         ChatMessage(role="user", content=prompt),
    #     ]
    #     resp = MistralAI(self.model).chat(response)
    #     return resp

    def generate_response(self, prompt):
        response = [
            ChatMessage(role="system", content="You are tasked with generating unit test cases, including descriptions and assert statements, for a given function. You will also calculate the test coverage percentage. Follow these instructions carefully:"),
            ChatMessage(role="user", content=prompt),
        ]
        resp = MistralAI(self.model).chat(response)
        return resp

    def generate_test_cases(self, func):
        if func not in self.list_of_functions:
            st.write("Incorrect function name")
            return

        snippet = self.functions[func]
        , file_extension = os.path.splitext(uploaded_file.name)
        lang = file_extension[1:]
       
        prompt = f"""1. You will be provided with the following inputs:
                           <function_code>
                           {snippet}
                           </function_code>

                           <language>{lang}</language> 

                    2. Analyze the provided function:
                       - Identify the function name, parameters, and return type
                       - Determine the main logic and branches in the function
                       - Note any potential edge cases or boundary conditions

                    3. Generate unit test cases:
                       - Create at least 3-5 test cases that cover different scenarios
                       - Include normal cases, edge cases, and potential error conditions
                       - Ensure that the test cases collectively cover all branches and logic paths in the function

                    4. For each test case, provide:
                       - A brief description of the test scenario
                       - Input values for the function parameters
                       - The expected output or behavior
                       - An assert statement in the appropriate syntax for the given programming language

                    5. Calculate the test coverage percentage:
                       - Determine the number of code paths or branches in the function
                       - Count how many of these paths are covered by your test cases
                       - Calculate the percentage: (covered paths / total paths) * 100

                    6. Present your output in the following format:
                           <unit_tests>
                           <test_case>
                           <description>Description of the test case</description>
                           <input>Input values for the function</input>
                           <expected_output>Expected output or behavior</expected_output>
                           <assert_statement>Assert statement in the appropriate language syntax</assert_statement>
                           </test_case>
                           <!-- Repeat for each test case -->
                           </unit_tests>

                           <test_coverage>
                           <percentage>Calculated test coverage percentage</percentage>
                           <explanation>Brief explanation of how the percentage was calculated</explanation>
                           </test_coverage>

                    Remember to adapt your assert statements and syntax to the specific programming language provided. If you're unsure about the exact syntax for a particular language, use a general pseudocode format that clearly conveys the assertion logic.
        """
        resp = self.generate_response(prompt)
        
        # post_prompt = f"""Except the last test case, display everything that is present in this end to end: \n\n{resp}\
        # \n\n - Do not add anything extra. Just copy and paste everything except the last test case.
        # \n\n - Do not mention the count of total number of test cases in the response.
        # \n\n - Do not mention this sentence - "I have excluded the last test case as per your request"
        # """
        # post_resp = self.generate_response(post_prompt)
        return resp

    def run(self):
        self.setup_streamlit_ui()
        if "messages" not in st.session_state:
            st.session_state.messages = []

        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        if func := st.chat_input("Enter the function name for generating test cases:"):
            st.session_state.messages.append({"role": "assistant", "content": f"Generating test cases for {func}"})
            st.success(f"Generating test cases for {func}")
            func = ''.join(func.split())
            test_cases = self.generate_test_cases(func)
            st.session_state.messages.append({"role": "assistant", "content": f"{test_cases}"})
            st.markdown(test_cases)

if __name__ == "__main__":
    test_case_generator = TestCaseGenerator()
    test_case_generator.run()