import os
import gradio as gr
import requests
from typing import Dict, List, Any, Optional
from dataclasses import asdict, dataclass
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_core.messages import HumanMessage, AIMessage
from mistralai import Mistral
import json
import base64
import io
from io import BytesIO
from bs4 import BeautifulSoup
import re
from PIL import Image

# Constants
MAX_IMAGE_SIZE_KB = 50  # 50KB maximum size for images
MAX_RESIZE_ATTEMPTS = 3  # Maximum number of resize attempts
from langchain_core.runnables.graph import MermaidDrawMethod

from dotenv import load_dotenv
load_dotenv()

# Initialize Mistral client
client = Mistral(api_key=os.getenv("MISTRAL_API_KEY"))

@dataclass
class SimulationState:
    """Enhanced state management for the simulation workflow"""
    url: str = ""
    content: str = ""
    content_type: str = ""  # 'article' or 'image'
    summary: str = ""
    agents: List[Dict[str, Any]] = None
    behaviors: List[str] = None
    emojis: Dict[str, str] = None
    simulation_code: str = ""
    documentation: str = ""
    messages: List = None
    # New fields for enhanced simulation
    environment: Dict[str, Any] = None
    hypothesis: str = ""
    story_narrative: str = ""
    research_question: str = ""
    expected_outcomes: List[str] = None
    agent_interactions: List[Dict[str, Any]] = None
    environmental_factors: List[Dict[str, Any]] = None
    # New fields for code validation
    code_errors: List[str] = None
    validation_attempts: int = 0
    max_validation_attempts: int = 3
    code_is_valid: bool = False

    def __post_init__(self):
        if self.agents is None:
            self.agents = []
        if self.behaviors is None:
            self.behaviors = []
        if self.emojis is None:
            self.emojis = {}
        if self.messages is None:
            self.messages = []
        if self.environment is None:
            self.environment = {}
        if self.expected_outcomes is None:
            self.expected_outcomes = []
        if self.agent_interactions is None:
            self.agent_interactions = []
        if self.environmental_factors is None:
            self.environmental_factors = []
        if self.code_errors is None:
            self.code_errors = []

    def to_dict(self):
        return asdict(self)

import time
from tenacity import retry, stop_after_attempt, wait_exponential

class MistralModelHub:
    def __init__(self, api_key: str):
        self.client = Mistral(api_key=api_key)
        
        # Available Mistral models
        self.models = {
            'large': MistralLLM(self.client, "mistral-large-latest"),
            'medium': MistralLLM(self.client, "magistral-small-2506"),
            'small': MistralLLM(self.client, "ministral-3b-2410"),
            'codestral': MistralLLM(self.client, "codestral-latest"),
            'pixtral': MistralLLM(self.client, "pixtral-12b-2409"),
            'nemo': MistralLLM(self.client, "open-mistral-nemo"),
            'saba': MistralLLM(self.client, "mistral-saba-latest"),
        }
    
    def __getattr__(self, name):
        if name in self.models:
            return self.models[name]
        raise AttributeError(f"Model '{name}' not available")

class MistralLLM:
    def __init__(self, client, model: str):
        self.model = model
        self.client = client

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10)
    )
    def chat(self, messages: List[Dict[str, str]], temperature: float = 0.7, max_tokens: Optional[int] = None) -> str:
        try:
            chat_messages = [
                {"role": msg["role"], "content": msg["content"]}
                for msg in messages
            ]
            print(f"Calling model: {self.model}")
            response = self.client.chat.complete(
                model=self.model,
                messages=chat_messages,
                temperature=temperature,
                max_tokens=max_tokens,
                response_format={"type": "text"}
            )
            
            if response.choices and len(response.choices) > 0:
                return response.choices[0].message.content
            return ""
            
        except Exception as e:
            print(f"Error calling Mistral API: {str(e)}")
            if hasattr(e, 'status_code') and e.status_code == 429:
                print("Rate limit exceeded. Waiting before retry...")
                time.sleep(5)
            raise

# Initialize LLM
llm_large = MistralModelHub(api_key=os.getenv("MISTRAL_API_KEY")).large
llm_codestral = MistralModelHub(api_key=os.getenv("MISTRAL_API_KEY")).codestral
llm_pixtral = MistralModelHub(api_key=os.getenv("MISTRAL_API_KEY")).pixtral
llm_medium = MistralModelHub(api_key=os.getenv("MISTRAL_API_KEY")).medium
llm_small = MistralModelHub(api_key=os.getenv("MISTRAL_API_KEY")).small
llm_saba = MistralModelHub(api_key=os.getenv("MISTRAL_API_KEY")).saba
llm_nemo = MistralModelHub(api_key=os.getenv("MISTRAL_API_KEY")).nemo

def resize_image_to_size_limit(image_data: bytes, content_type: str, max_size_kb: int = MAX_IMAGE_SIZE_KB, max_attempts: int = MAX_RESIZE_ATTEMPTS) -> str:
    """
    Resize image to fit within the specified size limit (in KB)
    Returns base64 encoded image data URL
    """
    def get_size_kb(data):
        return len(data) / 1024  # Convert bytes to KB
    
    current_data = image_data
    current_size_kb = get_size_kb(current_data)
    
    if current_size_kb <= max_size_kb:
        return f"data:{content_type};base64,{base64.b64encode(current_data).decode()}"
    
    # Try resizing up to max_attempts times
    for attempt in range(1, max_attempts + 1):
        try:
            # Open the image
            img = Image.open(BytesIO(current_data))
            
            # Calculate new dimensions (reduce by 20% each time)
            new_width = int(img.width * 0.8)
            new_height = int(img.height * 0.8)
            
            # Resize the image
            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
            
            # Convert back to bytes
            output = BytesIO()
            format = content_type.split('/')[-1].upper()
            if format not in ['JPEG', 'PNG', 'GIF']:  # Default to JPEG
                format = 'JPEG'
            img.save(output, format=format, quality=85, optimize=True)
            current_data = output.getvalue()
            
            # Check new size
            current_size_kb = get_size_kb(current_data)
            if current_size_kb <= max_size_kb:
                return f"data:{content_type};base64,{base64.b64encode(current_data).decode()}"
                
        except Exception as e:
            print(f"Error resizing image (attempt {attempt}): {str(e)}")
            break
    
    # If we get here, we couldn't resize within limits after max_attempts
    raise ValueError(f"Could not resize image below {max_size_kb}KB after {max_attempts} attempts")

def extract_content_node(state: SimulationState) -> SimulationState:
    """Extract content from URL - handles both articles and images. Only keep main article text, filter out navigation/copyright/unrelated sections."""
    try:
        response = requests.get(state.url, timeout=10)
        content_type = response.headers.get('content-type', '').lower()

        if 'image' in content_type:
            state.content_type = 'image'
            # Check image size and resize if needed
            image_data = response.content
            image_size_kb = len(image_data) / 1024  # Convert to KB
            if image_size_kb > MAX_IMAGE_SIZE_KB:
                print(f"Image size ({image_size_kb:.2f}KB) exceeds maximum allowed size ({MAX_IMAGE_SIZE_KB}KB). Resizing...")
                state.content = resize_image_to_size_limit(image_data, content_type)
            else:
                # Convert to base64 if within size limit
                image_data = base64.b64encode(image_data).decode()
                state.content = f"data:{content_type};base64,{image_data}"
            return state

        # Otherwise, treat as article
        state.content_type = 'article'
        soup = BeautifulSoup(response.text, 'html.parser')
        article_content = ''

        # Try to extract from <article> tag first
        article_tag = soup.find('article')
        if article_tag:
            article_content = article_tag.get_text(separator=' ', strip=True)
        else:
            # Fallback: find the largest <div> with <p> tags
            divs = soup.find_all('div')
            max_p_div = None
            max_p_count = 0
            for div in divs:
                p_count = len(div.find_all('p'))
                if p_count > max_p_count:
                    max_p_count = p_count
                    max_p_div = div
            if max_p_div:
                article_content = max_p_div.get_text(separator=' ', strip=True)

        # Fallback: all paragraphs
        if not article_content:
            paragraphs = soup.find_all('p')
            article_content = ' '.join([p.get_text().strip() for p in paragraphs])

        # Remove unwanted sections by keywords (navigation, copyright, etc.)
        unwanted_keywords = [
            'riproduzione riservata', 'copyright', 'da non perdere', 'prevPageLabel', 'nextPageLabel', 'condividi', 'link copiato', 'cookie', 'accetta', 'rifiuta'
        ]
        for keyword in unwanted_keywords:
            article_content = re.sub(rf'(?i){keyword}.*', '', article_content)

        # Remove extra whitespace and code blocks
        article_content = re.sub(r'\s+', ' ', article_content)
        article_content = re.sub(r'``````', '', article_content, flags=re.DOTALL)
        article_content = article_content.strip()
        state.content = article_content[:5000]  # Limit content size

    except Exception as e:
        state.content = f"Error extracting content: {str(e)}"
        state.content_type = 'error'
    return state

def process_image_node(state: SimulationState) -> SimulationState:
    """Process image content using Mistral's vision capabilities"""

    if state.content_type != 'image':
        return state

    messages = [
        {
            "role": "system",
            "content": """You are an expert image analyzer. Describe the image in detail, focusing on:
            1. Main subjects and objects
            2. Activities or interactions happening
            3. Social dynamics or relationships visible
            4. Setting and context
            5. Any text or symbols present

            Provide a comprehensive description that could be used to understand social dynamics for simulation purposes."""
        },
        {
            "role": "user",
            "content": f"Please describe this image in detail: {state.content}"
        }
    ]

    description = llm_pixtral.chat(messages)
    state.content = description
    state.content_type = 'article'  # Convert to article for further processing

    return state

def analyze_content_node(state: SimulationState) -> SimulationState:
    """Enhanced analysis to extract comprehensive simulation components"""

    system_prompt = """You are an expert in complex systems theory, social network analysis, and agent-based modeling. 
    
    Your task is to conduct a deep scientific analysis of the content to create a meaningful agent-based simulation. Consider:
    
    1. System Dynamics: Identify feedback loops, tipping points, and emergent behaviors
    2. Social Physics: Power dynamics, influence propagation, coalition formation
    3. Environmental Context: Physical, social, economic, and cultural environments
    4. Conflict and Cooperation: Competition for resources, alliance formation, negotiation
    5. Temporal Dynamics: How relationships and behaviors evolve over time
    
    Focus on creating agents with rich internal states and meaningful interactions that produce observable phenomena."""

    user_prompt = f"""
    Analyze this content for an advanced agent-based simulation:

    Content: {state.content}

    Provide a comprehensive JSON response with:
    {{
        "summary": "Deep analytical summary focusing on system dynamics",
        "research_domain": "Primary field of study (sociology, economics, politics, etc.)",
        "environment": {{
            "name": "Environment name",
            "description": "Detailed environment description",
            "resources": ["resource1", "resource2"],
            "constraints": ["constraint1", "constraint2"],
            "spatial_properties": "Physical or abstract space description",
            "temporal_dynamics": "How environment changes over time"
        }},
        "agents": [
            {{
                "name": "Agent name",
                "type": "Agent type (individual, organization, group, institution)",
                "description": "Detailed agent description",
                "attributes": {{
                    "influence_level": 1-10,
                    "resources": ["resource1", "resource2"],
                    "goals": ["goal1", "goal2"],
                    "constraints": ["constraint1", "constraint2"],
                    "decision_strategy": "How agent makes decisions",
                    "memory_span": "How long agent remembers interactions",
                    "cooperation_tendency": 1-10,
                    "aggression_level": 1-10,
                    "adaptation_rate": 1-10
                }},
                "initial_state": {{
                    "position": "starting position or status",
                    "energy": 1-100,
                    "relationships": {{}},
                    "knowledge": ["known_fact1", "known_fact2"]
                }}
            }}
        ],
        "agent_interactions": [
            {{
                "interaction_type": "cooperation/competition/negotiation/conflict/information_exchange",
                "participants": ["agent1", "agent2"],
                "conditions": "When this interaction occurs",
                "outcomes": {{
                    "positive": "What happens if interaction succeeds",
                    "negative": "What happens if interaction fails",
                    "environmental_impact": "How this affects the environment"
                }},
                "probability_factors": ["factor1", "factor2"]
            }}
        ],
        "environmental_factors": [
            {{
                "name": "Factor name",
                "type": "resource/constraint/event/pressure",
                "description": "How this factor affects the system",
                "impact_on_agents": "Specific effects on different agent types",
                "temporal_pattern": "constant/periodic/random/triggered"
            }}
        ],
        "system_dynamics": {{
            "feedback_loops": ["description of feedback loop"],
            "tipping_points": ["conditions that cause system change"],
            "equilibrium_states": ["possible stable states"],
            "emergent_behaviors": ["behaviors that emerge from interactions"]
        }}
    }}
    """

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]

    response = llm_medium.chat(messages, temperature=0.6)

    try:
        json_match = re.search(r'\{.*\}', response, re.DOTALL)
        if json_match:
            data = json.loads(json_match.group())
            state.summary = data.get("summary", "")
            state.environment = data.get("environment", {})
            state.agents = data.get("agents", [])
            state.agent_interactions = data.get("agent_interactions", [])
            state.environmental_factors = data.get("environmental_factors", [])
            
            # Extract system dynamics for later use
            system_dynamics = data.get("system_dynamics", {})
            state.behaviors = (
                system_dynamics.get("feedback_loops", []) +
                system_dynamics.get("emergent_behaviors", [])
            )
        else:
            state.summary = response
    except Exception as e:
        state.summary = f"Analysis error: {str(e)}"

    return state

def develop_hypothesis_node(state: SimulationState) -> SimulationState:
    """Develop scientific hypothesis and narrative story for the simulation"""

    system_prompt = """You are a research scientist developing hypotheses for agent-based social simulations. 
    
    Create a compelling research hypothesis that can be tested through simulation, along with a narrative story 
    that makes the simulation engaging and meaningful. The hypothesis should be:
    
    1. Testable through agent interactions
    2. Based on established social/behavioral theories
    3. Relevant to the analyzed content
    4. Capable of producing measurable outcomes
    
    The story should provide context and meaning to make the simulation educational and engaging."""

    agents_summary = "\n".join([
        f"- {agent['name']}: {agent['description']} (Influence: {agent['attributes']['influence_level']})"
        for agent in state.agents
    ])

    interactions_summary = "\n".join([
        f"- {interaction['interaction_type']}: {interaction['conditions']}"
        for interaction in state.agent_interactions
    ])

    user_prompt = f"""
    Develop a research hypothesis and story for this simulation:

    Summary: {state.summary}
    
    Environment: {state.environment.get('description', 'N/A')}
    
    Agents:
    {agents_summary}
    
    Key Interactions:
    {interactions_summary}

    Provide a JSON response with:
    {{
        "research_question": "Clear, testable research question",
        "hypothesis": "Scientific hypothesis that can be tested through simulation",
        "theoretical_framework": "Underlying social/behavioral theories",
        "story_narrative": "Engaging narrative that contextualizes the simulation",
        "expected_outcomes": [
            "Specific measurable outcome 1",
            "Specific measurable outcome 2",
            "Specific measurable outcome 3"
        ],
        "success_metrics": [
            "How to measure if hypothesis is supported",
            "Key indicators to track"
        ],
        "variables_to_manipulate": [
            {{
                "name": "Variable name",
                "description": "What this controls",
                "range": "Possible values",
                "impact": "Expected effect on system"
            }}
        ]
    }}
    """

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]

    response = llm_large.chat(messages, temperature=0.7)

    try:
        json_match = re.search(r'\{.*\}', response, re.DOTALL)
        if json_match:
            data = json.loads(json_match.group())
            state.research_question = data.get("research_question", "")
            state.hypothesis = data.get("hypothesis", "")
            state.story_narrative = data.get("story_narrative", "")
            state.expected_outcomes = data.get("expected_outcomes", [])
    except Exception as e:
        state.hypothesis = f"Hypothesis development error: {str(e)}"

    return state

def generate_emojis_node(state: SimulationState) -> SimulationState:
    """Generate emojis and visual representations for agents"""

    system_prompt = """You are a creative designer specializing in visual representation of social agents. 
    Generate appropriate emojis and visual symbols for each agent that represent their role, characteristics, and function in the social network."""

    agents_text = "\n".join([f"- {agent['name']}: {agent['description']}" for agent in state.agents])

    user_prompt = f"""
    Generate emojis for these social network agents:
    {agents_text}

    Provide a JSON response with emoji assignments:
    {{
        "agent_name": "🎭",
        "agent_name2": "👥"
    }}

    Choose emojis that best represent each agent's role and characteristics.
    """

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]

    response = llm_small.chat(messages, temperature=0.7)

    try:
        json_match = re.search(r'\{.*\}', response, re.DOTALL)
        if json_match:
            state.emojis = json.loads(json_match.group())
    except Exception as e:
        # Default emojis if parsing fails
        state.emojis = {agent['name']: "👤" for agent in state.agents}

    return state

def generate_simulation_node(state: SimulationState) -> SimulationState:
    """Generate advanced NetLogo-style simulation with comprehensive controls"""

    system_prompt = """You are an expert in advanced agent-based modeling and interactive simulations.

    Create a sophisticated HTML5 simulation with:
    
    1. ADVANCED CONTROLS:
       - Speed control (1-1000 scale)
       - Population controls for each agent type
       - Environmental parameter controls
       - Interaction probability controls
       - Resource availability controls
       - Distance threshold for interaction/cooperation/comunication 5-50 pixels
       - Scenario trigger buttons
    
    2. VISUAL FEATURES:
       - Dynamic environment visualization with color coding
       - Agent trails showing movement history
       - Network connections with varying thickness
       - Real-time statistics dashboard
       - Interactive agent information on hover/click
    
    3. SCIENTIFIC FEATURES:
       - Hypothesis testing interface
       - Data collection and export
       - Parameter sensitivity analysis
       - Scenario comparison tools
       - Statistical significance indicators
    
    4. TECHNICAL REQUIREMENTS:
       - Proper object-oriented agent architecture
       - Efficient spatial indexing for large populations
       - Event-driven interaction system
       - Configurable random seed for reproducibility
       - Performance monitoring and optimization

    FOCUS HEAVILY ON THE IMPLEMENTATION OF THE LOGIC AND THE AGENTS INTERACTIONS. DO NOT SKIP ANY IMPLEMENTAION. THE PROJECT MUST BE COMPLETE.
    
    The simulation must test the provided hypothesis and produce measurable outcomes."""

    agents_detail = "\n".join([
        f"""
        Agent: {agent['name']} ({state.emojis.get(agent['name'], '👤')})
        - Type: {agent['type']}
        - Goals: {agent['attributes']['goals']}
        - Decision Strategy: {agent['attributes']['decision_strategy']}
        - Cooperation: {agent['attributes']['cooperation_tendency']}/10
        - Aggression: {agent['attributes']['aggression_level']}/10
        - Resources: {agent['attributes']['resources']}
        """
        for agent in state.agents
    ])

    interactions_detail = "\n".join([
        f"""
        Interaction: {interaction['interaction_type']}
        - Participants: {interaction['participants']}
        - Conditions: {interaction['conditions']}
        - Success: {interaction['outcomes']['positive']}
        - Failure: {interaction['outcomes']['negative']}
        - Environmental Impact: {interaction['outcomes']['environmental_impact']}
        """
        for interaction in state.agent_interactions
    ])

    environment_detail = f"""
    Environment: {state.environment.get('name', 'Unknown')}
    - Description: {state.environment.get('description', '')}
    - Resources: {state.environment.get('resources', [])}
    - Constraints: {state.environment.get('constraints', [])}
    - Spatial Properties: {state.environment.get('spatial_properties', '')}
    """

    if isinstance(state, dict):
        research_question = state['research_question']
        expected_outcomes = state['expected_outcomes']
    else:
        research_question = state.research_question
        expected_outcomes = state.expected_outcomes
    user_prompt = f"""
    Create an advanced agent-based simulation for:

    HYPOTHESIS TO TEST: {state.hypothesis}
    
    STORY CONTEXT: {state.story_narrative}
    
    RESEARCH QUESTION: {research_question}

    ENVIRONMENT:
    {environment_detail}

    AGENTS:
    {agents_detail}

    INTERACTIONS:
    {interactions_detail}

    EXPECTED OUTCOMES: {expected_outcomes}

    Generate a complete HTML file with:

    1. CONTROL PANEL (left side):
       - Simulation speed slider (1-1000 fps)
       - Play/Pause/Reset/Step buttons
       - Population controls for each agent type (0-100)
       - Environmental parameter sliders based on the environment
       - Interaction probability controls
       - Distance threshold for interaction/cooperation/comunication 5-50 pixels
       - Scenario buttons for testing different conditions
       - Random seed input for reproducibility

    2. MAIN VISUALIZATION (center):
       - Large canvas (400x400) showing the environment. dO NOT MAKE THE CANVAS BIGGER THAN 400x400.
       - Dynamic background representing environmental state
       - Agents with emojis, trails, and status indicators
       - Connection lines showing relationships/interactions
       - Spatial zones for different environment areas
       - Ensure using a single emoji for each agent type. The emojis are provided for each agent.
       - Main visualization must not overlap with control panel or statistics panel.

    3. STATISTICS PANEL (right side):
       - STATISTIC PANEL MUST BE UPDATED IN REAL-TIME
       - Real-time metrics tracking hypothesis variables
       - Population counts and survival rates
       - Resource distribution graphs
       - Interaction frequency charts
       - Hypothesis testing results
       - Data export functionality

    4. ADVANCED FEATURES:
       - Agent inspector showing detailed state on click
       - Heatmaps for resource density and activity
       - Timeline scrubber for replay functionality
       - Parameter sweeping tools
       - Automated experiment runner

    Technical Implementation:
    - Use requestAnimationFrame for smooth animation
    - Implement spatial hashing for performance
    - Object-oriented agent and environment classes
    - Event system for interactions
    - Proper initialization and cleanup
    - Error handling and validation
    - Responsive design

    The simulation should automatically start with meaningful default values and immediately show the hypothesis being tested through agent behaviors.
    """

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]

    state.simulation_code = llm_codestral.chat(messages, temperature=0.2, max_tokens=8000)

    return state

def check_simulation_code_node(state: SimulationState) -> SimulationState:
    """Check and validate the generated simulation code using Codestral"""
    
    print(f"---CHECKING SIMULATION CODE (Attempt {state.validation_attempts + 1})---")
    
    # Increment validation attempts
    state.validation_attempts += 1
    
    system_prompt = """You are an expert code reviewer and HTML/JavaScript validator specializing in agent-based simulations.

    Your task is to:
    1. Analyze the provided HTML simulation code for syntax errors, logic issues, and functionality problems
    2. Check for proper implementation of agent behaviors, interactions, and controls
    3. Verify that all promised features are implemented correctly
    4. Identify any missing functionality or broken components
    5. Provide specific fixes and improvements

    Focus on:
    - HTML/CSS/JavaScript syntax validation
    - Proper agent implementation with movement, interactions, and behaviors
    - Working control panels with functional sliders and buttons
    - Real-time statistics and data display
    - Canvas rendering and animation loops
    - Event handling and user interactions
    - Performance and efficiency issues
    """

    # Extract HTML code if it's wrapped in markdown
    code_to_check = state.simulation_code
    if '```html' in code_to_check:
        html_match = re.search(r'```html\n(.*?)\n```', code_to_check, re.DOTALL)
        if html_match:
            code_to_check = html_match.group(1)

    user_prompt = f"""
    Analyze this agent-based simulation code and identify any issues:

    CODE TO CHECK:
    {code_to_check}

    EXPECTED FEATURES (check if properly implemented):
    - Hypothesis: {state.hypothesis}
    - Agents: {[agent['name'] for agent in state.agents]}
    - Emojis: {state.emojis}
    - Environment: {state.environment.get('name', 'Unknown')}
    - Interactions: {[interaction['interaction_type'] for interaction in state.agent_interactions]}

    Previous validation attempts: {state.validation_attempts - 1}
    Previous errors found: {state.code_errors}

    Provide a comprehensive JSON response:
    {{
        "is_valid": boolean,
        "syntax_errors": ["list of syntax errors found"],
        "logic_errors": ["list of logic/functionality errors"],
        "missing_features": ["list of features that should be implemented but are missing"],
        "performance_issues": ["list of performance problems"],
        "agent_issues": ["specific problems with agent implementation"],
        "interaction_issues": ["problems with agent interactions"],
        "control_issues": ["problems with UI controls"],
        "overall_assessment": "summary of code quality and functionality"
    }}
    """

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]

    response = llm_codestral.chat(messages, temperature=0.3)

    try:
        json_match = re.search(r'\{.*\}', response, re.DOTALL)
        if json_match:
            validation_result = json.loads(json_match.group())
            
            # Collect all errors
            all_errors = (
                validation_result.get("syntax_errors", []) +
                validation_result.get("logic_errors", []) +
                validation_result.get("missing_features", []) +
                validation_result.get("performance_issues", []) +
                validation_result.get("agent_issues", []) +
                validation_result.get("interaction_issues", []) +
                validation_result.get("control_issues", [])
            )
            
            state.code_errors = all_errors
            state.code_is_valid = validation_result.get("is_valid", False) and len(all_errors) == 0
            
            print(f"Code validation result: Valid={state.code_is_valid}, Errors found: {len(all_errors)}")
            
        else:
            state.code_errors = ["Failed to parse validation response"]
            state.code_is_valid = False
            
    except Exception as e:
        state.code_errors = [f"Validation error: {str(e)}"]
        state.code_is_valid = False

    return state

def fix_simulation_code_node(state: SimulationState) -> SimulationState:
    """Fix the simulation code based on validation errors"""
    
    print(f"---FIXING SIMULATION CODE (Attempt {state.validation_attempts})---")
    
    system_prompt = """You are an expert programmer specializing in HTML5 agent-based simulations.

    Your task is to fix the provided simulation code based on the identified errors and issues.
    Focus on:
    1. Fixing syntax errors and logic problems
    2. Implementing missing agent behaviors and interactions
    3. Ensuring proper canvas rendering and animation
    4. Making sure all UI controls work correctly
    5. Implementing real-time statistics updates
    6. Optimizing performance issues

    Provide ONLY the corrected HTML code without any explanations or markdown formatting.
    """

    # Extract HTML code if it's wrapped in markdown
    current_code = state.simulation_code
    if '```html' in current_code:
        html_match = re.search(r'```html\n(.*?)\n```', current_code, re.DOTALL)
        if html_match:
            current_code = html_match.group(1)

    errors_text = "\n".join([f"- {error}" for error in state.code_errors])

    user_prompt = f"""
    Fix this agent-based simulation code based on the identified errors:

    CURRENT CODE:
    {current_code}

    ERRORS TO FIX:
    {errors_text}

    REQUIREMENTS TO MAINTAIN:
    - Hypothesis: {state.hypothesis}
    - Agents: {[agent['name'] for agent in state.agents]}
    - Emojis: {state.emojis}
    - Environment: {state.environment.get('name', 'Unknown')}
    - Canvas size: 400x400 maximum
    - Real-time statistics panel
    - Working control panels
    - Proper agent movement and interactions

    Provide the complete fixed HTML code:
    """

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]

    fixed_code = llm_codestral.chat(messages, temperature=0.2, max_tokens=8000)
    
    # Clean up the response to ensure it's proper HTML
    if '```html' in fixed_code:
        html_match = re.search(r'```html\n(.*?)\n```', fixed_code, re.DOTALL)
        if html_match:
            fixed_code = html_match.group(1)
    
    state.simulation_code = fixed_code
    
    return state

def create_workflow() -> StateGraph:
    """Create the enhanced LangGraph workflow with code validation"""

    workflow = StateGraph(SimulationState)

    # Add all nodes
    workflow.add_node("extract_content", extract_content_node)
    workflow.add_node("process_image", process_image_node)
    workflow.add_node("analyze_content", analyze_content_node)
    workflow.add_node("develop_hypothesis", develop_hypothesis_node)
    workflow.add_node("generate_emojis", generate_emojis_node)
    workflow.add_node("generate_simulation", generate_simulation_node)
    workflow.add_node("check_simulation_code", check_simulation_code_node)
    workflow.add_node("fix_simulation_code", fix_simulation_code_node)

    # Define the enhanced flow
    workflow.set_entry_point("extract_content")

    # Conditional routing based on content type
    def route_content(state: SimulationState) -> str:
        if state.content_type == "image":
            return "process_image"
        else:
            state.content_type = "article"
            return "analyze_content"

    # Conditional routing for code validation
    def should_fix_code(state: SimulationState) -> str:
        if state.code_is_valid or state.validation_attempts >= state.max_validation_attempts:
            return "end"
        else:
            return "fix_code"

    workflow.add_conditional_edges(
        "extract_content",
        route_content,
        {
            "process_image": "process_image",
            "analyze_content": "analyze_content"
        }
    )

    # Enhanced workflow: analyze -> hypothesis -> emojis -> simulation -> validation loop
    workflow.add_edge("process_image", "analyze_content")
    workflow.add_edge("analyze_content", "develop_hypothesis")
    workflow.add_edge("develop_hypothesis", "generate_emojis")  
    workflow.add_edge("generate_emojis", "generate_simulation")
    workflow.add_edge("generate_simulation", "check_simulation_code")
    
    # Conditional edge for validation loop
    workflow.add_conditional_edges(
        "check_simulation_code",
        should_fix_code,
        {
            "fix_code": "fix_simulation_code",
            "end": END
        }
    )
    
    # After fixing code, check again
    workflow.add_edge("fix_simulation_code", "check_simulation_code")

    return workflow

# Create workflow and app
workflow = create_workflow()
memory = SqliteSaver.from_conn_string(":memory:")
app = workflow.compile()

def save_graph_as_png() -> bytes:
    """Save the LangGraph mermaid graph as PNG"""
    try:
        return app.get_graph().draw_mermaid_png(
            draw_method=MermaidDrawMethod.API,
            background_color="white",
            padding=10
        )
    except Exception as e:
        print(f"Error saving graph: {e}")
        return None

from gradio.themes.base import Base

# Create a custom light theme
custom_theme = gr.themes.Soft() # Could we enforce a light theme?

def create_gradio_interface():
    """Create the enhanced Gradio interface"""

    with gr.Blocks(title="Advanced Social Network Simulation Generator", theme=custom_theme) as demo:
        gr.Markdown("# 🌐 AI-Powered Advanced Social Network Simulation Generator")
        gr.Markdown(" * [Presentation Video](https://www.youtube.com/watch?v=8KRiPwacJnk)")
        gr.Markdown(" SNA (Social Network Analysis) By creating Agent-based Models (ABM) with AI agents, " \
        "you can generate advanced NetLogo-style social network simulations " \
        "with scientific hypothesis testing from articles or images." \
        "This tool uses the latest Mistral models to extract content, analyze it, "
        "and create complex simulations with rich agent interactions. By the way SNAstral is NOT [Netlogo](https://ccl.northwestern.edu/netlogo/). ;-)" )
        gr.Markdown("### 📚 Using Mistral Models:")
        for _, model in MistralModelHub(api_key=os.getenv("MISTRAL_API_KEY")).models.items():
            gr.Markdown(f"- {model.model}")
        gr.Markdown("Generate advanced NetLogo-style social network simulations with scientific hypothesis testing from articles or images using AI agents.")
        gr.Markdown("### ✨ NEW: Code validation with up to 3 retry attempts using Codestral")

        gr.Markdown("### 📚 Examples:")
        gr.Markdown("To test it use your favorite article. I tried with news articles. Here are some examples:")
        gr.Markdown("- https://www.ansa.it/sito/notizie/mondo/nordamerica/2025/06/10/dilagano-le-proteste-negli-usa150-arresti-a-san-francisco_3d6b794b-462c-4dd1-9c46-33ecf2bc5790.html")
        gr.Markdown("- https://www.foxnews.com/media/fox-news-beats-abc-nbc-cbs-during-weekday-primetime-while-cnn-has-lowest-rated-week-year")
        gr.Markdown(" It is not necessary but recommended to use an article with enough parties to analyze. If it does not have LLMs will hallucinate and make something anyway." )
        gr.Markdown(" Before generating the simulation you can edit the summary, agents, behaviors, emojis, hypothesis, story, environment and interactions. " )
        gr.Markdown("### DO NOT USE DARK THEME unless you don't care seeing the simulation results.")
        with gr.Row():
            with gr.Column(scale=2):
                # Input section
                url_input = gr.Textbox(
                    label="🔗 URL",
                    placeholder="Enter URL to article or image...",
                    lines=1
                )

                extract_btn = gr.Button("📄 Extract Content", variant="primary")

                # Content display
                content_display = gr.Markdown(label="📋 Extracted Content")

                # Editable summary and agents
                with gr.Group():
                    gr.Markdown("### ✏️ Edit Summary and Analysis")
                    summary_edit = gr.Textbox(
                        label="Summary",
                        lines=3,
                        placeholder="Summary will appear here..."
                    )

                    agents_edit = gr.Textbox(
                        label="Agents (JSON format - Edit as needed)",
                        lines=8,
                        placeholder="Agents JSON will appear here...",
                        info="Edit the JSON to modify agents before generating simulation"
                    )

                    behaviors_edit = gr.Textbox(
                        label="Behaviors (JSON format - Edit as needed)",
                        lines=4,
                        placeholder="Behaviors JSON will appear here...",
                        info="Edit the JSON to modify behaviors before generating simulation"
                    )

                # Enhanced fields
                with gr.Group():
                    gr.Markdown("### 🔬 Research Hypothesis")
                    hypothesis_edit = gr.Textbox(
                        label="Research Hypothesis",
                        lines=3,
                        placeholder="Hypothesis will appear here..."
                    )
                    
                    story_edit = gr.Textbox(
                        label="Story Narrative",
                        lines=4,
                        placeholder="Story narrative will appear here..."
                    )
                    
                    environment_edit = gr.Textbox(
                        label="Environment (JSON format)",
                        lines=6,
                        placeholder="Environment JSON will appear here..."
                    )
                    
                    interactions_edit = gr.Textbox(
                        label="Agent Interactions (JSON format)",
                        lines=8,
                        placeholder="Interactions JSON will appear here..."
                    )

                # Emoji selection
                with gr.Group():
                    gr.Markdown("### 🎭 Agent Emojis")
                    emojis_edit = gr.Textbox(
                        label="Emojis for Agents (JSON format - Edit as needed)",
                        lines=4,
                        placeholder="Emojis JSON will appear here...",
                        info="Edit the JSON to modify emojis before generating simulation"
                    )

                    generate_emojis_btn = gr.Button("🎨 Generate/Regenerate Emojis")

                # Generate simulation
                confirm_btn = gr.Button("🚀 Generate Advanced Simulation", variant="primary", size="lg")

            with gr.Column(scale=3):
                # Results section
                with gr.Tabs():
                    with gr.TabItem("🎮 Simulation"):
                        simulation_frame = gr.HTML(
                            label="Advanced Social Network Simulation",
                            value="<div style='padding: 20px; text-align: center; color: #666;'>Advanced simulation will appear here after generation...</div>"
                        )

                    with gr.TabItem("💻 Code Editor"):
                        simulation_code_editor = gr.Code(
                            label="Generated Simulation Code (Editable)",
                            language="html",
                            value="<!-- Advanced simulation code will appear here... -->",
                            lines=20,
                            interactive=True
                        )
                        
                        update_simulation_btn = gr.Button("🔄 Update Simulation from Code", variant="secondary")

                    with gr.TabItem("📖 Documentation"):
                        documentation_display = gr.Markdown(
                            label="Scientific Simulation Documentation",
                            value="Comprehensive documentation will be generated with the simulation..."
                        )

                    with gr.TabItem("🗺️ Workflow Graph"):
                        graph_display = gr.Image(
                            label="Enhanced LangGraph Workflow",
                            value=None
                        )

        # State management
        state = gr.State(SimulationState())

        # Event handlers
        def extract_content(url, current_state):
            current_state = SimulationState(url=url)
            config = {"configurable": {"thread_id": "simulation_thread"}}

            print("Current state:")
            print(f"{type(current_state)=}")
            if isinstance(current_state, SimulationState):
                current_state_dict = current_state.to_dict()
            else:
                current_state_dict = current_state
            for key, value in current_state_dict.items():
                print(f"{key}: {str(value)[:50]}")

            # Run extraction and analysis
            result = app.invoke(current_state, config)

            if not isinstance(result, dict):
                if isinstance(result, SimulationState):
                    result = result.to_dict()
                else:
                    raise ValueError(f"Unexpected result type: {type(result)}")

            print("Result:")
            print(f"{type(result)=}")
            for key, value in result.items():
                print(f"{key}: {str(value)[:50]}")

            # Update displays with enhanced content
            cleaned_content = re.sub(r'\s+', ' ', result['content']).strip()
            content_md = f"**Content Type:** {result['content_type']}\n\n**Content:**\n{cleaned_content[:1000]}..."

            return (
                result,  # Updated state
                content_md,  # Content display
                result['summary'],  # Summary
                json.dumps(result['agents'], indent=2),  # Agents as JSON string
                json.dumps(result['behaviors'], indent=2),  # Behaviors as JSON string
                result['hypothesis'],  # New: Hypothesis
                result['story_narrative'],  # New: Story
                json.dumps(result['environment'], indent=2),  # New: Environment
                json.dumps(result['agent_interactions'], indent=2),  # New: Interactions
                json.dumps(result['emojis'], indent=2, ensure_ascii=False)  # Emojis as JSON string
            )

        def generate_emojis_only(current_state, agents_json):
            # Update current state with potentially edited agents
            try:
                current_state.agents = json.loads(agents_json) if agents_json else []
            except:
                pass  # Keep existing agents if JSON parsing fails
                
            if not current_state.agents:
                return current_state, "{}"

            current_state = generate_emojis_node(current_state)
            return current_state, json.dumps(current_state.emojis, indent=2, ensure_ascii=False)

        def generate_final_simulation(current_state, summary, agents_json, behaviors_json, hypothesis, story, environment_json, interactions_json, emojis_json):
            # Update state with edited values
            current_state.summary = summary
            current_state.hypothesis = hypothesis
            current_state.story_narrative = story
            
            try:
                current_state.agents = json.loads(agents_json) if agents_json else []
            except json.JSONDecodeError as e:
                print(f"Error parsing agents JSON: {e}")
                current_state.agents = []
                
            try:
                current_state.behaviors = json.loads(behaviors_json) if behaviors_json else []
            except json.JSONDecodeError as e:
                print(f"Error parsing behaviors JSON: {e}")
                current_state.behaviors = []

            try:
                current_state.environment = json.loads(environment_json) if environment_json else {}
            except json.JSONDecodeError as e:
                print(f"Error parsing environment JSON: {e}")
                current_state.environment = {}

            try:
                current_state.agent_interactions = json.loads(interactions_json) if interactions_json else []
            except json.JSONDecodeError as e:
                print(f"Error parsing interactions JSON: {e}")
                current_state.agent_interactions = []
                
            try:
                current_state.emojis = json.loads(emojis_json) if emojis_json else {}
            except json.JSONDecodeError as e:
                print(f"Error parsing emojis JSON: {e}")
                current_state.emojis = {}

            # Reset validation state for new generation
            current_state.validation_attempts = 0
            current_state.code_errors = []
            current_state.code_is_valid = False

            # Generate simulation (now includes validation loop)
            current_state = generate_simulation_node(current_state)
            
            # The validation loop is now handled in the workflow
            config = {"configurable": {"thread_id": "validation_thread"}}
            validated_result = app.invoke(current_state, config)
            
            if isinstance(validated_result, dict):
                current_state.simulation_code = validated_result.get('simulation_code', current_state.simulation_code)
                current_state.validation_attempts = validated_result.get('validation_attempts', 0)
                current_state.code_errors = validated_result.get('code_errors', [])
                current_state.code_is_valid = validated_result.get('code_is_valid', False)
            
            # Generate documentation only after validation is complete
            if isinstance(current_state, dict):
                research_question = current_state['research_question']
                expected_outcomes = current_state['expected_outcomes']
            else:
                research_question = current_state.research_question
                expected_outcomes = current_state.expected_outcomes
                
            environment_detail = f"""
            Environment: {current_state.environment.get('name', 'Unknown')}
            - Description: {current_state.environment.get('description', '')}
            - Resources: {current_state.environment.get('resources', [])}
            - Constraints: {current_state.environment.get('constraints', [])}
            - Spatial Properties: {current_state.environment.get('spatial_properties', '')}
            """
            
            agents_detail = "\n".join([
                f"""
                Agent: {agent['name']} ({current_state.emojis.get(agent['name'], '👤')})
                - Type: {agent['type']}
                - Goals: {agent['attributes']['goals']}
                """
                for agent in current_state.agents
            ])
            
            interactions_detail = "\n".join([
                f"""
                Interaction: {interaction['interaction_type']}
                - Participants: {interaction['participants']}
                """
                for interaction in current_state.agent_interactions
            ])
            
            doc_prompt = f"""
            Create comprehensive documentation for this social network simulation based on:

            HYPOTHESIS: {current_state.hypothesis}
            STORY: {current_state.story_narrative}
            RESEARCH QUESTION: {research_question}
            ENVIRONMENT: {environment_detail}
            AGENTS: {agents_detail}
            INTERACTIONS: {interactions_detail}
            EXPECTED OUTCOMES: {expected_outcomes}
            
            VALIDATION RESULTS:
            - Code validation attempts: {current_state.validation_attempts}
            - Code is valid: {current_state.code_is_valid}
            - Errors found: {current_state.code_errors}

            Include:
            1. Overview and scientific purpose
            2. Hypothesis explanation and testing methodology
            3. Agent descriptions and behavioral models
            4. Environment and interaction explanations
            5. How to use advanced controls
            6. Parameters and their scientific effects
            7. Interpretation of results and statistical measures
            8. Educational objectives and learning outcomes
            9. Code quality assurance and validation process

            Make it accessible for users while being scientifically rigorous.
            """

            doc_messages = [
                {"role": "system", "content": "You are a scientific technical writer creating educational documentation."},
                {"role": "user", "content": doc_prompt}
            ]

            current_state.documentation = llm_nemo.chat(doc_messages, temperature=0.4)
            
            # Extract HTML from markdown for the iframe
            simulation_html = current_state.simulation_code
            if '```html' in simulation_html:
                # Extract content between ``````
                html_match = re.search(r'```html\n(.*?)\n```', simulation_html, re.DOTALL)
                if html_match:
                    simulation_html = html_match.group(1)
            # Use update_preview to wrap the HTML
            simulation_preview = update_preview(simulation_html)
            
            # Get workflow graph
            graph_png = None
            try:
                graph_bytes = save_graph_as_png()
                if graph_bytes:
                    from PIL import Image
                    import io
                    graph_png = Image.open(io.BytesIO(graph_bytes))
            except Exception as e:
                print(f"Error processing graph image: {e}")
                graph_png = None

            return (
                current_state,
                simulation_preview,  # Clean HTML for iframe (now wrapped)
                current_state.simulation_code,  # Original markdown for code editor
                current_state.documentation,
                graph_png
            )

        def update_preview(code):
            """Update the iframe preview using srcdoc for in-memory HTML preview."""
            if not code.strip():
                return "<div style='text-align:center; padding:50px;'>No code to preview</div>"
            # Escape quotes for srcdoc
            safe_code = code.replace('"', '&quot;').replace("'", "&#39;")
            return f'''<iframe style="width:100%;height:900px;border:1px solid #ccc;" srcdoc="{safe_code}"></iframe>'''

        def update_simulation_from_code(current_state, code):
            """Update simulation iframe from manually edited code"""
            current_state.simulation_code = code
            # Extract HTML from markdown if needed

            simulation_html = code
            if '```html' in code:
                html_match = re.search(r'```html\n(.*?)\n```', code, re.DOTALL)
                if html_match:  
                    simulation_html = html_match.group(1)
            # Use update_preview to wrap the HTML
            return current_state, update_preview(simulation_html)

        # Wire up events
        extract_btn.click(
            extract_content,
            inputs=[url_input, state],
            outputs=[state, content_display, summary_edit, agents_edit, behaviors_edit, hypothesis_edit, story_edit, environment_edit, interactions_edit, emojis_edit]
        )

        generate_emojis_btn.click(
            generate_emojis_only,
            inputs=[state, agents_edit],
            outputs=[state, emojis_edit]
        )

        confirm_btn.click(
            generate_final_simulation,
            inputs=[state, summary_edit, agents_edit, behaviors_edit, hypothesis_edit, story_edit, environment_edit, interactions_edit, emojis_edit],
            outputs=[state, simulation_frame, simulation_code_editor, documentation_display, graph_display]
        )
        
        update_simulation_btn.click(
            update_simulation_from_code,
            inputs=[state, simulation_code_editor],
            outputs=[state, simulation_frame]
        )

    return demo

# Launch the application
if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(share=False, debug=False)