#!/usr/bin/env python3
"""
Experiment 13 Real Trace Integration Utilities

Provides utilities for collecting, converting, and validating real Azure/AWS
production traces for Experiment 13 validation.

Author: CAAE Research
Date: January 23, 2026
"""

import json
import hashlib
import statistics
from datetime import datetime
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, asdict
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# ============================================================================
# DATA MODELS
# ============================================================================


@dataclass
class AzureTrace:
    """Azure OpenAI API trace schema."""

    trace_id: str
    timestamp_utc: str
    deployment_id: str
    model: str

    # Request details
    prompt_tokens: int
    completion_tokens: int
    temperature: float

    # Performance metrics
    queue_latency_ms: float
    ttft_ms: float
    total_latency_ms: float
    kv_cache_hits: int
    kv_cache_misses: int

    # System state
    concurrent_requests: int
    gpu_utilization_percent: float
    gpu_memory_used_gb: float

    # Outcome
    status: str  # "success", "timeout", "error", "cancelled"
    sla_deadline_ms: Optional[float] = None
    sla_met: Optional[bool] = None
    region: str = "eastus"


@dataclass
class AWSTrace:
    """AWS Bedrock trace schema."""

    trace_id: str
    timestamp: str
    model_id: str

    # Request details
    input_tokens: int
    output_tokens: int
    temperature: float

    # Performance metrics
    ttft_ms: float
    total_inference_time_ms: float
    tokens_per_second: float
    batch_size: int

    # Cache
    kv_cache_hit_rate: float
    cache_evictions: int

    # System state
    concurrent_calls: int
    queue_wait_ms: float
    memory_used_mb: float

    # SLA
    sla_deadline_ms: float
    sla_met: bool
    region: str = "us-east-1"


@dataclass
class ExperimentRequest:
    """Unified request format for Experiment 13."""

    request_id: str
    model: str
    model_size_gb: float
    context_tokens: int
    output_tokens: int
    total_tokens: int
    deadline_ms: float
    sla_type: str  # "critical", "standard", "batch"
    is_preempted: bool
    priority: str  # "high", "normal"
    actual_latency_ms: float
    ttft_ms: Optional[float] = None
    region: str = "unknown"
    source: str = "unknown"  # "azure" or "aws"


# ============================================================================
# SANITIZATION UTILITIES
# ============================================================================


def sanitize_azure_trace(raw_trace: Dict) -> Dict:
    """Remove PII and secrets from Azure trace."""
    sanitized = {
        "trace_id": "TRACE_"
        + hashlib.sha256(raw_trace.get("trace_id", "unknown").encode()).hexdigest()[:8],
        "timestamp_utc": raw_trace.get("timestamp_utc"),
        "deployment_id": "DEPLOYMENT_"
        + hashlib.sha256(
            raw_trace.get("deployment_id", "unknown").encode()
        ).hexdigest()[:8],
        "model": raw_trace.get("model", "unknown"),
        "prompt_tokens": raw_trace.get("prompt_tokens", 0),
        "completion_tokens": raw_trace.get("completion_tokens", 0),
        "temperature": raw_trace.get("temperature", 0.7),
        "queue_latency_ms": raw_trace.get("queue_latency_ms", 0),
        "ttft_ms": raw_trace.get("ttft_ms", 0),
        "total_latency_ms": raw_trace.get("total_latency_ms", 0),
        "kv_cache_hits": raw_trace.get("kv_cache_hits", 0),
        "kv_cache_misses": raw_trace.get("kv_cache_misses", 0),
        "concurrent_requests": raw_trace.get("concurrent_requests", 0),
        "gpu_utilization_percent": raw_trace.get("gpu_utilization_percent", 0),
        "gpu_memory_used_gb": raw_trace.get("gpu_memory_used_gb", 0),
        "status": raw_trace.get("status", "unknown"),
        "sla_deadline_ms": raw_trace.get("sla_deadline_ms"),
        "sla_met": raw_trace.get("sla_met"),
        "region": "REGION_"
        + hashlib.sha256(raw_trace.get("region", "unknown").encode()).hexdigest()[:4],
    }
    return sanitized


def sanitize_aws_trace(raw_trace: Dict) -> Dict:
    """Remove PII and secrets from AWS Bedrock trace."""
    sanitized = {
        "trace_id": "TRACE_"
        + hashlib.sha256(raw_trace.get("trace_id", "unknown").encode()).hexdigest()[:8],
        "timestamp": raw_trace.get("timestamp"),
        "model_id": raw_trace.get("model_id", "unknown"),
        "input_tokens": raw_trace.get("input_tokens", 0),
        "output_tokens": raw_trace.get("output_tokens", 0),
        "temperature": raw_trace.get("temperature", 0.7),
        "ttft_ms": raw_trace.get("ttft_ms", 0),
        "total_inference_time_ms": raw_trace.get("total_inference_time_ms", 0),
        "tokens_per_second": raw_trace.get("tokens_per_second", 0),
        "batch_size": raw_trace.get("batch_size", 1),
        "kv_cache_hit_rate": raw_trace.get("kv_cache_hit_rate", 0),
        "cache_evictions": raw_trace.get("cache_evictions", 0),
        "concurrent_calls": raw_trace.get("concurrent_calls", 0),
        "queue_wait_ms": raw_trace.get("queue_wait_ms", 0),
        "memory_used_mb": raw_trace.get("memory_used_mb", 0),
        "sla_deadline_ms": raw_trace.get("sla_deadline_ms", 0),
        "sla_met": raw_trace.get("sla_met", True),
        "region": "REGION_"
        + hashlib.sha256(raw_trace.get("region", "unknown").encode()).hexdigest()[:4],
    }
    return sanitized


# ============================================================================
# CONVERSION UTILITIES
# ============================================================================

MODEL_SIZE_MAPPING = {
    # Azure models
    "gpt-4": 180,
    "gpt-4-32k": 180,
    "gpt-35-turbo": 38,
    "gpt-35-turbo-16k": 38,
    "text-davinci-003": 175,
    # AWS/Open source models
    "claude-2": 65,
    "claude-instant": 32,
    "llama-2-7b": 14,
    "llama-2-13b": 26,
    "llama-2-70b": 140,
    "llama-70b": 140,
    "anthropic.claude-2": 65,
    "anthropic.claude-instant": 32,
    "meta.llama2-70b": 140,
}


def estimate_model_size(model_name: str) -> float:
    """Map model name to size in GB."""
    # Try exact match first
    if model_name in MODEL_SIZE_MAPPING:
        return MODEL_SIZE_MAPPING[model_name]

    # Try substring matching
    for key, size in MODEL_SIZE_MAPPING.items():
        if key.lower() in model_name.lower():
            return size

    # Default estimate based on model family
    if "70b" in model_name.lower():
        return 140
    elif "13b" in model_name.lower():
        return 26
    elif "7b" in model_name.lower():
        return 14
    elif "gpt-4" in model_name.lower():
        return 180
    elif "gpt" in model_name.lower():
        return 60  # Default medium size
    elif "claude" in model_name.lower():
        return 65

    return 50  # Default fallback


def infer_sla_type(latency_ms: float, concurrent: int = 1) -> str:
    """Infer SLA type from latency and concurrency."""
    # Critical: < 1000ms (interactive use)
    # Standard: 1000-3000ms (normal use)
    # Batch: > 3000ms (async processing)

    if latency_ms < 1000 or concurrent < 5:
        return "critical"
    elif latency_ms < 3000:
        return "standard"
    else:
        return "batch"


def convert_azure_to_experiment_format(azure_trace: Dict) -> ExperimentRequest:
    """Convert Azure trace to experiment request format."""

    prompt_tokens = azure_trace.get("prompt_tokens", 0)
    completion_tokens = azure_trace.get("completion_tokens", 0)
    total_tokens = prompt_tokens + completion_tokens

    latency_ms = azure_trace.get("total_latency_ms", 1000)

    # Infer SLA deadline (typically 2-3x actual latency in Azure)
    deadline_ms = latency_ms * 2.5 if latency_ms > 0 else 3000

    model_name = azure_trace.get("model", "unknown")

    return ExperimentRequest(
        request_id=azure_trace.get("trace_id", "unknown"),
        model=model_name,
        model_size_gb=estimate_model_size(model_name),
        context_tokens=prompt_tokens,
        output_tokens=completion_tokens,
        total_tokens=total_tokens,
        deadline_ms=deadline_ms,
        sla_type=infer_sla_type(latency_ms, azure_trace.get("concurrent_requests", 1)),
        is_preempted=azure_trace.get("status") == "cancelled",
        priority="high" if deadline_ms < 1500 else "normal",
        actual_latency_ms=latency_ms,
        ttft_ms=azure_trace.get("ttft_ms"),
        region=azure_trace.get("region", "unknown"),
        source="azure",
    )


def convert_aws_to_experiment_format(aws_trace: Dict) -> ExperimentRequest:
    """Convert AWS trace to experiment request format."""

    input_tokens = aws_trace.get("input_tokens", 0)
    output_tokens = aws_trace.get("output_tokens", 0)
    total_tokens = input_tokens + output_tokens

    latency_ms = aws_trace.get("total_inference_time_ms", 1000)

    # Infer SLA deadline
    deadline_ms = latency_ms * 2.0 if latency_ms > 0 else 5000

    model_name = aws_trace.get("model_id", "unknown")

    return ExperimentRequest(
        request_id=aws_trace.get("trace_id", "unknown"),
        model=model_name,
        model_size_gb=estimate_model_size(model_name),
        context_tokens=input_tokens,
        output_tokens=output_tokens,
        total_tokens=total_tokens,
        deadline_ms=deadline_ms,
        sla_type=infer_sla_type(latency_ms, aws_trace.get("concurrent_calls", 1)),
        is_preempted=not aws_trace.get("sla_met", True),
        priority="high" if deadline_ms < 2000 else "normal",
        actual_latency_ms=latency_ms,
        ttft_ms=aws_trace.get("ttft_ms"),
        region=aws_trace.get("region", "unknown"),
        source="aws",
    )


# ============================================================================
# VALIDATION UTILITIES
# ============================================================================


def validate_trace_collection(
    traces: List[Dict], min_samples: int = 1000
) -> Tuple[bool, List[str]]:
    """Validate collected traces meet quality requirements."""

    if not traces:
        return False, ["No traces provided"]

    issues = []

    # Check minimum sample size
    if len(traces) < min_samples:
        issues.append(f"Only {len(traces)} traces (need {min_samples}+)")

    # Check required fields
    required_fields = [
        "trace_id",
        "timestamp_utc" if "timestamp_utc" in traces[0] else "timestamp",
        "total_latency_ms",
        "total_tokens",
    ]
    for field in required_fields:
        if not all(
            field in t or "latency_ms" in t or "inference_time_ms" in t for t in traces
        ):
            issues.append(f"Missing required field: {field}")

    # Check timestamp recency (should be recent)
    try:
        timestamps = []
        for t in traces:
            ts_str = t.get("timestamp_utc") or t.get("timestamp")
            if ts_str:
                ts = datetime.fromisoformat(ts_str.replace("Z", "+00:00"))
                timestamps.append(ts)

        if timestamps:
            oldest = min(timestamps)
            age_days = (datetime.now(oldest.tzinfo) - oldest).days

            if age_days > 30:
                issues.append(f"Traces are {age_days} days old (prefer < 7 days)")
    except Exception as e:
        issues.append(f"Could not parse timestamps: {e}")

    # Check latency distribution (should be realistic)
    try:
        latencies = []
        for t in traces:
            lat = t.get("total_latency_ms") or t.get("total_inference_time_ms")
            if lat:
                latencies.append(lat)

        if latencies:
            p50 = statistics.median(latencies)
            p99 = sorted(latencies)[int(len(latencies) * 0.99)]

            if p50 > 0:
                ratio = p99 / p50
                if ratio > 10 or ratio < 1.5:
                    issues.append(f"Unusual P99/P50 ratio: {ratio:.1f}x (expect 2-5x)")
    except Exception as e:
        issues.append(f"Could not analyze latency distribution: {e}")

    return len(issues) == 0, issues


def calculate_percentile(data: List[float], percentile: float) -> float:
    """Calculate percentile of data."""
    if not data:
        return 0.0

    sorted_data = sorted(data)
    index = int(len(sorted_data) * percentile / 100)
    return float(sorted_data[min(index, len(sorted_data) - 1)])


# ============================================================================
# COMPARISON UTILITIES
# ============================================================================


def compare_simulation_to_reality(
    simulated_latencies: List[float], actual_latencies: List[float]
) -> Dict:
    """Compare simulation fidelity against real traces."""

    if not actual_latencies:
        return {"error": "No actual latencies provided"}

    sim_mean = statistics.mean(simulated_latencies) if simulated_latencies else 0
    actual_mean = statistics.mean(actual_latencies)

    sim_p99 = (
        calculate_percentile(simulated_latencies, 99) if simulated_latencies else 0
    )
    actual_p99 = calculate_percentile(actual_latencies, 99)

    mean_error = (
        abs(sim_mean - actual_mean) / actual_mean * 100 if actual_mean > 0 else 0
    )
    p99_error = abs(sim_p99 - actual_p99) / actual_p99 * 100 if actual_p99 > 0 else 0

    fidelity_score = 100 - min(mean_error, 50)  # Cap error at 50%

    return {
        "mean_latency": {
            "simulated_ms": round(sim_mean, 2),
            "actual_ms": round(actual_mean, 2),
            "error_percent": round(mean_error, 2),
        },
        "p99_latency": {
            "simulated_ms": round(sim_p99, 2),
            "actual_ms": round(actual_p99, 2),
            "error_percent": round(p99_error, 2),
        },
        "fidelity_score": round(fidelity_score, 2),
        "fidelity_status": (
            "✅ Good"
            if fidelity_score > 80
            else "⚠️ Fair" if fidelity_score > 60 else "❌ Poor"
        ),
    }


# ============================================================================
# FILE I/O UTILITIES
# ============================================================================


def load_traces_from_file(filepath: str) -> List[Dict]:
    """Load traces from JSON file."""
    try:
        with open(filepath, "r") as f:
            data = json.load(f)

        if isinstance(data, dict) and "traces" in data:
            return data["traces"]
        elif isinstance(data, list):
            return data
        else:
            logger.error(f"Unexpected trace format in {filepath}")
            return []
    except Exception as e:
        logger.error(f"Failed to load traces: {e}")
        return []


def save_converted_requests(requests: List[ExperimentRequest], output_filepath: str):
    """Save converted requests to JSON file."""
    try:
        data = {
            "count": len(requests),
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "requests": [asdict(r) for r in requests],
        }

        Path(output_filepath).parent.mkdir(parents=True, exist_ok=True)
        with open(output_filepath, "w") as f:
            json.dump(data, f, indent=2)

        logger.info(f"Saved {len(requests)} converted requests to {output_filepath}")
    except Exception as e:
        logger.error(f"Failed to save requests: {e}")


# ============================================================================
# MAIN UTILITY FUNCTIONS
# ============================================================================


def process_azure_traces(input_file: str, output_file: str) -> int:
    """Load, sanitize, and convert Azure traces."""
    logger.info(f"Processing Azure traces from {input_file}")

    raw_traces = load_traces_from_file(input_file)
    logger.info(f"Loaded {len(raw_traces)} raw traces")

    # Sanitize
    sanitized = [sanitize_azure_trace(t) for t in raw_traces]
    logger.info(f"Sanitized {len(sanitized)} traces")

    # Validate
    valid, issues = validate_trace_collection(sanitized)
    if not valid:
        logger.warning(f"Validation issues: {issues}")

    # Convert
    converted = [convert_azure_to_experiment_format(t) for t in sanitized]

    # Save
    save_converted_requests(converted, output_file)

    return len(converted)


def process_aws_traces(input_file: str, output_file: str) -> int:
    """Load, sanitize, and convert AWS traces."""
    logger.info(f"Processing AWS traces from {input_file}")

    raw_traces = load_traces_from_file(input_file)
    logger.info(f"Loaded {len(raw_traces)} raw traces")

    # Sanitize
    sanitized = [sanitize_aws_trace(t) for t in raw_traces]
    logger.info(f"Sanitized {len(sanitized)} traces")

    # Validate
    valid, issues = validate_trace_collection(sanitized)
    if not valid:
        logger.warning(f"Validation issues: {issues}")

    # Convert
    converted = [convert_aws_to_experiment_format(t) for t in sanitized]

    # Save
    save_converted_requests(converted, output_file)

    return len(converted)


if __name__ == "__main__":
    # Example usage
    logger.info("Real Trace Integration Utilities loaded")
