import os
import json
import numpy as np
import requests

import streamlit as st
import streamlit_nested_layout
from streamlit_ace import st_ace
from utils import *
# from io import StringIO 
# import base64
# from pathlib import Path

import warnings
warnings.simplefilter("ignore", UserWarning)

API = 'http://4.193.50.237:5000/api'

@st.experimental_memo
def load_session():
    return requests.Session()


def translate_panel(seed, length):
    with st.container():
        st.header("Code Translation")
            
        # Expand description panel 
        with st.expander('_:pencil: Quickly select demo example below_'):
            option = st.selectbox(
                '',
                ('Select Example', 'Example 1', 'Example 2', 'Example 3'),
                key='trans'
            )
            example_code = ""
            language = "python"
            index=0
            uploaded_file = st.file_uploader("or Upload your file: ", key='trans_file')
            if uploaded_file:
                # with open(uploaded_file, encoding='utf8') as f:
                example_code = uploaded_file.read().decode()

            if option == 'Example 1':
                example_code = TRANS_EXAMPLE1
                index=1

            elif option == "Example 2":
                example_code = TRANS_EXAMPLE2
                language = "java"
                index=0

            elif option == "Example 3":
                example_code = TRANS_EXAMPLE3
                language = "java"
                index=0

        if example_code  == '':
            content = st_ace('', placeholder="Input your code in HERE", language='python', auto_update=True, key="<trans-code>")
        else:
            content = st_ace(example_code, placeholder="Input your code in HERE", language=language, auto_update=True)
        
        target_lgs = st.selectbox(
            '*Select your target translate languages:*',
            ('Python', 'Java'),
            index=index
        )
        
        if st.button("**:repeat: Translate Code**"):
            with st.spinner('Translating...'):
                source_lgs = detect_lang(API, content)
                
                if target_lgs == source_lgs:
                    st.warning('Selected language can not be equal to source language', icon="⚠️")
                    trans_code = content
                else:
                    trans_code = request_api(
                    url=API, 
                    data={
                        'input': content,
                        'task': 'translate',
                        'source_lgs': source_lgs,
                        'target_lgs': target_lgs,
                        'seed': seed,
                        'length': length
                    },
                )
                    trans_code = json.loads(trans_code.text)
                    trans_code = trans_code['output']
                    
            st.write("Translated Code")
            st.code(trans_code, language=target_lgs.lower())
        

def complete_panel(seed, length):
    with st.container():
        st.header("Code Completion")
            
        # Expand description panel 
        with st.expander('_:pencil: Quickly select demo example below_'):
            option = st.selectbox(
                '',
                ('Select Example', 'Example 1', 'Example 2', 'Example 3'),
                key='complete'
            )
            example_code = ""

            if option == 'Example 1':
                example_code = COM_EXAMPLE1

            if option == "Example 2":
                example_code = COM_EXAMPLE2

            if option == "Example 3":
                example_code = COM_EXAMPLE3

        if example_code  == '':
            content = st_ace('', placeholder="Input your code in HERE", language='python', auto_update=True, key="<complete-code>")
        else:
            content = st_ace(example_code, placeholder="Input your code in HERE", language='python', auto_update=True)
        
        if st.button("**:writing_hand: Complete Code**"):
            with st.spinner('Completing ...'):
                source_lgs = detect_lang(API, content)
                trans_code = request_api(
                    url=API, 
                    data={
                        'input': content,
                        'task': 'completion',
                        'source_lgs': source_lgs,
                        'target_lgs': None,
                        'seed': seed,
                        'length': length
                    },
                )
                trans_code = json.loads(trans_code.text)
                    
            st.write("Completed Code")
            code = str(content).strip() + trans_code['output']
            st.code(code, language=source_lgs.lower())


def repair_panel(seed, length):
    with st.container():
        st.header("Code Repair")
            
        # Expand description panel 
        with st.expander('_:pencil: Quickly select demo example below_'):
            option = st.selectbox(
                '',
                ('Select Example', 'Example 1', 'Example 2', 'Example 3'),
                key='<example-repair>'
            )
            example_code = ""

            if option == 'Example 1':
                example_code = RE_EXAMPLE1

            if option == "Example 2":
                example_code = RE_EXAMPLE2

            if option == "Example 3":
                example_code = RE_EXAMPLE3

        if example_code  == '':
            content = st_ace('', placeholder="Input your code in HERE", language='python', auto_update=True, key="<repair-code>")
        else:
            content = st_ace(example_code, placeholder="Input your code in HERE", language='python', auto_update=True)
        
        if st.button("**:hammer_and_wrench: Repair Code**"):
            with st.spinner('Repairing . . .'):
                source_lgs = detect_lang(API, content)
                trans_code = request_api(
                    url=API, 
                    data={
                        'input': content,
                        'task': 'code_repair',
                        'source_lgs': source_lgs,
                        'target_lgs': None,
                        'seed': seed,
                        'length': length
                    },
                )
                
                trans_code = json.loads(trans_code.text)
                    
            st.write("Repaired Code")
            code = trans_code['output']
            st.code(code, language=source_lgs.lower())


def codegen_panel(seed, length):
    with st.container():
        st.header("Code Generation")
            
        # Expand description panel 
        with st.expander('_:pencil: Quickly select demo example below_'):
            option = st.selectbox(
                '',
                ('Select Example', 'Example 1', 'Example 2', 'Example 3'),
                key='<example-code-gen>'
            )
            example_code = ""
            uploaded_file = st.file_uploader("or Upload your file: ", key='gen_file')
            if uploaded_file:
                example_code = uploaded_file.read().decode()

            if option == 'Example 1':
                example_code = CGEN_EXAMPLE1

            if option == "Example 2":
                example_code = CGEN_EXAMPLE2

            if option == "Example 3":
                example_code = CGEN_EXAMPLE3

        if example_code  == '':
            content = st_ace('', placeholder="Input your code in HERE", auto_update=True, key="<gen-code-code>")
        else:
            content = st_ace(example_code, placeholder="Input your code in HERE", auto_update=True)
        
        target_lgs = st.selectbox(
            '*Select your target translate languages:*',
            ('Python', 'Java'),
            index=0,
            key="gen_select"
        )

        if st.button("**:hammer: Generate Code**"):
            with st.spinner('Generating . . .'):
                trans_code = request_api(
                    url=API, 
                    data={
                        'input': content,
                        'task': 'text2code',
                        'source_lgs': target_lgs.lower(),
                        'target_lgs': None,
                        'seed': seed,
                        'length': length
                    },
                )
                
                trans_code = json.loads(trans_code.text)
                    
            st.write("Generated Code")
            code = trans_code['output']
            st.code(code, language=target_lgs.lower())
    

def codesum_panel(seed, length):
    with st.container():
        st.header("Code Summarization")
            
        # Expand description panel 
        with st.expander('_:pencil: Quickly select demo example below_'):
            option = st.selectbox(
                '',
                ('Select Example', 'Example 1', 'Example 2', 'Example 3'),
                key='<example-sum>'
            )
            example_code = ""
            uploaded_file = st.file_uploader("or Upload your file: ", key='sum_file')
            if uploaded_file:
                # with open(uploaded_file, encoding='utf8') as f:
                example_code = uploaded_file.read().decode()

            if option == 'Example 1':
                example_code = SUM_EXAMPLE1

            if option == "Example 2":
                example_code = SUM_EXAMPLE2

            if option == "Example 3":
                example_code = SUM_EXAMPLE3

        if example_code  == '':
            content = st_ace('', placeholder="Input your code in HERE", language='python', auto_update=True, key="<sum-code>")
        else:
            content = st_ace(example_code, placeholder="Input your code in HERE", language='python', auto_update=True)
        
        task = 'code2text'
        with st.expander('Advanced Option'):
            param = st.checkbox('Summarize with param')
            if param:
                task += '_param'
        if st.button("**:hammer_and_wrench: Summarize Code**"):

            with st.spinner('Summarizing . . .'):
                source_lgs = detect_lang(API, content)
                trans_code = request_api(
                    url=API, 
                    data={
                        'input': content,
                        'task': task,
                        'source_lgs': source_lgs,
                        'target_lgs': None,
                        'seed': seed,
                        'length': length
                    },
                )
                trans_code = json.loads(trans_code.text)
                    
            st.write("Summarized Code")
            code = trans_code['output']
            st.code(code, language=source_lgs.lower())


def testgen_panel(seed, length):
    with st.container():
        st.header("Test Generation")
            
        # Expand description panel 
        with st.expander('_:pencil: Quickly select demo example below_'):
            option = st.selectbox(
                '',
                ('Select Example', 'Example 1', 'Example 2', 'Example 3'),
                key='<example-test>'
            )
            example_code = ""

            if option == 'Example 1':
                example_code = TEST_EXAMPLE1

            if option == "Example 2":
                example_code = TEST_EXAMPLE2

            if option == "Example 3":
                example_code = TEST_EXAMPLE3
        
        st.warning('Feature only support Java', icon="⚠️")
        
        if example_code  == '':
            content = st_ace('', placeholder="Input your code in HERE", language='python', auto_update=True, key="<test-code>")
        else:
            content = st_ace(example_code, placeholder="Input your code in HERE", language='python', auto_update=True)
        
        if st.button("**:hammer_and_wrench: Generate Test**"):
            with st.spinner('Generating . . .'):
                source_lgs = detect_lang(API, content)
                trans_code = request_api(
                    url=API, 
                    data={
                        'input': content,
                        'task': 'testgen',
                        'source_lgs': source_lgs,
                        'target_lgs': None,
                        'seed': seed,
                        'length': length
                    },
                )
                print(trans_code.text)
                trans_code = json.loads(trans_code.text)
                    
            st.write("Generated Test")
            code = trans_code['output']
            st.code(code, language=source_lgs.lower())


def member_infor(image_path, INFOR):
    col_1, col_2, col_3, col_4 = st.columns([1,1,2,6])
    avatar_image = Image.open(image_path)
    with col_2:
            st.image(avatar_image, width=100)
    with col_4:
        st.markdown(INFOR, unsafe_allow_html=True)

def logo(image_path):
    logo = load_image_from_local(image_path)
    st.image(logo, width=100)

def main_app():
    st.set_page_config(
        page_title="Docify-Lab",
        page_icon="🧊",
        layout="wide",
        # initial_sidebar_state="expanded"
    )
    
    col_title1, col_title2 = st.columns([1.5, 20])
    avatar_image = Image.open('./assets/ava_app.png')
    # st.markdown("<p style='text-align: center; color: grey;'>"+img_to_html('./assets/ava_app.png')+"</p>", unsafe_allow_html=True)

    with col_title1:
        st.image(avatar_image, width=70)
        # st.markdown("<p style='text-align: center; color: grey; width: 64px;'>"+img_to_html('./assets/ava_app.png')+"</p>", unsafe_allow_html=True)

    with col_title2:
        st.title("                 Welcome to Docify-Lab")
    
    # layout
    col1, col2 = st.columns([6, 4])
    # sess = load_session()
    
    # Advanced option
    st.sidebar.write("Advanced Custom")
    length = st.sidebar.slider('Generate length', 16, 1024, 512, 32)
    seed = st.sidebar.number_input('Model seed:', value=42, min_value=1, step=1)
    
    with col1:

        tab1, tab2, tab3, tab4, tab5, tab6 = st.tabs(
            ["Code Translation", "Test Generation", "Code Repair", 
            "Code Summarization", "Code Generation", "Code Completion"]
        )
        
        with tab1:
            translate_panel(seed, length)
        with tab2:
            testgen_panel(seed, length)
        with tab3:
            repair_panel(seed, length)
        with tab4:
            codesum_panel(seed, length)
        with tab5:
            codegen_panel(seed, length)
        with tab6:
            complete_panel(seed, length)
    
    with col2:
        c1, c2, c3, c4, c5, c6, c7 = st.columns([1,1,1,1,1,1,1])

        with c2:
            logo("./assets/fpt_logo.png")
        with c4:
            logo("./assets/AI_center_logo.png")
        with c6:
            logo("./assets/Docify-Logo.png")
        # st.markdown(SIDEBAR_INFO, unsafe_allow_html=True)


        with st.expander("What is Docify-Lab?", expanded=True):
            st.markdown(STORY, unsafe_allow_html=True)
        
        with st.expander("Docify Lab - FPT Software AI Center", expanded=True):
            tab_t1, tab_1, tab_t2, tab_t3, tab_t4, tab_t5, tab_t6 = st.tabs(
                ["About Us", "Leader","Member 1", "Member 2", 
                "Member 3", "Member 4", "Member 5"],
            )

            with tab_t1:
                st.markdown(GROUP_INFO, unsafe_allow_html=True)
            with tab_1:
                member_infor('./assets/Anh_Nghi_ava.png', LEADER_INFOR)
            with tab_t2:
                member_infor('./assets/Nam_ava.png', NAM_INFOR)
            with tab_t3:
                member_infor('./assets/Dung_ava.jpg', DUNG_INFOR)
            with tab_t4:
                member_infor('./assets/Minh_ava.JPG', MINH_INFOR)
            with tab_t5:
                member_infor('./assets/VA_ava.jpg', VA_INFOR)
            with tab_t6:
                member_infor('./assets/KA_ava.jpg', KA_INFOR)
    
if __name__ == '__main__':
    main_app()