#!/usr/bin/env python3
"""
Script to fix the OpinionDynamicsEngine to accept personas directly.
Run this from your AI_Personas directory on your MacBook.
"""

import sys
import os

def update_dynamics_file():
    """Update src/influence/dynamics.py"""

    filepath = "src/influence/dynamics.py"

    print(f"Updating {filepath}...")

    # Read the file
    with open(filepath, 'r') as f:
        content = f.read()

    # Old function signature
    old_signature = """    def run_dynamics(
        self,
        persona_ids: List[str],
        question: str,
        max_rounds: int = 5,
        convergence_threshold: float = 0.1,
        network_type: str = "scale_free",
        progress_callback: Optional[Callable[[str], None]] = None,
    ) -> List[RoundResult]:
        """
        Run opinion dynamics simulation.

        Args:
            persona_ids: List of persona IDs to include
            question: The proposal/question to discuss
            max_rounds: Maximum number of rounds
            convergence_threshold: Stop if total change < threshold
            network_type: Network topology ("scale_free", "small_world", "fully_connected")
            progress_callback: Function to call with progress updates

        Returns:
            List of RoundResult for each round
        """
        # Load personas and build influence network
        personas = [self.persona_db.get_persona(pid) for pid in persona_ids]
        influence_network = InfluenceNetwork(personas, network_type=network_type)"""

    # New function signature
    new_signature = """    def run_dynamics(
        self,
        question: str,
        max_rounds: int = 5,
        convergence_threshold: float = 0.1,
        network_type: str = "scale_free",
        persona_ids: Optional[List[str]] = None,
        personas: Optional[List[Persona]] = None,
        progress_callback: Optional[Callable[[str], None]] = None,
    ) -> List[RoundResult]:
        """
        Run opinion dynamics simulation.

        Args:
            question: The proposal/question to discuss
            max_rounds: Maximum number of rounds
            convergence_threshold: Stop if total change < threshold
            network_type: Network topology ("scale_free", "small_world", "fully_connected")
            persona_ids: List of persona IDs to include (for small group mode)
            personas: Pre-loaded personas list (for population mode)
            progress_callback: Function to call with progress updates

        Returns:
            List of RoundResult for each round
        """
        # Load personas from IDs if not provided directly
        if personas is None:
            if persona_ids is None:
                raise ValueError("Either persona_ids or personas must be provided")
            personas = [self.persona_db.get_persona(pid) for pid in persona_ids]

        # Build influence network
        influence_network = InfluenceNetwork(personas, network_type=network_type)"""

    # Replace
    if old_signature in content:
        content = content.replace(old_signature, new_signature)
        print("✓ Updated run_dynamics method signature")
    else:
        print("⚠ Could not find old signature - file may already be updated or different than expected")
        return False

    # Write back
    with open(filepath, 'w') as f:
        f.write(content)

    print(f"✓ Successfully updated {filepath}")
    return True


def update_ui_file():
    """Update pages/3_🌐_Opinion_Equilibria.py"""

    filepath = "pages/3_🌐_Opinion_Equilibria.py"

    print(f"\nUpdating {filepath}...")

    # Read the file
    with open(filepath, 'r') as f:
        lines = f.readlines()

    # Check if already has PopulationNetwork import
    has_population_import = any("from src.influence.population_network import PopulationNetwork" in line for line in lines)

    if has_population_import:
        print("✓ File already has PopulationNetwork import")
    else:
        # Find the line with "from src.influence.models import OpinionPosition"
        for i, line in enumerate(lines):
            if "from src.influence.models import OpinionPosition" in line:
                # Add imports after this line
                lines.insert(i + 1, "from src.influence.population_network import PopulationNetwork\n")
                lines.insert(i + 2, "from src.population.variant_generator import VariationLevel\n")
                print("✓ Added PopulationNetwork and VariationLevel imports")
                break

    # Check if has PersonaDatabase fixes
    old_db_load = "    persona_db.load_from_directory"
    needs_db_fix = any(old_db_load in line for line in lines)

    if needs_db_fix:
        # Fix PersonaDatabase calls
        for i, line in enumerate(lines):
            if "persona_db.load_from_directory" in line:
                # Comment out this line
                lines[i] = "    # " + line.lstrip()
            elif "persona_db.list_personas()" in line:
                lines[i] = line.replace("list_personas()", "get_all_personas()")
        print("✓ Fixed PersonaDatabase method calls")

    # Write back
    with open(filepath, 'w') as f:
        f.writelines(lines)

    print(f"✓ Successfully updated {filepath}")
    return True


if __name__ == "__main__":
    print("=" * 60)
    print("Fixing OpinionDynamicsEngine to support both modes")
    print("=" * 60)

    # Check we're in the right directory
    if not os.path.exists("src/influence/dynamics.py"):
        print("\n❌ Error: Cannot find src/influence/dynamics.py")
        print("Please run this script from the AI_Personas directory:")
        print("  cd /Users/matthewstroud/AI_Personas")
        print("  python fix_dynamics.py")
        sys.exit(1)

    # Update files
    success = True
    success = update_dynamics_file() and success
    success = update_ui_file() and success

    if success:
        print("\n" + "=" * 60)
        print("✅ All files updated successfully!")
        print("=" * 60)
        print("\nNext steps:")
        print("1. Restart your Streamlit app")
        print("2. Navigate to the Opinion Equilibria page")
        print("3. Try the new Population Network mode!")
    else:
        print("\n⚠ Some updates failed - check the messages above")
        sys.exit(1)