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

Metrics generators for load testing.
"""

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

from sqlalchemy import text

from mcpgateway.db import (
    ToolMetric,
    ResourceMetric,
    PromptMetric,
    ServerMetric,
    A2AAgentMetric,
)

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


class ToolMetricsGenerator(BaseGenerator):
    """Generate tool usage metrics."""

    def get_count(self) -> int:
        """Get total number of metric records to generate."""
        tool_count = self.get_scale_config("gateways", 100) * self.get_scale_config("tools_per_gateway_avg", 50)
        metrics_per_tool = self.get_scale_config("metrics_per_tool_avg", 100)
        return int(tool_count * metrics_per_tool)

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

    def generate(self) -> Generator[ToolMetric, None, None]:
        """Generate tool metric records.

        Yields:
            ToolMetric instances
        """
        tool_result = self.db.execute(text("SELECT id FROM tools ORDER BY created_at"))
        tool_ids = [row[0] for row in tool_result.fetchall()]

        if not tool_ids:
            self.logger.warning("No tools found - cannot generate metrics")
            return

        min_metrics = self.get_scale_config("metrics_per_tool_min", 10)
        max_metrics = self.get_scale_config("metrics_per_tool_max", 500)
        avg_metrics = self.get_scale_config("metrics_per_tool_avg", 100)

        # Get temporal distribution
        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 tool_id in tool_ids:
            num_metrics = max(min_metrics, min(max_metrics, int(random.gauss(avg_metrics, avg_metrics / 3))))
            timestamps = exponential_decay_temporal(num_metrics, start_date, end_date, recent_percent)

            for timestamp in timestamps:
                # 95% success rate
                is_success = random.random() < 0.95
                response_time = random.expovariate(1 / 150) if is_success else random.expovariate(1 / 500)

                yield ToolMetric(
                    tool_id=tool_id,
                    timestamp=timestamp,
                    response_time=round(response_time, 2),
                    is_success=is_success,
                    error_message=None if is_success else random.choice([
                        "Connection timeout",
                        "Rate limit exceeded",
                        "Invalid request",
                        "Gateway unavailable",
                    ]),
                )


class ResourceMetricsGenerator(BaseGenerator):
    """Generate resource access metrics."""

    def get_count(self) -> int:
        """Get total number of metric records to generate."""
        user_count = self.get_scale_config("users", 100)
        resources_per_user = self.get_scale_config("resources_per_user_avg", 100)
        metrics_per_resource = self.get_scale_config("metrics_per_resource_avg", 50)
        return int(user_count * resources_per_user * metrics_per_resource)

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

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

        Yields:
            ResourceMetric 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()]

        if not resource_ids:
            self.logger.warning("No resources found - cannot generate metrics")
            return

        min_metrics = self.get_scale_config("metrics_per_resource_min", 5)
        max_metrics = self.get_scale_config("metrics_per_resource_max", 200)
        avg_metrics = self.get_scale_config("metrics_per_resource_avg", 50)

        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 resource_id in resource_ids:
            num_metrics = max(min_metrics, min(max_metrics, int(random.gauss(avg_metrics, avg_metrics / 3))))
            timestamps = exponential_decay_temporal(num_metrics, start_date, end_date, recent_percent)

            for timestamp in timestamps:
                is_success = random.random() < 0.97
                response_time = random.expovariate(1 / 80) if is_success else random.expovariate(1 / 300)

                yield ResourceMetric(
                    resource_id=resource_id,
                    timestamp=timestamp,
                    response_time=round(response_time, 2),
                    is_success=is_success,
                    error_message=None if is_success else random.choice([
                        "Resource not found",
                        "Access denied",
                        "Resource unavailable",
                    ]),
                )


class PromptMetricsGenerator(BaseGenerator):
    """Generate prompt usage metrics."""

    def get_count(self) -> int:
        """Get total number of metric records to generate."""
        user_count = self.get_scale_config("users", 100)
        prompts_per_user = self.get_scale_config("prompts_per_user_avg", 100)
        metrics_per_prompt = self.get_scale_config("metrics_per_prompt_avg", 30)
        return int(user_count * prompts_per_user * metrics_per_prompt)

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

    def generate(self) -> Generator[PromptMetric, None, None]:
        """Generate prompt metric records.

        Yields:
            PromptMetric instances
        """
        prompt_result = self.db.execute(text("SELECT id FROM prompts ORDER BY created_at"))
        prompt_ids = [row[0] for row in prompt_result.fetchall()]

        if not prompt_ids:
            self.logger.warning("No prompts found - cannot generate metrics")
            return

        min_metrics = self.get_scale_config("metrics_per_prompt_min", 5)
        max_metrics = self.get_scale_config("metrics_per_prompt_max", 150)
        avg_metrics = self.get_scale_config("metrics_per_prompt_avg", 30)

        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 prompt_id in prompt_ids:
            num_metrics = max(min_metrics, min(max_metrics, int(random.gauss(avg_metrics, avg_metrics / 3))))
            timestamps = exponential_decay_temporal(num_metrics, start_date, end_date, recent_percent)

            for timestamp in timestamps:
                is_success = random.random() < 0.96
                response_time = random.expovariate(1 / 100) if is_success else random.expovariate(1 / 400)

                yield PromptMetric(
                    prompt_id=prompt_id,
                    timestamp=timestamp,
                    response_time=round(response_time, 2),
                    is_success=is_success,
                    error_message=None if is_success else "Prompt execution failed",
                )


class ServerMetricsGenerator(BaseGenerator):
    """Generate server usage metrics."""

    def get_count(self) -> int:
        """Get total number of metric records to generate."""
        user_count = self.get_scale_config("users", 100)
        servers_per_user = self.get_scale_config("servers_per_user_avg", 10)
        metrics_per_server = self.get_scale_config("metrics_per_server_avg", 200)
        return int(user_count * servers_per_user * metrics_per_server)

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

    def generate(self) -> Generator[ServerMetric, None, None]:
        """Generate server metric records.

        Yields:
            ServerMetric 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 metrics")
            return

        min_metrics = self.get_scale_config("metrics_per_server_min", 20)
        max_metrics = self.get_scale_config("metrics_per_server_max", 1000)
        avg_metrics = self.get_scale_config("metrics_per_server_avg", 200)

        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_metrics = max(min_metrics, min(max_metrics, int(random.gauss(avg_metrics, avg_metrics / 3))))
            timestamps = exponential_decay_temporal(num_metrics, start_date, end_date, recent_percent)

            for timestamp in timestamps:
                is_success = random.random() < 0.98
                response_time = random.expovariate(1 / 200) if is_success else random.expovariate(1 / 800)

                yield ServerMetric(
                    server_id=server_id,
                    timestamp=timestamp,
                    response_time=round(response_time, 2),
                    is_success=is_success,
                    error_message=None if is_success else "Server error",
                )


class A2AAgentMetricsGenerator(BaseGenerator):
    """Generate A2A agent usage metrics."""

    def get_count(self) -> int:
        """Get total number of metric records to generate."""
        user_count = self.get_scale_config("users", 100)
        agents_per_user = self.get_scale_config("a2a_agents_per_user_avg", 2)
        metrics_per_agent = self.get_scale_config("metrics_per_a2a_agent_avg", 150)
        return int(user_count * agents_per_user * metrics_per_agent)

    def get_dependencies(self) -> List[str]:
        """Depends on A2A agents."""
        return ["A2AAgentGenerator"]

    def generate(self) -> Generator[A2AAgentMetric, None, None]:
        """Generate A2A agent metric records.

        Yields:
            A2AAgentMetric instances
        """
        agent_result = self.db.execute(text("SELECT id FROM a2a_agents ORDER BY created_at"))
        agent_ids = [row[0] for row in agent_result.fetchall()]

        if not agent_ids:
            self.logger.info("No A2A agents found - skipping metrics generation")
            return

        min_metrics = self.get_scale_config("metrics_per_a2a_agent_min", 10)
        max_metrics = self.get_scale_config("metrics_per_a2a_agent_max", 500)
        avg_metrics = self.get_scale_config("metrics_per_a2a_agent_avg", 150)

        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 agent_id in agent_ids:
            num_metrics = max(min_metrics, min(max_metrics, int(random.gauss(avg_metrics, avg_metrics / 3))))
            timestamps = exponential_decay_temporal(num_metrics, start_date, end_date, recent_percent)

            for timestamp in timestamps:
                is_success = random.random() < 0.94
                response_time = random.expovariate(1 / 1200) if is_success else random.expovariate(1 / 3000)
                tokens_used = random.randint(50, 2000) if is_success else 0

                yield A2AAgentMetric(
                    a2a_agent_id=agent_id,
                    timestamp=timestamp,
                    response_time=round(response_time, 2),
                    is_success=is_success,
                    interaction_type=random.choice(["text_generation", "code_generation", "analysis", "translation"]),
                    error_message=None if is_success else random.choice([
                        "API rate limit exceeded",
                        "Authentication failed",
                        "Model unavailable",
                        "Request timeout",
                    ]),
                )