import os
# import langchain_mistralai
from langchain.agents import AgentType, initialize_agent
import gradio as gr
# from mistralai import Mistral
import pandas as pd
from collections import Counter
from openai import AzureOpenAI
import time
import re

# Memory per persona
structured_memory = {
    "👨‍💼 Scrum Master": {},
    "💡 Product Owner": {},
    "👨‍💻 Developer": {},
    "🎨 Designer": {},
}
selected_items = []

def prepare_model(selected_model):
    if selected_model == 'o4-mini':
        client = AzureOpenAI(
                api_version="2024-12-01-preview",
                azure_endpoint="https://nextgenagilehu4401821853.openai.azure.com/",
                api_key="2ziHxIi2Pz511IscFVVjBpgHkm2nDXNfPAvXyFWpsmCHlpHwAOuJJQQJ99BEACHYHv6XJ3w3AAAAACOGmC8m"  # Use a string instead of AzureKeyCredential
                )
    elif selected_model == 'o3-mini':
        client = AzureOpenAI(
                api_version="2024-12-01-preview",
                azure_endpoint="https://nextgenagilehu4401821853.openai.azure.com/",
                api_key="2ziHxIi2Pz511IscFVVjBpgHkm2nDXNfPAvXyFWpsmCHlpHwAOuJJQQJ99BEACHYHv6XJ3w3AAAAACOGmC8m"  # Use a string instead of AzureKeyCredential
                )
    return client 

def clear_selections():
    global selected_items
    selected_items.clear()
    return ""

def add_selection(option):
    global selected_items
    selected_items += [option]
    flat_list = [item for sublist in selected_items for item in sublist]
    counts = Counter(flat_list)
    formatted_output = ", ".join(f"{count} {option.lower()}" for option, count in counts.items())
    return formatted_output

# Persona system prompts with emojis
persona_prompts = {
    "👨‍💼 Scrum Master": "You are an experienced Scrum Master responsible for creating detailed JIRA tickets across parameters like EPIC, Feature, Story, Tasks, Subtasks, assignee based on job designation,sprint and story points",
    "💡 Product Owner": "You are a Product Owner focused on defining product features, value delivery, and aligning JIRA tickets with business outcomes. Depending on requirements you can focus upon these aspects: Description,Business Context,User Value,Scope,Out of Scope,Dependencies,Acceptance Criteria,Definition of Done,User Impact Metrics,Next Steps",
    "👨‍💻 Developer": "You are a Software Developer breaking down technical objectives into implementation tasks, bug fixes, and engineering subtasks.Depending on requirements you can focus upon these aspects: Description, Acceptance Criteria, Technical Details:Frontend,Backend,Database,Integrations,Logging & Monitoring, Component/s, Definition of Done, Sub-tasks,Story Points",
    "🎨 Designer": "You are a UX/UI Designer structuring JIRA tickets around user flows, UI components, and design assets."
}

# Default prompt examples
persona_default_prompts = {
    "👨‍💼 Scrum Master": "Objective: Automate JIRA story creation for a new agile project",
    "💡 Product Owner": "Objective: Define user stories for a travel booking feature",
    "👨‍💻 Developer": "Objective: Break down the user login system into dev tasks",
    "🎨 Designer": "Objective: Create tasks for designing the homepage UI"
}

model_select = {
    "O-4 Mini": "o4-mini",
    "O-3 Mini" : "o3-mini"
}
# Process function with memory per persona
def process_files(message, chat_history, persona,model):
    selected_model = model_select.get(model)
    client = prepare_model(selected_model)
    system_prompt = persona_prompts.get(persona)
    memory = structured_memory[persona]

    if message.lower().startswith("objective:"):
        objective = message.split(":", 1)[-1].strip()
        if len(selected_items) > 0:
            flat_list = [item for sublist in selected_items for item in sublist]
            counts = Counter(flat_list)
            formatted_output = ", ".join(f"{count} {option.lower()}" for option, count in counts.items())
            full_prompt = f"Create a complete JIRA ticket for the {objective} where your team structure is: {formatted_output}"
        else:
            full_prompt = f"Create a complete JIRA ticket for the {objective}"

        chat_response = client.chat.completions.create(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": full_prompt}
            ],
            max_completion_tokens=100000,
            model=selected_model
        )
        bot_message = chat_response.choices[0].message.content
        memory["current"] = bot_message

    elif any(keyword in message.lower() for keyword in ["improve", "update", "change"]):
        if "current" not in memory:
            bot_message = "Please provide a project objective first using 'Objective: <your project>'."
        else:
            previous_output = memory["current"]
            improvement_prompt = f"Improve the following JIRA ticket structure based on the user's instruction.\n\nStructure:\n{previous_output}\n\nUser Request:\n{message}"
            chat_response = client.chat.completions.create(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": improvement_prompt}
                ],
                max_completion_tokens=100000,
                model=selected_model
            )
            bot_message = chat_response.choices[0].message.content
            memory["current"] = bot_message
    else:
        bot_message = "Please start with 'Objective: <project>' or a request to update something."

    chat_history.append({"role": "user", "content": message})
    chat_history.append({"role": "assistant", "content": bot_message})
    time.sleep(1)
    return "", chat_history

def insert_prefix(prefix_text, textbox_content):
    if not textbox_content.lower().startswith(prefix_text.lower()):
        return f"{prefix_text} {textbox_content}"
    else:
        return textbox_content

def set_default_prompt(persona):
    return persona_default_prompts.get(persona, "")

# Gradio App
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):  # Sidebar
            gr.Image("/afh/projects/NextGenAgileTool-aefe44c3-e188-4674-9bd7-fd55842e362e/data/logo-removebg-preview.png", width=135, height=100,show_download_button=False , container= False,show_fullscreen_button=False, show_label=False)
            # gr.Markdown("### *NextGenAgile* – AI Jira Assistant")
            gr.Markdown("### 💽 Choose Model")

            model = gr.Dropdown(
                choices=[
                    "O-4 Mini",
                    "O-3 Mini"
                ],
                value="",
                label=""
            )

            gr.Markdown("### 🧠 Choose Persona")

            persona = gr.Dropdown(
                choices=[
                    "👨‍💼 Scrum Master",
                    "💡 Product Owner",
                    "👨‍💻 Developer",
                    "🎨 Designer"
                ],
                value="👨‍💼 Scrum Master",
                label="🧑 Persona"
            )

        with gr.Column(scale=9):
            
            chatbot = gr.Chatbot(
                label="Chat History",
                type="messages",
                height=500,
                value=[
                    {"role": "assistant", "content": "Hello! 👋 Start with an 'Objective' or ask to 'Improve' an earlier ticket."}
                ]
            )

            with gr.Row():
                objective_btn = gr.Button("➡️ Objective")
                improve_btn = gr.Button("✏️ Improve")
                update_btn = gr.Button("♻️ Update")
                change_btn = gr.Button("🔄 Change")

            with gr.Row():
                msg = gr.Textbox(
                    placeholder="Type your request here...",
                    label="Your Prompt",
                    lines=2,
                    submit_btn=True
                )

            clear = gr.ClearButton([msg, chatbot])

            # Hook logic
            msg.submit(process_files, [msg, chatbot, persona,model], [msg, chatbot])
            persona.change(set_default_prompt, persona, msg)

            objective_btn.click(insert_prefix, inputs=[gr.State("Objective:"), msg], outputs=msg)
            improve_btn.click(insert_prefix, inputs=[gr.State("Improve:"), msg], outputs=msg)
            update_btn.click(insert_prefix, inputs=[gr.State("Update:"), msg], outputs=msg)
            change_btn.click(insert_prefix, inputs=[gr.State("Change:"), msg], outputs=msg)

# if __name__ == "__main__":
#     demo.launch(debug=True, share=True)


server_info = demo.launch(share=True)

# Extract the public URL using regex
match = re.search(r"Running on public URL: (https?://[^\s]+)", server_info)
public_url = match.group(1) if match else "URL not found"

# Save only the public URL to a file
with open("gradio_link.txt", "w") as f:
    f.write(public_url)