import os

import pandas as pd
import streamlit as st
from streamlit_extras.grid import grid
from streamlit_card import card

from lida_ko import Manager, TextGenerationConfig
from lida_ko.datamodel import Goal
from lida_ko.utils import clean_code_snippet

openai_api_key = os.environ["OPENAI_API_KEY"]
selected_dataset = st.session_state.selected_dataset
selected_method = st.session_state.selected_method
selected_model = st.session_state.selected_model
use_cache = st.session_state.use_cache
temperature = st.session_state.temperature
lida_manager: Manager = st.session_state.lida_manager
summary = st.session_state.summary 
selected_goal_object: Goal = st.session_state.selected_goal_object
selected_dataframe: pd.DataFrame = st.session_state.selected_dataframe
num_visualizations = st.session_state.num_visualizations

def generate_visualizations(code = None, feedback=None):
    if code and feedback:
        visualizations = lida_manager.repair(
            code=code,
            goal=selected_goal_object,
            summary=summary,
            feedback=feedback,
            textgen_config=textgen_config,
            library=st.session_state.selected_library
        )
    else:
        visualizations = lida_manager.visualize(
                summary=summary,
                goal=selected_goal_object,
                textgen_config=textgen_config,
                library=st.session_state.selected_library)
    return visualizations

st.title("📊 데이터 시각화 만들기")
st.write("")
st.empty()

if not selected_goal_object:
    st.error("**ERROR**: 🚨 데이터 분석 목표를 설정해주세요.")
    st.stop()

if selected_goal_object:
    # Update the visualization generation call to use the selected library.

    textgen_config = TextGenerationConfig(
        n=num_visualizations, temperature=temperature,
        model=selected_model,
        use_cache=use_cache)

    # **** lida.visualize *****
    if not st.session_state.update_viz:
        visualizations = generate_visualizations()
        st.session_state.visualizations = visualizations
    else:
        st.session_state.update_viz = False

    col1, col2, col3 = st.columns([5, 0.5, 5])
    def render_visualization(idx, viz):
        st.write(f'### 🌟 시각화 {idx + 1}')
        if viz:
            with st.spinner("인공지능이 시각화를 생성중입니다..."):
                try:
                    if st.session_state.selected_library == "plotly":
                        data = st.session_state.selected_dataframe
                        # extract the code from the generated responses and execute it
                        temp_namespace = {
                            'data': data,
                        }
                        exec(clean_code_snippet(viz['code']), temp_namespace)
                        fig = st.plotly_chart(temp_namespace['chart'])
                    else:
                        from PIL import Image
                        import io
                        import base64

                        imgdata = base64.b64decode(viz.raster)
                        img = Image.open(io.BytesIO(imgdata))
                        st.image(img, caption=f"Visualization {idx + 1}", use_column_width=True)
                except Exception as e:
                    st.error(f"Error loading visualization: {e}")
        with st.popover("🧑‍💻 코드 확인하기", use_container_width=True):
            if isinstance(viz, dict):
                code_string = viz['code']
            else:
                code_string = viz.code
            st.code(clean_code_snippet(code_string))
        with st.popover("🗨️ 변경 요청하기", use_container_width=True):
            chat_message = st.chat_input("(구현 중)변경하고 싶은 내용을 자연어 로 입력해주세요",key=f"chat_message_{idx}",
                                         disabled=True)
            if chat_message:
                st.session_state.visualizations = generate_visualizations(viz['code'], chat_message)
                st.session_state.update_viz = True
                st.rerun()
                render_visualization(idx, st.session_state.visualizations[idx])

        return fig

    with col1:
        idx = 0
        selected_viz = st.session_state.visualizations[idx]
        render_visualization(idx, selected_viz)

    with col2:
        st.empty()
        
    with col3:
        idx = 1
        selected_viz = visualizations[idx]
        render_visualization(idx, selected_viz)