import random

from locust import HttpUser, between, task


class TuringApiUser(HttpUser):
    # Wait time between requests
    wait_time = between(1, 5)

    # List of supported languages
    languages = ["python", "java", "pharo"]

    # Code examples
    code_snippets = {
        "python": ["def init(self): pass", "print('Hello World')", "import os"],
        "java": ["public static void main(String[] args)", "System.out.println(e);", "private int x = 0;"],
        "pharo": ["Transcript show: 'Hello'.", "^ self size", "Object subclass: #Name"]
    }

    @task(1)
    def health_check(self):
        """
        Checks if the API is alive.
        """
        self.client.get("/")

    @task(3)
    def predict_code_classification(self):
        """
        Sends a prediction request with language as a QUERY PARAMETER.
        """
        # Randomly selects one of the three languages
        selected_lang = random.choice(self.languages)
        texts = self.code_snippets.get(selected_lang, ["generic code"])
        
    
        json_payload = {
            "texts": texts
        }
        
        query_params = {
            "language": selected_lang
        }
        
        headers = {'Content-Type': 'application/json'}
        
        # Perform the request
        self.client.post(
            "/predict", 
            params=query_params,  
            json=json_payload,   
            headers=headers, 
            name="/predict"
        )