from dataclasses import dataclass
from typing import Dict, List, Optional
import time
import json
import subprocess
import statistics

@dataclass
class WorkloadMetrics:
    # Prompt handling metrics
    short_prompt_latency: float
    long_prompt_latency: float
    input_tokens_per_second: float
    output_tokens_per_second: float

    # Task-specific performance
    code_generation_speed: float
    math_computation_speed: float
    multilingual_performance: Dict[str, float]

    # System prompt impact
    baseline_latency: float
    system_prompt_overhead: float

    # Generation characteristics
    token_consistency: float  # Variance in token generation speed
    completion_success_rate: float
    average_generation_length: float

@dataclass
class WorkloadConfig:
    short_prompt_len: int = 100
    long_prompt_len: int = 2000
    target_output_len: int = 500
    system_prompt: str = "You are a helpful AI assistant."
    custom_prompts: Optional[Dict[str, str]] = None

class WorkloadBenchmark:
    def __init__(
        self,
        model_name: str,
        config: Optional[WorkloadConfig] = None
    ):
        self.model_name = model_name
        self.config = config or WorkloadConfig()

        # Default prompts for different tasks
        self.default_prompts = {
            "short": {
                "general": "Explain what photosynthesis is.",
                "code": "Write a Python function to calculate fibonacci numbers.",
                "math": "Solve the quadratic equation: 2x² + 5x - 3 = 0",
                "multilingual": {
                    "en": "Describe a beautiful sunset.",
                    "es": "Describe una puesta de sol hermosa.",
                    "fr": "Décrivez un beau coucher de soleil.",
                    "de": "Beschreiben Sie einen schönen Sonnenuntergang."
                }
            },
            "long": {
                "general": """Write a comprehensive essay about the impact of artificial intelligence
                on society, covering aspects such as employment, ethics, education, healthcare, and
                future implications. Include specific examples and potential solutions to challenges.""",
                "code": """Create a complete Python class for a library management system. Include
                methods for adding books, removing books, checking out books, returning books, and
                generating reports. Include proper documentation and error handling.""",
                "math": """Solve the following calculus problems step by step:
                1. Find the derivative of f(x) = x³ln(x) + e^(x²)
                2. Calculate the integral of g(x) = sin²(x)cos(x)
                3. Find the local maxima and minima of h(x) = x³ - 3x² + 1""",
            }
        }

        # Use custom prompts if provided
        self.prompts = self.config.custom_prompts or self.default_prompts

    def _run_generation(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        max_tokens: Optional[int] = None
    ) -> Dict:
        """Run a single generation and return metrics."""
        try:
            cmd = [
                "cortex", "run",
                "-d",
                self.model_name,
                "--max-tokens", str(max_tokens) if max_tokens else str(self.config.target_output_len),
                "--json"
            ]

            input_data = {
                "messages": [
                    {"role": "system", "content": system_prompt} if system_prompt else None,
                    {"role": "user", "content": prompt}
                ]
            }

            start_time = time.time()
            result = subprocess.run(
                cmd,
                input=json.dumps(input_data),
                capture_output=True,
                text=True,
                check=True
            )
            end_time = time.time()

            response = json.loads(result.stdout)

            return {
                "latency": end_time - start_time,
                "input_tokens": response.get("usage", {}).get("prompt_tokens", 0),
                "output_tokens": response.get("usage", {}).get("completion_tokens", 0),
                "success": True,
                "output": response.get("choices", [{}])[0].get("message", {}).get("content", "")
            }
        except Exception as e:
            return {
                "latency": time.time() - start_time,
                "input_tokens": 0,
                "output_tokens": 0,
                "success": False,
                "error": str(e)
            }

    def _measure_prompt_performance(self) -> tuple[float, float, float, float]:
        """Measure performance with different prompt lengths."""
        # Short prompt performance
        short_results = self._run_generation(self.prompts["short"]["general"])

        # Long prompt performance
        long_results = self._run_generation(self.prompts["long"]["general"])

        # Calculate tokens per second
        input_tps = (
            short_results["input_tokens"] + long_results["input_tokens"]
        ) / (short_results["latency"] + long_results["latency"])

        output_tps = (
            short_results["output_tokens"] + long_results["output_tokens"]
        ) / (short_results["latency"] + long_results["latency"])

        return (
            short_results["latency"],
            long_results["latency"],
            input_tps,
            output_tps
        )

    def _measure_task_performance(self) -> tuple[float, float, Dict[str, float]]:
        """Measure performance on specific tasks."""
        # Code generation
        code_result = self._run_generation(self.prompts["short"]["code"])

        # Math computation
        math_result = self._run_generation(self.prompts["short"]["math"])

        # Multilingual performance
        multi_results = {}
        for lang, prompt in self.prompts["short"]["multilingual"].items():
            result = self._run_generation(prompt)
            multi_results[lang] = result["latency"]

        return (
            code_result["latency"],
            math_result["latency"],
            multi_results
        )

    def _measure_system_prompt_impact(self) -> tuple[float, float]:
        """Measure impact of system prompt on performance."""
        # Without system prompt
        baseline = self._run_generation(
            self.prompts["short"]["general"],
            system_prompt=None
        )

        # With system prompt
        with_system = self._run_generation(
            self.prompts["short"]["general"],
            system_prompt=self.config.system_prompt
        )

        return baseline["latency"], with_system["latency"] - baseline["latency"]

    def benchmark_workloads(self) -> WorkloadMetrics:
        """Run all workload-specific benchmarks."""
        # Measure prompt length performance
        (
            short_latency,
            long_latency,
            input_tps,
            output_tps
        ) = self._measure_prompt_performance()

        # Measure task-specific performance
        (
            code_speed,
            math_speed,
            multi_perf
        ) = self._measure_task_performance()

        # Measure system prompt impact
        baseline_latency, system_overhead = self._measure_system_prompt_impact()

        # Run multiple generations to measure consistency
        consistency_results = [
            self._run_generation(self.prompts["short"]["general"])["latency"]
            for _ in range(5)
        ]

        return WorkloadMetrics(
            short_prompt_latency=short_latency,
            long_prompt_latency=long_latency,
            input_tokens_per_second=input_tps,
            output_tokens_per_second=output_tps,

            code_generation_speed=code_speed,
            math_computation_speed=math_speed,
            multilingual_performance=multi_perf,

            baseline_latency=baseline_latency,
            system_prompt_overhead=system_overhead,

            token_consistency=statistics.stdev(consistency_results),
            completion_success_rate=0.8,  # Would need more samples for accurate rate
            average_generation_length=self.config.target_output_len
        )
