import uuid

from locust import HttpUser, between, task


class ConsultancyStudent(HttpUser):
    # Simulate a student taking 5-15 seconds to read a response
    wait_time = between(5, 15)

    def on_start(self):
        """Called when a 'student' starts their session."""
        self.user_id = str(uuid.uuid4())
        self.common_headers = {"Content-Type": "application/json"}

    @task(3)  # This task happens 3x more often than the others
    def ask_about_germany(self):
        payload = {
            "prompt": "What are the requirements for a student visa in Germany?",
            "userId": self.user_id,
        }
        with self.client.post(
            "/api/chat/chat",
            json=payload,
            headers=self.common_headers,
            catch_response=True,
        ) as response:
            if response.status_code == 200:
                response.success()
            else:
                response.failure(f"Failed with status: {response.status_code}")

    @task(1)
    def ask_about_costs(self):
        # Follow-up question to test LangGraph's memory/state
        payload = {
            "prompt": "And how much is the blocked account amount currently?",
            "userId": self.user_id,
        }
        self.client.post("/api/chat/chat", json=payload, headers=self.common_headers)

    @task(1)
    def small_talk(self):
        # Test the 'greeting' or 'smalltalk' intent node (LLM only, no RAG)
        self.client.post(
            "/api/chat/chat",
            json={"prompt": "Thank you, you are helpful!", "userId": self.user_id},
        )