import json
import logging
import multiprocessing
import os

import gradio as gr

from swiftsage.agents import SwiftSage
from swiftsage.utils.commons import PromptTemplate, api_configs, setup_logging
from pkg_resources import resource_filename

#ENGINE = "Together"
#SWIFT_MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo"
#FEEDBACK_MODEL_ID = "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"
#SAGE_MODEL_ID = "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo"
 

# ENGINE = "Groq"
# SWIFT_MODEL_ID = "llama-3.1-8b-instant"
# FEEDBACK_MODEL_ID = "llama-3.1-8b-instant"
# SAGE_MODEL_ID = "llama-3.1-70b-versatile"

ENGINE = "SambaNova"
SWIFT_MODEL_ID = "Meta-Llama-3.1-8B-Instruct"
FEEDBACK_MODEL_ID = "Meta-Llama-3.1-70B-Instruct"
SAGE_MODEL_ID = "Meta-Llama-3.1-405B-Instruct"

def solve_problem(problem, max_iterations, reward_threshold, swift_model_id, sage_model_id, feedback_model_id, use_retrieval, start_with_sage, swift_temperature, swift_top_p, sage_temperature, sage_top_p, feedback_temperature, feedback_top_p):
    global ENGINE
    # Configuration for each LLM
    max_iterations = int(max_iterations)
    reward_threshold = int(reward_threshold)

    swift_config = {
        "model_id": swift_model_id,
        "api_config": api_configs[ENGINE],
        "temperature": float(swift_temperature),
        "top_p": float(swift_top_p),
        "max_tokens": 8192,
    }

    feedback_config = {
        "model_id": feedback_model_id,
        "api_config": api_configs[ENGINE],
        "temperature": float(feedback_temperature),
        "top_p": float(feedback_top_p),
        "max_tokens": 8192,
    }

    sage_config = {
        "model_id": sage_model_id,
        "api_config": api_configs[ENGINE],
        "temperature": float(sage_temperature),
        "top_p": float(sage_top_p),
        "max_tokens": 8192,
    }

    # specify the path to the prompt templates
    # prompt_template_dir = './swiftsage/prompt_templates'
    # prompt_template_dir = resource_filename('swiftsage', 'prompt_templates')

    # Try multiple locations for the prompt templates
    possible_paths = [
        resource_filename('swiftsage', 'prompt_templates'),
        os.path.join(os.path.dirname(__file__), '..', 'swiftsage', 'prompt_templates'),
        os.path.join(os.path.dirname(__file__), 'swiftsage', 'prompt_templates'),
        '/app/swiftsage/prompt_templates',  # For Docker environments
    ]

    prompt_template_dir = None
    for path in possible_paths:
        if os.path.exists(path):
            prompt_template_dir = path
            break

    dataset = [] 
    embeddings = [] # TODO: for retrieval augmentation (not implemented yet now)
    s2 = SwiftSage(
        dataset,
        embeddings,
        prompt_template_dir,
        swift_config,
        sage_config,
        feedback_config,
        use_retrieval=use_retrieval,
        start_with_sage=start_with_sage,
    )

    reasoning, solution, messages = s2.solve(problem, max_iterations, reward_threshold)
    reasoning = reasoning.replace("The generated code is:", "\n---\nThe generated code is:").strip()
    solution = solution.replace("Answer (from running the code):\n ", " ").strip()
    # generate HTML for the log messages and display them with wrap and a scroll bar and a max height in the code block with log style 
    
    log_messages = "<pre style='white-space: pre-wrap; max-height: 500px; overflow-y: scroll;'><code class='log'>" + "\n".join(messages) + "</code></pre>"
    return reasoning, solution, log_messages


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    # gr.Markdown("## SwiftSage: A Multi-Agent Framework for Reasoning")
    # use the html and center the title 
    gr.HTML("<h1 style='text-align: center;'>🏦 Bank Failure Predictor</h1>")
    gr.HTML("<span>This tool predicts the likelihood of bank failure based on balance sheet data.</span>")

    with gr.Row(): 
        swift_model_id = gr.Textbox(label="😄 Swift Model ID", value=SWIFT_MODEL_ID)
        feedback_model_id = gr.Textbox(label="🤔 Feedback Model ID", value=FEEDBACK_MODEL_ID)
        sage_model_id = gr.Textbox(label="😎 Sage Model ID", value=SAGE_MODEL_ID)
        # the following two should have a smaller width
    
    with gr.Accordion(label="⚙️ Advanced Options", open=False):
        with gr.Row():
            with gr.Column():
                max_iterations = gr.Textbox(label="Max Iterations", value="5")
                reward_threshold = gr.Textbox(label="feedback Threshold", value="8")
            # TODO: add top-p and temperature for each module for controlling 
            with gr.Column():
                top_p_swift = gr.Textbox(label="Top-p for Swift", value="0.9")
                temperature_swift = gr.Textbox(label="Temperature for Swift", value="0.5")
            with gr.Column():
                top_p_sage = gr.Textbox(label="Top-p for Sage", value="0.9")
                temperature_sage = gr.Textbox(label="Temperature for Sage", value="0.5")
            with gr.Column():
                top_p_feedback = gr.Textbox(label="Top-p for Feedback", value="0.9")
                temperature_feedback = gr.Textbox(label="Temperature for Feedback", value="0.5")

            use_retrieval = gr.Checkbox(label="Use Retrieval Augmentation", value=False, visible=False)
            start_with_sage = gr.Checkbox(label="Start with Sage", value=False, visible=False)

    problem = gr.Textbox(label="Input balance sheet data or parameters", value="Enter the bank's financial data here...", lines=5)

    solve_button = gr.Button("🔮 Predict Failure Chance")
    reasoning_output = gr.Textbox(label="Prediction steps with Code", interactive=False)
    solution_output = gr.Textbox(label="Prediction Result", interactive=False)

    # add a log display for showing the log messages
    with gr.Accordion(label="📜 Log Messages", open=False):
        log_output = gr.HTML("<p>No log messages yet.</p>")

    solve_button.click(
        solve_problem,
        inputs=[problem, max_iterations, reward_threshold, swift_model_id, sage_model_id, feedback_model_id, use_retrieval, start_with_sage, temperature_swift, top_p_swift, temperature_sage, top_p_sage, temperature_feedback, top_p_feedback],
        outputs=[reasoning_output, solution_output, log_output],
    )



if __name__ == '__main__':
    # make logs dir if it does not exist
    if not os.path.exists('logs'):
        os.makedirs('logs')
    multiprocessing.set_start_method('spawn')
    demo.launch(share=True, show_api=True)