import os,re
import gradio as gr
import nest_asyncio
from langchain import PromptTemplate
from llama_index.core import StorageContext, load_index_from_storage
import networkx as nx
from pyvis.network import Network
from IPython.display import HTML, Markdown, display

# nest_asyncio.apply()

kg_index_path = "./telcom_full_property_kg_processed_schema/" 
kg_plot_path = kg_index_path+"full_kg.html"
os.environ["OPENAI_API_KEY"] = os.getenv('oai')


index = load_index_from_storage(
    StorageContext.from_defaults(persist_dir="./telcom_full_property_kg_processed_schema")
)
query_engine = index.as_query_engine(
    include_text=True,
    similarity_top_k=7,
)
retriever = index.as_retriever(
    include_text=True,  # include source text, default True
    similarity_top_k=7,
)


teamplate_prompt_upsell = '''You are a virtual assistant for a telecom company, designed to assist users with their queries and potentially upsell services. Your task is to analyze the customer's data from context, their query, and offer the most appropriate assistance.

First, you will be given the customer's data context. This information will help you understand the customer's current plan and usage patterns:

When interacting with a customer, you will receive a query with their details like name or phone number.
<query>
{QUERY}
</query>

Analyze the query to determine the type of assistance required. Categorize it into one of the following:
1. Technical Support
2. Billing Inquiry
3. Plan Information
4. Service Upgrade
5. General Inquiry

Based on the query type and customer data, provide an appropriate response. Your response should:
1. Address the customer's immediate concern
2. Be clear and concise
3. Use a friendly and causal tone
4. Make sure to provide facts and relations for each response
5. Use Emojis to engage the customer in conversation

If the query presents an opportunity for upselling, consider recommending relevant services or upgrades based on the customer's current plan and usage patterns. However, ensure that your primary focus remains on resolving the customer's initial query.

Format your response as follows:

<response>
<query_type>[Categorized query type]</query_type>
<answer>[Your detailed response addressing the customer's query]</answer>
<reference>[Provide the reference documents used for generating the response]</reference>
<facts>[Provide the facts used for generating the response]</facts>
<upsell_opportunity>[If applicable, provide a brief upsell recommendation]</upsell_opportunity>
</response>

Remember to always prioritize customer satisfaction and only suggest upsells when they genuinely benefit the customer.
'''




def parse_response_with_regex(xml_response):
    # Define regex patterns for each tag
    query_type_pattern = re.compile(r'<query_type>(.*?)</query_type>', re.DOTALL)
    answer_pattern = re.compile(r'<answer>(.*?)</answer>', re.DOTALL)
    reference_pattern = re.compile(r'<reference>(.*?)</reference>', re.DOTALL)
    facts_pattern = re.compile(r'<facts>(.*?)</facts>', re.DOTALL)
    upsell_opportunity_pattern = re.compile(r'<upsell_opportunity>(.*?)</upsell_opportunity>', re.DOTALL)

    # Extract data using regex
    query_type = query_type_pattern.search(xml_response).group(1).strip()
    answer = answer_pattern.search(xml_response).group(1).strip()
    reference = reference_pattern.search(xml_response).group(1).strip()
    facts = facts_pattern.search(xml_response).group(1).strip()
    upsell_opportunity = upsell_opportunity_pattern.search(xml_response).group(1).strip()

    # Format the extracted information
    formatted_response = f"""
    ### Query Type
    {query_type}

    ### Answer
    {answer}

    ### Reference
    {reference}

    ### Facts
    {facts}

    ### Upsell Opportunity
    {upsell_opportunity}
    """
    return formatted_response

def extract_pattern_triplet(text):
    # Define the regex pattern to match the desired format
    pattern = re.compile(r'\b\w+\b\s*->\s*\b\w+\b\s*->\s*\b\w+\b')
    # Find all matches in the text
    matches = pattern.findall(text)
    return "\n <br> ".join(matches)

def query_tqa(query):
    data = {
        'QUERY': query
    }
    prompt = PromptTemplate(template=teamplate_prompt_upsell, input_variables=["QUERY"])
    query_ready = prompt.format(**data)
    response = query_engine.query(query_ready)
    nodes = retriever.retrieve(query_ready)
    parsed_resp = parse_response_with_regex(str(response))

    reference = []
    reference_text = []
    for node in nodes:
        reference.append(extract_pattern_triplet(node.text))
        reference_text.append(node.text)

    return parsed_resp, response, reference , reference_text

def plot_full_kg():
    """Plot the full knowledge graph and return the HTML representation."""
    return HTML(filename=kg_plot_path)

with gr.Blocks() as demo:
    gr.Markdown("<h1>Telcom Graph-RAG v0.1</h1>")

    with gr.Tab("Virtual Assistant"):
        with gr.Row():
            query_input = gr.Textbox(label="Input Your Query..")
        with gr.Row():
            model_output = gr.Textbox(label="Response")
            model_metadata = gr.Textbox(label="Raw Response")
        with gr.Row():
            reference = gr.HTML(label="Extracted Reference")
            reference_text = gr.Textbox(label="Extracted Reference raw")


        ask_button = gr.Button("Ask TelcomVA!!")

    with gr.Accordion("Explore KG!", open=False):
        gr.Markdown("This KG is built using a subset of Github repositories. ")
        kg_output = gr.HTML()
        plot_button = gr.Button("Plot Full KG!!")

    ask_button.click(query_tqa, inputs=[query_input], outputs=[model_output,model_metadata,reference,reference_text])
    plot_button.click(plot_full_kg, outputs=kg_output)
    examples = gr.Examples(
    examples=[
        ["what are the upselling ideas for roaming package you can recommend for customer Rina Wati."],
    ],
    inputs=[query_input]
    )

demo.launch(auth=(os.getenv('id'), os.getenv('pass')), share=True)
# demo.launch(share=False)