# -*- coding: utf-8 -*-
"""Location: ./tests/load/generators/sessions.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

MCP session and subscription generators for load testing.
"""

import json
import random
import uuid
from datetime import datetime, timedelta
from typing import Generator, List

from sqlalchemy import text

from mcpgateway.db import (
    SessionRecord,
    SessionMessageRecord,
    ResourceSubscription,
)

from ..utils.distributions import exponential_decay_temporal
from .base import BaseGenerator


class MCPSessionGenerator(BaseGenerator):
    """Generate MCP protocol sessions."""

    def get_count(self) -> int:
        """Get total number of MCP sessions to generate."""
        server_count = self.get_scale_config("servers_per_user_avg", 10) * self.get_scale_config("users", 100)
        sessions_per_server = self.get_scale_config("mcp_sessions_per_server_avg", 10)
        return int(server_count * sessions_per_server)

    def get_dependencies(self) -> List[str]:
        """Depends on servers."""
        return ["ServerGenerator"]

    def generate(self) -> Generator[SessionRecord, None, None]:
        """Generate MCP session records.

        Yields:
            SessionRecord instances
        """
        server_result = self.db.execute(text("SELECT id FROM servers ORDER BY created_at"))
        server_ids = [row[0] for row in server_result.fetchall()]

        if not server_ids:
            self.logger.warning("No servers found - cannot generate MCP sessions")
            return

        min_sessions = self.get_scale_config("mcp_sessions_per_server_min", 1)
        max_sessions = self.get_scale_config("mcp_sessions_per_server_max", 50)
        avg_sessions = self.get_scale_config("mcp_sessions_per_server_avg", 10)

        start_date = datetime.fromisoformat(self.config.get("temporal", {}).get("start_date", "2023-01-01"))
        end_date = datetime.fromisoformat(self.config.get("temporal", {}).get("end_date", datetime.now().isoformat()))
        recent_percent = self.config.get("temporal", {}).get("recent_data_percent", 80) / 100

        for server_id in server_ids:
            num_sessions = max(min_sessions, min(max_sessions, int(random.gauss(avg_sessions, avg_sessions / 3))))
            timestamps = exponential_decay_temporal(num_sessions, start_date, end_date, recent_percent)

            for timestamp in timestamps:
                session_id = str(uuid.uuid4())
                last_accessed = timestamp + timedelta(minutes=random.randint(1, 120))

                data = {
                    "server_id": server_id,
                    "client_info": {
                        "name": random.choice(["mcp-client", "claude-desktop", "custom-client"]),
                        "version": f"{random.randint(1,3)}.{random.randint(0,9)}.{random.randint(0,20)}",
                    },
                    "transport": random.choice(["sse", "websocket", "stdio"]),
                }

                yield SessionRecord(
                    session_id=session_id,
                    created_at=timestamp,
                    last_accessed=last_accessed,
                    data=json.dumps(data),
                )


class MCPMessageGenerator(BaseGenerator):
    """Generate MCP protocol messages."""

    def get_count(self) -> int:
        """Get total number of MCP messages to generate."""
        server_count = self.get_scale_config("servers_per_user_avg", 10) * self.get_scale_config("users", 100)
        sessions_per_server = self.get_scale_config("mcp_sessions_per_server_avg", 10)
        messages_per_session = self.get_scale_config("mcp_messages_per_session_avg", 20)
        return int(server_count * sessions_per_server * messages_per_session)

    def get_dependencies(self) -> List[str]:
        """Depends on MCP sessions."""
        return ["MCPSessionGenerator"]

    def generate(self) -> Generator[SessionMessageRecord, None, None]:
        """Generate MCP message records.

        Yields:
            SessionMessageRecord instances
        """
        session_result = self.db.execute(text("SELECT session_id, created_at FROM mcp_sessions ORDER BY created_at"))
        sessions = [(row[0], datetime.fromisoformat(row[1]) if isinstance(row[1], str) else row[1]) for row in session_result.fetchall()]

        if not sessions:
            self.logger.warning("No MCP sessions found - cannot generate messages")
            return

        min_messages = self.get_scale_config("mcp_messages_per_session_min", 5)
        max_messages = self.get_scale_config("mcp_messages_per_session_max", 100)
        avg_messages = self.get_scale_config("mcp_messages_per_session_avg", 20)

        message_types = [
            {"jsonrpc": "2.0", "method": "tools/list", "params": {}},
            {"jsonrpc": "2.0", "method": "tools/call", "params": {"name": "example_tool"}},
            {"jsonrpc": "2.0", "method": "resources/list", "params": {}},
            {"jsonrpc": "2.0", "method": "resources/read", "params": {"uri": "file:///example"}},
            {"jsonrpc": "2.0", "method": "prompts/list", "params": {}},
            {"jsonrpc": "2.0", "method": "prompts/get", "params": {"name": "example_prompt"}},
        ]

        for session_id, session_created in sessions:
            num_messages = max(min_messages, min(max_messages, int(random.gauss(avg_messages, avg_messages / 3))))

            for i in range(num_messages):
                created_at = session_created + timedelta(seconds=random.randint(1, 3600))
                last_accessed = created_at + timedelta(milliseconds=random.randint(10, 5000))

                yield SessionMessageRecord(
                    session_id=session_id,
                    message=json.dumps(random.choice(message_types)),
                    created_at=created_at,
                    last_accessed=last_accessed,
                )


class ResourceSubscriptionGenerator(BaseGenerator):
    """Generate resource subscriptions."""

    def get_count(self) -> int:
        """Get total number of resource subscriptions to generate."""
        user_count = self.get_scale_config("users", 100)
        resources_per_user = self.get_scale_config("resources_per_user_avg", 100)
        subscription_rate = self.get_scale_config("resource_subscription_rate", 0.1)  # 10% subscribed
        return int(user_count * resources_per_user * subscription_rate)

    def get_dependencies(self) -> List[str]:
        """Depends on resources and users."""
        return ["ResourceGenerator", "UserGenerator"]

    def generate(self) -> Generator[ResourceSubscription, None, None]:
        """Generate resource subscription records.

        Yields:
            ResourceSubscription instances
        """
        resource_result = self.db.execute(text("SELECT id FROM resources ORDER BY created_at"))
        resource_ids = [row[0] for row in resource_result.fetchall()]

        user_result = self.db.execute(text("SELECT email FROM email_users ORDER BY created_at"))
        user_emails = [row[0] for row in user_result.fetchall()]

        if not resource_ids or not user_emails:
            self.logger.warning("No resources or users found - cannot generate subscriptions")
            return

        subscription_rate = self.get_scale_config("resource_subscription_rate", 0.1)
        start_date = datetime.fromisoformat(self.config.get("temporal", {}).get("start_date", "2023-01-01"))
        end_date = datetime.fromisoformat(self.config.get("temporal", {}).get("end_date", datetime.now().isoformat()))

        # Create subscriptions for a random sample of resources
        num_subscriptions = int(len(resource_ids) * subscription_rate)
        sampled_resources = random.sample(resource_ids, min(num_subscriptions, len(resource_ids)))

        for resource_id in sampled_resources:
            # Each resource can have multiple subscribers
            num_subscribers = random.randint(1, min(5, len(user_emails)))
            selected_subscribers = random.sample(user_emails, num_subscribers)

            for subscriber_email in selected_subscribers:
                created_at = start_date + (end_date - start_date) * random.random()
                last_notification = created_at + timedelta(days=random.randint(1, 30)) if random.random() < 0.5 else None

                yield ResourceSubscription(
                    resource_id=resource_id,
                    subscriber_id=subscriber_email,
                    created_at=created_at,
                    last_notification=last_notification,
                )