import streamlit as st
import json
import re
import networkx as nx
import matplotlib.pyplot as plt
from transformers import pipeline

# Initialize two text2text-generation pipelines using Flan-T5-large.
model_request_to_graph = pipeline('text2text-generation', model='google/flan-t5-large')
model_reply_with_graph = pipeline('text2text-generation', model='google/flan-t5-large')

def generate_input_graph(request_text):
    """
    Model 1: Convert the customer request into a structured graph.
    The prompt now includes an explicit example to guide the model.
    """
    prompt = (
        "You are an assistant that converts a natural language customer request into a structured graph. "
        "Output only valid JSON with exactly two keys: 'nodes' and 'edges'.\n"
        "Example output:\n"
        '{"nodes": ["City A", "City B", "City C", "City D", "City E"], '
        '"edges": [["City A", "City B"], ["City B", "City C"], ["City C", "City D"], ["City D", "City E"]]} \n\n'
        "Do not include any extra text, explanation, or commentary.\n\n"
        f"Customer Request: \"{request_text}\"\n\n"
        "Structured Graph JSON:"
    )
    output = model_request_to_graph(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]['generated_text']
    return output

def generate_reply_and_graph(request_text):
    """
    Model 2: Generate a detailed response along with a structured graph.
    The prompt now includes an explicit example of the expected JSON format.
    """
    prompt = (
        "You are an assistant that responds to a customer request with a detailed reply and a structured graph. "
        "First, provide a helpful textual reply to the request. Then output a valid JSON object with exactly two keys: 'nodes' and 'edges'.\n"
        "Example JSON output:\n"
        '{"nodes": ["City A", "City B", "City C", "City D", "City E"], '
        '"edges": [["City A", "City B"], ["City B", "City C"], ["City C", "City D"], ["City D", "City E"]]} \n\n'
        "Do not include any extra text or commentary outside the JSON.\n\n"
        f"Customer Request: \"{request_text}\"\n\n"
        "Detailed Response and Structured Graph JSON:"
    )
    output = model_reply_with_graph(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]['generated_text']
    return output

def try_parse_json(text):
    """
    Attempt to extract a valid JSON substring from the model output using a non-greedy match.
    Returns the parsed JSON if valid, otherwise logs the error and returns None.
    """
    st.write("Raw model output:", text)
    # Non-greedy match for JSON object
    json_pattern = re.compile(r'\{.*?\}', re.DOTALL)
    match = json_pattern.search(text)
    if match:
        json_str = match.group()
        try:
            parsed = json.loads(json_str)
            # Check for required keys and non-empty lists
            if 'nodes' in parsed and 'edges' in parsed and parsed['nodes'] and parsed['edges']:
                return parsed
            else:
                st.error("Parsed JSON is missing required keys or contains empty lists.")
                st.write("Extracted JSON:", json_str)
                return None
        except Exception as e:
            st.error(f"JSON parsing error: {e}")
            st.write("Extracted text:", json_str)
            return None
    else:
        st.error("No JSON object found in the model output.")
    return None

def build_and_visualize_graph(graph_data, title="Graph Visualization"):
    """
    Build a NetworkX directed graph from a dictionary with 'nodes' and 'edges'
    and visualize it using Matplotlib.
    """
    if not graph_data or 'nodes' not in graph_data or 'edges' not in graph_data:
        st.error("Invalid graph data format.")
        return
    G = nx.DiGraph()
    for node in graph_data['nodes']:
        G.add_node(node)
    for edge in graph_data['edges']:
        if isinstance(edge, list) and len(edge) == 2:
            G.add_edge(edge[0], edge[1])
        else:
            st.warning(f"Skipping invalid edge format: {edge}")
    plt.figure(figsize=(8, 6))
    pos = nx.spring_layout(G, seed=42)
    nx.draw(G, pos, with_labels=True, node_color='skyblue', edge_color='gray', node_size=2000, font_size=10)
    plt.title(title)
    st.pyplot(plt)

def main():
    st.title("Two-LLM Pipeline: Request & Reply with Graph Visualization")
    st.write(
        "Enter a customer request to see it transformed into an input graph and then receive a detailed reply "
        "with its own structured graph. Both graphs are built using non-empty 'nodes' and 'edges'."
    )
    
    customer_request = st.text_area("Enter Customer Request:", height=150)
    
    if st.button("Process Request"):
        if not customer_request.strip():
            st.error("Please enter a valid customer request.")
            return
        
        # --- Stage 1: Request to Input Graph ---
        st.subheader("Stage 1: Customer Request → Input Graph")
        model1_output = generate_input_graph(customer_request)
        st.code(model1_output, language="json")
        input_graph_data = try_parse_json(model1_output)
        if input_graph_data:
            st.write("Input Graph Data:")
            st.json(input_graph_data)
            build_and_visualize_graph(input_graph_data, title="Input Graph (Customer Request)")
        else:
            st.warning("Could not parse a valid graph structure from Model 1 output.")
        
        # --- Stage 2: Detailed Reply with Structured Graph ---
        st.subheader("Stage 2: Detailed Response with Structured Graph")
        model2_output = generate_reply_and_graph(customer_request)
        st.code(model2_output, language="json")
        output_graph_data = try_parse_json(model2_output)
        if output_graph_data:
            st.write("Output Graph Data:")
            st.json(output_graph_data)
            build_and_visualize_graph(output_graph_data, title="Output Graph (Structured Response)")
        else:
            st.warning("Could not parse a valid graph structure from Model 2 output.")

if __name__ == "__main__":
    main()