import streamlit as st
import json
from jsonschema import validate, ValidationError
import re
import networkx as nx
import matplotlib.pyplot as plt

# Mock existing resources (customizable)
existing_resources = {
    "storageAccount1": {
        "type": "Microsoft.Storage/storageAccounts",
        "location": "East US",
        "status": "existing",
        "properties": {
            "accountType": "Standard_LRS"
        }
    }
}

# ARM template schema (simplified version)
arm_schema = {
    "type": "object",
    "properties": {
        "$schema": {"type": "string"},
        "contentVersion": {"type": "string"},
        "parameters": {"type": "object"},
        "variables": {"type": "object"},
        "resources": {
            "type": "array",
            "items": {
                "type": "object",
                "required": ["type", "name", "location"],
                "properties": {
                    "type": {"type": "string"},
                    "name": {"type": "string"},
                    "location": {"type": "string"},
                    "properties": {"type": "object"},
                    "dependsOn": {"type": "array", "items": {"type": "string"}}
                }
            }
        },
        "outputs": {"type": "object"}
    },
    "required": ["$schema", "contentVersion", "resources"]
}

def validate_template(template):
    try:
        validate(instance=template, schema=arm_schema)
        return True, "Template is valid."
    except ValidationError as e:
        return False, f"Template validation error: {e.message}"

def resolve_template_expressions(template):
    def resolve_expression(match):
        expr = match.group(1)
        parts = expr.split('.')
        if parts[0] == 'parameters':
            return str(template.get('parameters', {}).get(parts[1], {}).get('defaultValue', ''))
        elif parts[0] == 'variables':
            return str(template.get('variables', {}).get(parts[1], ''))
        return match.group(0)

    template_str = json.dumps(template)
    resolved_str = re.sub(r'\[(\w+\.\w+)\]', resolve_expression, template_str)
    return json.loads(resolved_str)

def estimate_cost(resources):
    # This is a simplified cost estimation. In a real scenario, you'd need to
    # integrate with Azure's pricing API or use a more comprehensive pricing model.
    cost_estimates = {
        "Microsoft.Storage/storageAccounts": 10,  # $10 per month
        "Microsoft.Compute/virtualMachines": 50,  # $50 per month
        "Microsoft.Web/sites": 30,  # $30 per month
        "Microsoft.Sql/servers": 40,  # $40 per month
        # Add more resource types and their estimated costs
    }
    
    total_cost = 0
    for resource in resources:
        resource_type = resource.get('type')
        if resource_type in cost_estimates:
            total_cost += cost_estimates[resource_type]
        else:
            total_cost += 5  # Default cost for unknown resource types
    
    return total_cost

def create_dependency_graph(resources):
    G = nx.DiGraph()
    for resource in resources:
        resource_name = resource.get('name')
        G.add_node(resource_name)
        dependencies = resource.get('dependsOn', [])
        for dep in dependencies:
            G.add_edge(dep, resource_name)
    return G

def plot_dependency_graph(G):
    plt.figure(figsize=(12, 8))
    pos = nx.spring_layout(G)
    nx.draw(G, pos, with_labels=True, node_color='lightblue', 
            node_size=3000, font_size=8, font_weight='bold')
    edge_labels = nx.get_edge_attributes(G, 'weight')
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
    plt.title("Resource Dependency Graph")
    return plt

def lint_template(template):
    lint_results = []
    
    # Check 1: Ensure all resources have tags
    for resource in template.get('resources', []):
        if 'tags' not in resource:
            lint_results.append(f"Warning: Resource '{resource.get('name')}' does not have any tags.")
    
    # Check 2: Ensure parameters have descriptions
    for param_name, param_value in template.get('parameters', {}).items():
        if 'metadata' not in param_value or 'description' not in param_value['metadata']:
            lint_results.append(f"Warning: Parameter '{param_name}' does not have a description.")
    
    # Check 3: Warn about hardcoded values in resource properties
    for resource in template.get('resources', []):
        for prop, value in resource.get('properties', {}).items():
            if isinstance(value, (str, int, float)) and not isinstance(value, bool):
                lint_results.append(f"Info: Consider using a parameter for the hardcoded value in resource '{resource.get('name')}', property '{prop}'.")
    
    return lint_results

def generate_system_diagram(resources):
    diagram = ["graph TD"]
    resource_nodes = {}
    
    # Create nodes for each resource
    for i, resource in enumerate(resources):
        resource_type = resource['type'].split('/')[-1]
        node_id = f"R{i}"
        resource_nodes[resource['name']] = node_id
        diagram.append(f"{node_id}[{resource_type}<br/>{resource['name']}]")
    
    # Create connections based on dependencies
    for resource in resources:
        if 'dependsOn' in resource:
            for dependency in resource['dependsOn']:
                if dependency in resource_nodes:
                    diagram.append(f"{resource_nodes[dependency]} --> {resource_nodes[resource['name']]}")
    
    return "\n".join(diagram)

def simulate_deployment(arm_template):
    try:
        template = json.loads(arm_template)
        
        # Validate template
        is_valid, validation_message = validate_template(template)
        if not is_valid:
            return [], [], [], [], 0, None, [], "", validation_message

        # Resolve expressions
        resolved_template = resolve_template_expressions(template)

        resources_to_create = []
        resources_to_update = []
        resources_to_delete = []
        resource_details = []

        for resource in resolved_template.get('resources', []):
            resource_name = resource.get('name')
            resource_type = resource.get('type')
            location = resource.get('location')
            properties = resource.get('properties', {})

            if resource_name in existing_resources:
                existing = existing_resources[resource_name]
                if (existing['type'] == resource_type and
                    existing['location'] == location and
                    existing['properties'] == properties):
                    resources_to_update.append(resource_name)
                else:
                    resources_to_create.append(resource_name)
            else:
                resources_to_create.append(resource_name)

            resource_details.append({
                "name": resource_name,
                "type": resource_type,
                "location": location,
                "properties": properties,
                "dependsOn": resource.get('dependsOn', [])
            })

        for resource_name in existing_resources:
            if resource_name not in [r.get('name') for r in resolved_template.get('resources', [])]:
                resources_to_delete.append(resource_name)

        # Estimate cost
        estimated_cost = estimate_cost(resource_details)

        # Create dependency graph
        dependency_graph = create_dependency_graph(resolved_template.get('resources', []))

        # Lint template
        lint_results = lint_template(template)

        # Generate system diagram
        system_diagram = generate_system_diagram(resource_details)

        return resources_to_create, resources_to_update, resources_to_delete, resource_details, estimated_cost, dependency_graph, lint_results, system_diagram, "Simulation completed successfully."
    except json.JSONDecodeError:
        return [], [], [], [], 0, None, [], "", "Invalid JSON format. Please check your ARM template."

# Streamlit UI
st.title("Comprehensive ARM Template Simulator")
st.subheader("Paste your ARM template below to simulate its deployment.")

# Input box for ARM template
template_input = st.text_area("Paste ARM Template JSON here:", height=300)

# Button to simulate the evaluation
if st.button("Simulate Template"):
    if template_input:
        resources_to_create, resources_to_update, resources_to_delete, resource_details, estimated_cost, dependency_graph, lint_results, system_diagram, message = simulate_deployment(template_input)
        
        st.subheader("Simulation Results:")
        st.write(message)

        if resources_to_create or resources_to_update or resources_to_delete:
            st.write("### Resources to be Created:")
            if resources_to_create:
                st.write(resources_to_create)
            else:
                st.write("No new resources will be created.")

            st.write("### Resources to be Updated:")
            if resources_to_update:
                st.write(resources_to_update)
            else:
                st.write("No resources will be updated.")

            st.write("### Resources to be Deleted:")
            if resources_to_delete:
                st.write(resources_to_delete)
            else:
                st.write("No resources will be deleted.")

            st.write("### Detailed Resource Information:")
            for resource in resource_details:
                st.write(f"**Name:** {resource['name']}")
                st.write(f"**Type:** {resource['type']}")
                st.write(f"**Location:** {resource['location']}")
                st.write("**Properties:**")
                st.json(resource['properties'])
                st.write("---")

            st.write(f"### Estimated Monthly Cost: ${estimated_cost}")

            st.write("### Resource Dependency Graph:")
            if dependency_graph:
                fig = plot_dependency_graph(dependency_graph)
                st.pyplot(fig)
            else:
                st.write("No dependencies found between resources.")

            st.write("### System Diagram:")
            if system_diagram:
                st.mermaid(system_diagram)
            else:
                st.write("Unable to generate system diagram.")

            st.write("### Template Lint Results:")
            if lint_results:
                for result in lint_results:
                    st.write(result)
            else:
                st.write("No linting issues found.")
        else:
            st.write("No changes detected in the simulation.")
    else:
        st.error("Please paste an ARM template to simulate.")