import os
import pandas as pd
import streamlit as st
from datetime import datetime
from groq import Groq
from logic import LLMClient, CodeProcessor
from batch_code_logic_csv import csv_read_batch_code
import zipfile
import io
import google.generativeai as genai
import re


client = Groq(api_key=os.getenv("GROQ_API_KEY"))

GOOGLE_API_KEY=os.getenv('GOOGLE_API_KEY')
llm_obj = LLMClient(client,GOOGLE_API_KEY)
processor = CodeProcessor(llm_obj)

st.title("Code Analysis with LLMs")

st.sidebar.title("Input Options")

code_input_method = st.sidebar.radio("How would you like to provide your code?",
                                     ("Upload CSV file", "Upload Code File"))

code_dict = {}

if code_input_method == "Upload CSV file":
    uploaded_file = st.sidebar.file_uploader("Upload your CSV/Excel file", type=["csv", "xlsx"])

    if uploaded_file is not None:
        dataframe = pd.read_csv(uploaded_file)
        code_dict = csv_read_batch_code(dataframe)

elif code_input_method == "Upload Code File":
    uploaded_file = st.sidebar.file_uploader("Upload your code file", type=["py", "txt"])
    if uploaded_file is not None:
        code_text = uploaded_file.read().decode("utf-8")

        code_lines = [line.strip() for line in code_text.splitlines()]
        code_lines = [re.sub(r"\s+", " ", line) for line in code_lines]

        # Join the cleaned lines back if needed as a single string
        clean_code_text = "\n".join(code_lines)  # Combine lines into one string

        code_dict = {"single_code": clean_code_text}

model_choice = st.sidebar.selectbox("Select LLM Model",
                                    ["llama-3.2-90b-text-preview", "llama-3.2-90b-text-preview", "llama3-8b-8192","llama-3.1-70b-versatile","gemma2-9b-it","gemini-pro"])

if code_dict:
    unique_key = st.sidebar.selectbox("Select a Key for Analysis", list(code_dict.keys()))

    if st.sidebar.button("Analyze Code") and unique_key:
        code_text = code_dict[unique_key]
        markdown_output = processor.process_code(code_text, model_choice)

        with st.expander(f"Analysis for {unique_key}"):
            st.markdown(markdown_output)

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        st.download_button(
            label=f"Download {unique_key} Result as Markdown",
            data=markdown_output,
            file_name=f"code_analysis_{unique_key}_{timestamp}.md",
            mime="text/markdown"
        )

    if st.sidebar.button("Batch Predict"):
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        all_markdowns = {}
        for key, code_text in code_dict.items():
            markdown_output = processor.process_code(code_text, model_choice)
            all_markdowns[key] = markdown_output

            with st.expander(f"Analysis for {key}"):
                st.markdown(markdown_output)

            st.download_button(
                label=f"Download {key} Result as Markdown",
                data=markdown_output,
                file_name=f"code_analysis_{key}_{timestamp}.md",
                mime="text/markdown"
            )

        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w") as zip_file:
            for key, markdown_output in all_markdowns.items():
                zip_file.writestr(f"code_analysis_{key}_{timestamp}.md", markdown_output)

        st.download_button(
            label="Download All as Zip",
            data=zip_buffer.getvalue(),
            file_name=f"code_analysis_batch_{timestamp}.zip",
            mime="application/zip"
        )
else:
    st.write("Please upload your file to analyze.")