import json
import matplotlib.pyplot as plt
import streamlit as st
from utils import client

def analyze_vfx_potential(thread_id, additional_context=None):
    run = client.beta.threads.runs.create(
        thread_id=thread_id,
        assistant_id="asst_WZPuGYRu6zU3XomrOietVAS5"
    )
    
    while run.status in ['queued', 'in_progress', 'cancelling']:
        run = client.beta.threads.runs.retrieve(
            thread_id=thread_id,
            run_id=run.id
        )
    
    if run.status == 'completed':
        messages = client.beta.threads.messages.list(thread_id=thread_id)
        analysis = next((msg.content[0].text.value for msg in reversed(list(messages)) if msg.role == "assistant"), "")
        return analysis
    else:
        return f"Error: Run status is {run.status}"

def process_vfx_potential(analysis):
    try:
            analysis = json.loads(analysis)
        
            st.subheader("VFX Potential Analysis")
        
        # Overall VFX Score
            scene_names = list(analysis.keys())
            vfx_potential = [analysis[scene]['vfx_potential'] for scene in scene_names]
            vfx_complexity = [analysis[scene]['vfx_complexity'] for scene in scene_names]
            budget_percentage = [analysis[scene]['budget_percentage'] for scene in scene_names]
            
            # Create subplots for the charts
            fig, ax = plt.subplots(3, 1, figsize=(10, 12), sharex=True)
            
            # Plot VFX potential
            ax[0].bar(scene_names, vfx_potential, color='b', alpha=0.7)
            ax[0].set_title('VFX Potential for Each Scene')
            ax[0].set_ylabel('VFX Potential (0 to 1)')
            ax[0].set_ylim(0, 1)
            
            # Plot VFX complexity
            ax[1].bar(scene_names, vfx_complexity, color='g', alpha=0.7)
            ax[1].set_title('VFX Complexity for Each Scene')
            ax[1].set_ylabel('VFX Complexity (0 to 1)')
            ax[1].set_ylim(0, 1)
            
            # Plot Budget Percentage Allocation
            ax[2].bar(scene_names, budget_percentage, color='r', alpha=0.7)
            ax[2].set_title('Budget Allocation for Each Scene')
            ax[2].set_ylabel('Budget Percentage (%)')
            ax[2].set_ylim(0, max(budget_percentage) + 5)
            ax[2].set_xlabel('Scene Names')

            # Rotate x-axis labels for better readability
            plt.xticks(rotation=45, ha='right')
            
            # Adjust layout to prevent overlap
            plt.tight_layout()
            st.pyplot(fig)
        
    except Exception as e:
        st.error(f"Error processing data for VFX Potential: {e}")
        st.write("Please check the structure of the JSON data:")
        st.json(analysis)