import os
import json
import uuid
import pandas as pd
import gradio as gr
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Import the LangGraph agent
from agent import agent

# ---------------------------------------------------------
# 1. UI STATE & RENDERING HELPERS (Zero Business Logic)
# ---------------------------------------------------------

def get_progress_html():
    """Builds the HTML progress bar by checking for local checkpoint files."""
    phases = {
        "Phase 1: Familiarisation": "summaries.json",
        "Phase 2: Initial Codes": "labels.json",
        "Phase 3: Themes": "themes.json",
        "Phase 5.5: PAJAIS": "taxonomy_map.json",
        "Phase 6: Report": "comparison.csv"
    }

    html = "<div style='display: flex; justify-content: space-between; padding: 15px; background: #2b2b2b; color: white; border-radius: 8px; font-family: sans-serif;'>"
    for name, file in phases.items():
        status = "✅" if os.path.exists(file) else "⬜"
        html += f"<span>{status} <b>{name}</b></span>"
    html += "</div>"
    return html

def load_review_table():
    """Loads the highest priority JSON into the Review Table format based on Phase progress."""
    columns = ["#", "Topic Label", "Top Evidence", "Sentences", "Papers", "Approve", "Rename To", "Reasoning"]
    empty_df = pd.DataFrame(columns=columns)

    try:
        # Priority 1: Phase 5.5 Taxonomy Map
        if os.path.exists("taxonomy_map.json"):
            with open("taxonomy_map.json", "r") as f: data = json.load(f)
            rows = [[i, k, f"→ {v.get('pajais_match', 'NOVEL')} | {v.get('reasoning', '')}", "", "", "", "", ""] for i, (k, v) in enumerate(data.items())]
            return pd.DataFrame(rows, columns=columns)

        # Priority 2: Phase 3 Themes
        if os.path.exists("themes.json"):
            with open("themes.json", "r") as f: data = json.load(f)
            rows = [[i, k, " | ".join(v.get("top_sentences", [])), v.get("size", ""), v.get("papers_count", ""), "", "", ""] for i, (k, v) in enumerate(data.items())]
            return pd.DataFrame(rows, columns=columns)

        # Priority 3: Phase 2 Labels
        if os.path.exists("labels.json"):
            with open("labels.json", "r") as f: data = json.load(f)
            rows = [[i, v.get("label", "Unknown"), f"Category: {v.get('category', '')} | {v.get('reasoning', '')}", "", "", "", "", ""] for i, (k, v) in enumerate(data.items())]
            return pd.DataFrame(rows, columns=columns)

        # Priority 4: Phase 2 Raw Summaries
        if os.path.exists("summaries.json"):
            with open("summaries.json", "r") as f: data = json.load(f)
            rows = [[i, f"Topic {k}", " | ".join(v.get("top_sentences", [])), v.get("size", ""), v.get("papers_count", ""), "", "", ""] for i, (k, v) in enumerate(data.items())]
            return pd.DataFrame(rows, columns=columns)

    except Exception:
        pass # UI helpers shouldn't crash the app if a JSON is malformed while writing

    return empty_df

def get_available_downloads():
    """Returns a list of all current checkpoint files for the Download tab."""
    target_files = ["processed_data.json", "summaries.json", "emb.npy", "charts.html", "labels.json", "themes.json", "taxonomy_map.json", "comparison.csv", "narrative.txt"]
    return [f for f in target_files if os.path.exists(f)]

def load_chart_view(chart_name):
    """Loads a mock view of the requested chart (assumes charts.html holds them)."""
    if os.path.exists("charts.html"):
        with open("charts.html", "r") as f:
            return f.read() # In a real app, you'd parse out the specific div
    return "<div style='padding: 20px;'>Charts not yet generated. Complete Phase 2.</div>"

# ---------------------------------------------------------
# 2. AGENT INTERACTION HANDLERS
# ---------------------------------------------------------

def interact_with_agent(user_message, chat_history, thread_id):
    """Sends a message to LangGraph and extracts the AI response."""

    # 1. First, append the user's message as a dictionary
    chat_history.append({"role": "user", "content": user_message})

    # 2. Next, append the assistant's thinking state as a dictionary
    chat_history.append({"role": "assistant", "content": "⏳ Thinking..."})

    yield chat_history, get_progress_html(), load_review_table(), get_available_downloads()

    config = {"configurable": {"thread_id": thread_id}}

    try:
        # Invoke the LangGraph agent
        result = agent.invoke({"messages": [("user", user_message)]}, config=config)

        # Extract the last AI message
        ai_response = result["messages"][-1].content

        # 3. Update the last assistant message with the real response
        chat_history[-1] = {"role": "assistant", "content": ai_response}

    except Exception as e:
        # 4. Handle errors using the dictionary format as well
        chat_history[-1] = {"role": "assistant", "content": f"❌ Error communicating with agent: {str(e)}"}

    # Return updated states
    yield chat_history, get_progress_html(), load_review_table(), get_available_downloads()

def handle_csv_upload(file_obj, chat_history, thread_id):
    if file_obj is None:
        return chat_history, get_progress_html(), load_review_table(), get_available_downloads()
    # Save the uploaded file as 'Scopus.csv'
    import shutil
    shutil.copy(file_obj.name, "Scopus.csv")
    # Now trigger the agent
    yield from interact_with_agent("Analyze my Scopus CSV", chat_history, thread_id)

def handle_submit_review(df, chat_history, thread_id):
    """Converts the Gradio Dataframe edits into a string and sends to the agent."""
    # Filter only rows where the user provided input in the editable columns
    review_data = df[df["Approve"].astype(str).str.strip() != ""].to_dict(orient="records")

    if not review_data:
        msg = "I submitted the table, but I didn't make any edits."
    else:
        msg = f"Here is my submitted review table data:\n{json.dumps(review_data, indent=2)}"

    yield from interact_with_agent(msg, chat_history, thread_id)


# ---------------------------------------------------------
# 3. GRADIO UI LAYOUT
# ---------------------------------------------------------

with gr.Blocks(title="B&C Thematic Analysis Agent", theme=gr.themes.Soft()) as demo:
    # State object to keep track of the LangGraph memory thread for this session
    session_thread_id = gr.State(value=lambda: str(uuid.uuid4()))

    gr.Markdown("# 🧠 Braun & Clarke (2006) Thematic Analysis AI Agent")

    # --- PHASE PROGRESS PIPELINE ---
    progress_bar = gr.HTML(value=get_progress_html())

    with gr.Row():
        # --- SECTION 1 & 2: INPUT AND CONVERSATION ---
        with gr.Column(scale=1, variant="panel"):
            gr.Markdown("### ① DATA INPUT")
            csv_upload = gr.File(label="Upload Scopus CSV", file_types=[".csv"])

            gr.Markdown("### ② AGENT CONVERSATION")
            chatbot = gr.Chatbot(label="Agent Dialogue", height=400)
            with gr.Row():
                user_input = gr.Textbox(show_label=False, placeholder="Type your message...", scale=4)
                send_btn = gr.Button("Send", variant="primary", scale=1)

        # --- SECTION 3: RESULTS ---
        with gr.Column(scale=2, variant="panel"):
            gr.Markdown("### ③ RESULTS")
            with gr.Tabs():

                # TAB A: Review Table
                with gr.TabItem("Review Table"):
                    gr.Markdown("*Edit the 'Approve', 'Rename To', and 'Reasoning' columns to guide the agent.*")
                    review_table = gr.Dataframe(
                        value=load_review_table(),
                        headers=["#", "Topic Label", "Top Evidence", "Sentences", "Papers", "Approve", "Rename To", "Reasoning"],
                        datatype=["number", "str", "str", "str", "str", "str", "str", "str"],
                        interactive=True,
                        wrap=True
                    )
                    submit_review_btn = gr.Button("Submit Review", variant="primary")

                # TAB B: Charts
                with gr.TabItem("Charts"):
                    chart_dropdown = gr.Dropdown(choices=["Intertopic Map", "Bar Chart", "Hierarchy", "Heatmap"], label="Select Chart", value="Bar Chart")
                    chart_html = gr.HTML(value=load_chart_view("Bar Chart"))

                # TAB C: Download
                with gr.TabItem("Download"):
                    download_files = gr.File(label="Checkpoint Files", file_count="multiple", value=get_available_downloads())

    # ---------------------------------------------------------
    # 4. EVENT WIRING
    # ---------------------------------------------------------

    # Common outputs to update across all actions
    update_outputs = [chatbot, progress_bar, review_table, download_files]

    # File Upload Trigger
    csv_upload.upload(
        handle_csv_upload,
        inputs=[csv_upload, chatbot, session_thread_id],
        outputs=update_outputs
    )

    # Chat Send Trigger
    send_btn.click(
        interact_with_agent,
        inputs=[user_input, chatbot, session_thread_id],
        outputs=update_outputs
    ).then(lambda: "", None, user_input) # Clear textbox

    user_input.submit(
        interact_with_agent,
        inputs=[user_input, chatbot, session_thread_id],
        outputs=update_outputs
    ).then(lambda: "", None, user_input)

    # Submit Review Trigger
    submit_review_btn.click(
        handle_submit_review,
        inputs=[review_table, chatbot, session_thread_id],
        outputs=update_outputs
    )

    # Chart Dropdown Trigger
    chart_dropdown.change(
        load_chart_view,
        inputs=[chart_dropdown],
        outputs=[chart_html]
    )

if __name__ == "__main__":
    demo.launch()