import gradio as gr
import utils
from app import question_graph, MemorySaver
from translation import LANGUAGES
import uuid
import os
from dotenv import load_dotenv
import tempfile

load_dotenv()

service_account_json = os.environ.get("GOOGLE_SERVICE_ACCOUNT_JSON")
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")

if service_account_json:
    # Write the JSON content to a temporary file
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json")
    temp_file.write(service_account_json.encode("utf-8"))
    temp_file.close()
    # Set the GOOGLE_APPLICATION_CREDENTIALS environment variable to the temp file path
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = temp_file.name

# -----------------------------
# In-Memory Storage for Feedback Flow
# -----------------------------
session_data = {}
session_data["en_topics"] = utils.get_topic_names("en")
session_data["az_topics"] = utils.get_topic_names("az")
session_data["ru_topics"] = utils.get_topic_names("ru")


# -----------------------------
# Graph Invocation Function
# -----------------------------
def run_graph(topic, subtopic, difficulty, lang):
    session_id = str(uuid.uuid4())  # Unique session per question
    memory = MemorySaver()

    if lang == "az" or lang == "ru":
        topic, subtopic, difficulty = utils.get_eng_names(
            topic, subtopic, difficulty, lang
        )

    print("Topic:", topic, "Subtopic:", subtopic, "Difficulty:", difficulty)

    inputs = {
        "topic": topic,
        "subtopic": subtopic,
        "difficulty": difficulty,
        "description": "",
        "num_questions": 3,
        "messages": [],
        "context": [],
        "relevant_questions": [],
        "question": "",
        "steps": [],
        "tool_requests": [],
        "tool_results": [],
        "verified": False,
        "solution": "",
        "answer": "",
        "error": "",
        "human_feedback": "",
    }

    # Run up to feedback node (interrupt point)
    partial_result = question_graph.invoke(
        inputs, config={"configurable": {"thread_id": session_id}}
    )
    session_data[session_id] = {
        "memory": memory,
        "inputs": partial_result,  # Store intermediate state
    }

    generated_question = partial_result.get("question", "No question generated.")

    # Translate Question if English is not selected
    if lang != "en":
        generated_question = utils.translate_text(
            generated_question, PROJECT_ID, target_lang=lang, source_lang="en"
        )

    return session_id, generated_question, LANGUAGES[lang]["feedback_placeholder"]


# -----------------------------
# Handle Feedback and Rerun Graph
# -----------------------------
def handle_feedback(session_id, feedback_text, lang):
    data = session_data.get(session_id)
    if not data:
        return session_id, "Session expired or invalid.", ""

    memory = data["memory"]
    feedback_text = feedback_text.strip()
    thread = {"configurable": {"thread_id": session_id}}

    if feedback_text:
        # Update state with feedback at node 'feedback'
        feedback_text = utils.translate_text(
            feedback_text, PROJECT_ID, target_lang="en"
        )

        question_graph.update_state(
            thread, {"human_feedback": feedback_text}, as_node="feedback"
        )

        # Continue from feedback
        final_result = question_graph.invoke(
            None, config={"configurable": {"thread_id": session_id}}
        )

        new_question = final_result.get("question", "No question generated.")
        session_data[session_id]["inputs"] = final_result

        if final_result.get("steps"):
            return session_id, new_question, "Proceeding to solve question..."
        else:
            return (
                session_id,
                new_question,
                "Feedback applied. Refine again or click 'Solve Question'.",
            )
    else:
        # No feedback provided, proceed immediately
        question_graph.update_state(
            thread,
            {"human_feedback": None},
            as_node="feedback",
        )

        final_result = question_graph.invoke(
            None, config={"configurable": {"thread_id": session_id}}
        )

        solution_md = final_result.get("solution", "No solution generated.")
        answer = final_result.get("answer", "No answer.")

        # Translate solution and answer if lang is not English
        if lang != "en":
            solution_md = utils.translate_text(
                solution_md, PROJECT_ID, target_lang=lang, source_lang="en"
            )
            answer = utils.translate_text(
                answer, PROJECT_ID, target_lang=lang, source_lang="en"
            )

        return session_id, solution_md, answer


# -----------------------------
# Solve Question – Full Graph Run
# -----------------------------
def solve_question(session_id, lang):
    data = session_data.get(session_id)
    thread = {"configurable": {"thread_id": session_id}}
    if not data:
        return "Session expired or invalid.", "", ""

    state = data["inputs"]
    memory = data["memory"]

    question_graph.update_state(
        thread,
        {"human_feedback": None},
        as_node="feedback",
    )

    final_result = question_graph.invoke(
        None, config={"configurable": {"thread_id": session_id}}
    )

    print("Final Result", final_result)
    solution_md = final_result.get("solution", "No solution generated.")
    answer = final_result.get("answer", "No answer.")

    # Translate solution and answer if lang is not English
    if lang != "en":
        solution_md = utils.translate_text(
            solution_md, PROJECT_ID, target_lang=lang, source_lang="en"
        )
        answer = utils.translate_text(
            answer, PROJECT_ID, target_lang=lang, source_lang="en"
        )

    return solution_md


# Language toggle updates UI
def update_ui(lang):
    return (
        LANGUAGES[lang]["title"],  # Markdown title
        gr.update(visible=(lang == "az")),  # Show Azerbaijani UI
        gr.update(visible=(lang == "en")),  # Show English UI
        gr.update(visible=(lang == "ru")),  # Show Russian UI
    )


def update_subtopics(selected_topic, lang):
    subtopics_list = utils.get_subtopics(selected_topic, lang)
    return gr.update(
        choices=subtopics_list, value=subtopics_list[0] if subtopics_list else ""
    )


# -----------------------------
# Gradio UI Layout
# -----------------------------
with gr.Blocks(title="LangGraph Math Solver") as demo:
    lang_toggle = gr.Dropdown(
        label="🌍 Select Language", choices=["en", "az", "ru"], value="en"
    )

    title = gr.Markdown(LANGUAGES["en"]["title"])

    # English UI Components
    with gr.Column(visible=True) as english_ui:
        with gr.Row():
            topic_input_en = gr.Dropdown(
                label=LANGUAGES["en"]["topic_label"],
                choices=session_data["en_topics"],
                value=None,
            )
            subtopic_input_en = gr.Dropdown(
                label=LANGUAGES["en"]["subtopic_label"], choices=[], value=None
            )
            difficulty_input_en = gr.Dropdown(
                label=LANGUAGES["en"]["difficulty_label"],
                choices=["easy", "medium", "hard"],
                value="medium",
            )

        generate_btn_en = gr.Button(LANGUAGES["en"]["generate_btn"])
        question_output_en = gr.Markdown(
            label=LANGUAGES["en"]["question_output"],
            render=True,
            latex_delimiters=[
                {"left": "$$", "right": "$$", "display": True},
                {"left": "$", "right": "$", "display": False},
            ],
        )
        feedback_note_en = gr.Textbox(
            label=LANGUAGES["en"]["feedback_label"],
            placeholder=LANGUAGES["en"]["feedback_placeholder"],
        )

        with gr.Row():
            feedback_btn_en = gr.Button(LANGUAGES["en"]["feedback_btn"])
            solve_btn_en = gr.Button(LANGUAGES["en"]["solve_btn"])

        solution_output_en = gr.Markdown(
            label=LANGUAGES["en"]["solution_output"],
            render=True,
            latex_delimiters=[
                {"left": "$$", "right": "$$", "display": True},
                {"left": "$", "right": "$", "display": False},
                {"left": "[", "right": "]", "display": False},
            ],
        )
        # final_answer_en = gr.Markdown(
        #     label=LANGUAGES["en"]["final_answer"], render=True
        # )

    # Azerbaijani UI Components (Initially Hidden)
    with gr.Column(visible=False) as azeri_ui:
        with gr.Row():
            topic_input_az = gr.Dropdown(
                label=LANGUAGES["az"]["topic_label"],
                choices=session_data["az_topics"],
                value=None,
            )
            subtopic_input_az = gr.Dropdown(
                label=LANGUAGES["az"]["subtopic_label"], choices=[], value=None
            )
            difficulty_input_az = gr.Dropdown(
                label=LANGUAGES["az"]["difficulty_label"],
                choices=["asan", "orta", "çətin"],
                value="orta",
            )

        generate_btn_az = gr.Button(LANGUAGES["az"]["generate_btn"])
        question_output_az = gr.Markdown(
            label=LANGUAGES["az"]["question_output"],
            render=True,
            latex_delimiters=[
                {"left": "$$", "right": "$$", "display": True},
                {"left": "$", "right": "$", "display": False},
            ],
        )
        feedback_note_az = gr.Textbox(
            label=LANGUAGES["az"]["feedback_label"],
            placeholder=LANGUAGES["az"]["feedback_placeholder"],
        )

        with gr.Row():
            feedback_btn_az = gr.Button(LANGUAGES["az"]["feedback_btn"])
            solve_btn_az = gr.Button(LANGUAGES["az"]["solve_btn"])

        solution_output_az = gr.Markdown(
            label=LANGUAGES["az"]["solution_output"],
            render=True,
            latex_delimiters=[
                {"left": "$$", "right": "$$", "display": True},
                {"left": "$", "right": "$", "display": False},
                {"left": "[", "right": "]", "display": False},
            ],
        )
        # final_answer_az = gr.Markdown(
        #     label=LANGUAGES["az"]["final_answer"], render=True
        # )

    # Russian UI Components (Initially Hidden)
    with gr.Column(visible=False) as russian_ui:
        with gr.Row():
            topic_input_ru = gr.Dropdown(
                label=LANGUAGES["ru"]["topic_label"],
                choices=session_data["ru_topics"],
                value=None,
            )
            subtopic_input_ru = gr.Dropdown(
                label=LANGUAGES["ru"]["subtopic_label"], choices=[], value=None
            )
            difficulty_input_ru = gr.Dropdown(
                label=LANGUAGES["ru"]["difficulty_label"],
                choices=["легкий", "средний", "сложный"],
                value="средний",
            )

        generate_btn_ru = gr.Button(LANGUAGES["ru"]["generate_btn"])
        question_output_ru = gr.Markdown(
            label=LANGUAGES["ru"]["question_output"],
            render=True,
            latex_delimiters=[
                {"left": "$$", "right": "$$", "display": True},
                {"left": "$", "right": "$", "display": False},
            ],
        )
        feedback_note_ru = gr.Textbox(
            label=LANGUAGES["ru"]["feedback_label"],
            placeholder=LANGUAGES["ru"]["feedback_placeholder"],
        )

        with gr.Row():
            feedback_btn_ru = gr.Button(LANGUAGES["ru"]["feedback_btn"])
            solve_btn_ru = gr.Button(LANGUAGES["ru"]["solve_btn"])

        solution_output_ru = gr.Markdown(
            label=LANGUAGES["ru"]["solution_output"],
            render=True,
            latex_delimiters=[
                {"left": "$$", "right": "$$", "display": True},
                {"left": "$", "right": "$", "display": False},
                {"left": "[", "right": "]", "display": False},
            ],
        )
        # final_answer_ru = gr.Markdown(
        #     label=LANGUAGES["ru"]["final_answer"], render=True
        # )

    # Language Toggle
    lang_toggle.change(
        update_ui,
        inputs=[lang_toggle],
        outputs=[
            title,
            azeri_ui,
            english_ui,
            russian_ui,
        ],  # Using .update() to toggle visibility
    )

    # Topic Selection Events
    topic_input_en.change(
        update_subtopics,
        inputs=[topic_input_en, lang_toggle],
        outputs=[subtopic_input_en],
    )

    topic_input_az.change(
        update_subtopics,
        inputs=[topic_input_az, lang_toggle],
        outputs=[subtopic_input_az],
    )

    topic_input_ru.change(
        update_subtopics,
        inputs=[topic_input_ru, lang_toggle],
        outputs=[subtopic_input_ru],
    )

    # Hidden session ID
    session_state = gr.State()

    # Events for English UI
    generate_btn_en.click(
        run_graph,
        inputs=[topic_input_en, subtopic_input_en, difficulty_input_en, lang_toggle],
        outputs=[session_state, question_output_en, feedback_note_en],
    )

    feedback_btn_en.click(
        handle_feedback,
        inputs=[session_state, feedback_note_en, lang_toggle],
        outputs=[session_state, question_output_en, feedback_note_en],
    )

    solve_btn_en.click(
        solve_question,
        inputs=[session_state, lang_toggle],
        outputs=[solution_output_en],
    )

    # Events for Azerbaijani UI
    generate_btn_az.click(
        run_graph,
        inputs=[topic_input_az, subtopic_input_az, difficulty_input_az, lang_toggle],
        outputs=[session_state, question_output_az, feedback_note_az],
    )

    feedback_btn_az.click(
        handle_feedback,
        inputs=[session_state, feedback_note_az, lang_toggle],
        outputs=[session_state, question_output_az, feedback_note_az],
    )

    solve_btn_az.click(
        solve_question,
        inputs=[session_state, lang_toggle],
        outputs=[solution_output_az],
    )

    # Events for Russian UI
    generate_btn_ru.click(
        run_graph,
        inputs=[topic_input_ru, subtopic_input_ru, difficulty_input_ru, lang_toggle],
        outputs=[session_state, question_output_ru, feedback_note_ru],
    )

    feedback_btn_ru.click(
        handle_feedback,
        inputs=[session_state, feedback_note_ru, lang_toggle],
        outputs=[session_state, question_output_ru, feedback_note_ru],
    )

    solve_btn_ru.click(
        solve_question,
        inputs=[session_state, lang_toggle],
        outputs=[solution_output_ru],
    )


# -----------------------------
# Launch App
# -----------------------------
if __name__ == "__main__":
    demo.launch()