import os
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
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

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

uploaded_files = st.file_uploader("Upload a python(.py) file", type=".py", 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...")

        st.success("Fetching list of functions...")
        file_path = f"./data/{uploaded_file.name}"
        def extract_functions_from_file(file_path):
            with open(file_path, "r") as file:
                file_content = file.read()
            
            parsed_content = ast.parse(file_content)
            functions = {}
            
            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)
                    functions[func_name] = func_body
            
            return functions
        
        functions = extract_functions_from_file(file_path)
        list_of_functions = list(functions.keys())
        st.write(list_of_functions)

        def res(prompt, model):

            # 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)
            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"

                # 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 5 to 10 unit test cases. Include couple of edge cases as well.
                        \n\n Politely refuse if the function is not suitable for generating test cases.
                        \n\n There should be no duplicate test cases. 
                        \n\n Avoid generating repeated statements.
                    """

                # print(prompt)
              
                resp = res(prompt, model)
                st.session_state.messages.append({"role": "assistant", "content": f"{resp}"})
                st.markdown(resp)
                # st.session_state.messages.append({"role": "assistant", "content": f"{resp['response']}"})
                # st.markdown(resp['response'])