from datetime import datetime
import gradio as gr
import json
import time
from legal_assitr_prototype.engine.engine import SYSTEM_PROMPT, generate, generate_user_prompt

from legal_assitr_prototype.templates.nda.nda_form import NDA_TEMPLATE_PROMPT
from legal_assitr_prototype.templates.templates_sidebar import templates_sidebar

DESCRIPTION = """
# Legal Assitr PoC ⚖️ - This is a PoC for Legal Assitr
### Legal Assist is an AI tool designed to assist users in creating legal documents by leveraging artificial intelligence (AI) technology.

`retro56/zs-writer` and `retro56/legalwriter-7b-instruct`(best) perform the best but try out different models to see how they react to the same prompt.

### Built by [Ankur Debnath](https://github.com/r3tr056) with 💙
"""

MAX_MAX_NEW_TOKENS = 10
DEFAULT_MAX_NEW_TOKENS = 10

curr_date = datetime.now().strftime("%m/%d/%Y")
shared = dict()

def main():

    def complete(system_prompt, effective_date, disclosing_party, receiving_party, confidential_info, obligations, term, breach, city):

        form_prompt = NDA_TEMPLATE_PROMPT.format(
            effective_date=effective_date,
            disclosing_party=disclosing_party,
            receiving_party=receiving_party,
            confidential_info=confidential_info,
            obligations=obligations,
            term=term,
            breach=breach,
            city=city
        )
        
        result = generate(system_prompt=system_prompt, user_prompt=form_prompt)
        result = f"""{result}"""
        return result

    with gr.Blocks(theme='soft') as demo:
        gr.Markdown(DESCRIPTION)
        with gr.Row():

            with gr.Column(scale=2):
                with gr.Blocks() as nda_form:
                    with gr.Group():
                        effective_date = gr.Textbox(
                            label="Effective Date",
                            placeholder=curr_date,
                            value=curr_date
                        )
                        city = gr.Textbox(
                            label="City"
                        )
                        purpose = gr.Textbox(
                            label="Purpose of the NDA",
                            placeholder="Enter the purpose of the NDA"
                        )
                        parties = gr.Textbox(
                            label="Parties Involved",
                            placeholder="Parties Involved..."
                        )
                        disclosing_party = gr.Textbox(
                            label="Disclosing Party",
                            placeholder="Disclosing Party with Address and Contact..."
                        )
                        receiving_party = gr.Textbox(
                            label="Receiving Party",
                            placeholder="Receiving Party with Address And Contact..."
                        )
                        obligations = gr.Textbox(
                            label="Obligations",
                            placeholder="Obligations"
                        )
                        mutual_or_not = gr.Textbox(
                            label="Mutual or One Way NDA",
                            placeholder="Whether the NDA is mutual or One-Way..."
                        )
                        confidential_info = gr.Textbox(
                            label="Confidential Information",
                            placeholder="Explain the Confidential Information"
                        )
                        term = gr.Textbox(
                            label="Term",
                            placeholder="How long the Confidentiality lasts",
                        )
                        breach = gr.Textbox(
                            label="Consequences of breach",
                            placeholder="Consequences of breach"
                        )

                with gr.Group():
                    with gr.Accordion(label="AI Settings", open=False):
                        system_prompt = gr.Textbox(
                            label='System Prompt',
                            value=SYSTEM_PROMPT,
                            lines=5
                        )

                        max_new_tokens = gr.Slider(
                            label='Max new tokens',
                            minimum=1,
                            maximum=MAX_MAX_NEW_TOKENS,
                            step=1,
                            value=DEFAULT_MAX_NEW_TOKENS,
                        )
                        temp = gr.Slider(
                            label='Temperature',
                            minimum=0.1,
                            maximum=4.0,
                            step=0.1,
                            value=1.0,
                        )
                        top_p = gr.Slider(
                            label='Top-p (nucleus sampling)',
                            minimum=0.05,
                            maximum=1.0,
                            step=0.05,
                            value=0.95,
                        )
                        top_k = gr.Slider(
                            label='Top-k',
                            minimum=1,
                            maximum=1000,
                            step=1,
                            value=50,
                        )

            with gr.Column(scale=4):
                with gr.Row():
                    with gr.Tab("Raw"):
                        raw_editor_textbox = gr.Textbox(
                            label="Generated Document...",
                            lines=35,
                            elem_id='textbox-default',
                            elem_classes=['textbox_default_output', 'add_scrollbar'],
                            show_copy_button=True
                        )
                    with gr.Tab("Prettified"):
                        md_render_btn = gr.Button('Render')
                        md_editor_box = gr.Markdown()

                    with gr.Tab("Generated Tokens"):
                        shared["tokens-default"] = gr.Textbox(
                            show_copy_button=True,
                            lines=35,
                            label="Tokens",
                            elem_classes=['textbox_tokens', 'add_scrollbar', 'monospace']
                        )
                with gr.Row():
                    generate_btn = gr.Button(value="Generate Content", variant='primary')
                    render_md = gr.Button(value="Prettify", variant='secondary')
        
        generate_btn.click(fn=complete, inputs=[
            system_prompt,
            effective_date,
            disclosing_party,
            receiving_party,
            confidential_info,
            obligations,
            term,
            breach,
            city
        ], outputs=[raw_editor_textbox])
        

    demo.launch()


if __name__ == "__main__":
    main()