import pandas as pd
import asyncio
import json
import os
import time
from typing import Optional, Generator, AsyncGenerator
from pydantic import BaseModel

class SimulationConfig(BaseModel):
    speed: float = 1.0  # Multiplier, or seconds delay? Let's say speed multiplier (1x, 2x...)
    # But for simplicity, let's treat it as "transactions per second" or delay.
    # The requirement says "Speed Slider: Adjusts how fast...".
    # Let's map speed 1-10 to delay. 
    # Speed 1 = 2 sec delay, Speed 10 = 0.1 sec delay.
    
class SimulationStatus(BaseModel):
    status: str  # "running", "paused", "stopped"
    current_index: int
    total_transactions: int
    speed: float

class SimulationManager:
    def __init__(self):
        self.dataset: Optional[pd.DataFrame] = None
        self.current_index = 0
        self.is_running = False
        self.speed = 1.0  # Default 1x speed
        self.delay = 1.0  # Seconds between transactions
        self.dataset_path = None
        
    def load_dataset(self, path: str):
        self.dataset_path = path
        # Check if file exists, if not try to find it
        if not os.path.exists(path):
            # Try finding it in common locations
            possible_paths = [
                path,
                f"dataset/{path}",
                f"../dataset/{path}",
                "dataset/test_dataset.csv.gz",
                "dataset/test_dataset.csv"
            ]
            for p in possible_paths:
                if os.path.exists(p):
                    self.dataset_path = p
                    break
        
        if self.dataset_path and os.path.exists(self.dataset_path):
            print(f"Loading simulation dataset from {self.dataset_path}")
            # Load only necessary columns to save memory if needed
            self.dataset = pd.read_csv(self.dataset_path)
            # Ensure it's sorted by step
            if 'step' in self.dataset.columns:
                self.dataset = self.dataset.sort_values('step')
            print(f"Loaded {len(self.dataset)} transactions for simulation")
        else:
            print("Simulation dataset not found!")

    def start(self):
        if self.dataset is None:
            raise ValueError("Dataset not loaded")
        self.is_running = True
        return self.get_status()

    def stop(self):
        self.is_running = False
        return self.get_status()

    def reset(self):
        self.current_index = 0
        self.is_running = False
        return self.get_status()

    def set_speed(self, speed: float):
        self.speed = max(0.1, min(speed, 10.0)) # Clamp between 0.1 and 10
        # Formula: Higher speed = Lower delay.
        # Base delay = 2 seconds (at 1x)
        # Delay = 2 / speed
        self.delay = 2.0 / self.speed
        return self.get_status()

    def get_status(self):
        return {
            "status": "running" if self.is_running else "paused",
            "current_index": self.current_index,
            "total_transactions": len(self.dataset) if self.dataset is not None else 0,
            "speed": self.speed
        }

    async def stream_generator(self) -> AsyncGenerator[str, None]:
        """
        Yields SSE formatted data:
        data: {json_content}\n\n
        """
        while True:
            if not self.is_running:
                await asyncio.sleep(0.5)
                # yield Comment to keep connection alive
                yield ": keep-alive\n\n"
                continue

            if self.dataset is None or self.current_index >= len(self.dataset):
                self.is_running = False
                yield f"data: {json.dumps({'event': 'finished'})}\\n\\n"
                break

            # Get current transaction
            row = self.dataset.iloc[self.current_index]
            transaction = row.to_dict()
            
            # Convert numpy types to native python types for JSON serialization
            for key, value in transaction.items():
                if hasattr(value, 'item'):
                    transaction[key] = value.item()
            
            # Create payload
            payload = {
                "transaction": transaction,
                "index": self.current_index,
                "total": len(self.dataset)
            }
            
            yield f"data: {json.dumps(payload)}\\n\\n"
            
            self.current_index += 1
            
            # Wait based on speed
            await asyncio.sleep(self.delay)

simulation_manager = SimulationManager()