#!/usr/bin/env python3
"""
Performance optimizations for FlowTorch production.
- GPU prewarm
- Concurrency guard
- Model caching
"""

import os
import time
import threading
from typing import Optional
import logging

logger = logging.getLogger(__name__)


class ConcurrencyGuard:
    """
    Thread-safe concurrency limiter.
    Prevents resource exhaustion from concurrent requests.
    """

    def __init__(self, max_concurrent: int = 4):
        """
        Initialize concurrency guard.

        Args:
            max_concurrent: Max concurrent inference requests
        """
        self.max_concurrent = max_concurrent
        self.semaphore = threading.Semaphore(max_concurrent)
        self.active_count = 0
        self.lock = threading.Lock()

    def __enter__(self):
        """Acquire concurrency slot."""
        acquired = self.semaphore.acquire(blocking=True, timeout=30.0)

        if not acquired:
            raise RuntimeError("Concurrency limit timeout - server overloaded")

        with self.lock:
            self.active_count += 1

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Release concurrency slot."""
        self.semaphore.release()

        with self.lock:
            self.active_count -= 1

    def get_active_count(self) -> int:
        """Get current active request count."""
        with self.lock:
            return self.active_count


class GPUPrewarmer:
    """
    GPU model prewarm manager.
    Loads model into GPU memory on startup for faster first inference.
    """

    def __init__(self):
        self.prewarmed = False
        self.prewarm_time: Optional[float] = None
        self.device = None

    def prewarm(self, device_hint: str = "cpu") -> bool:
        """
        Prewarm GPU/CPU with model.

        Args:
            device_hint: 'cuda' or 'cpu'

        Returns:
            True if prewarm successful, False otherwise
        """
        if self.prewarmed:
            logger.info("Model already prewarmed")
            return True

        start = time.time()

        try:
            # Check if CUDA available
            if device_hint == "cuda":
                import torch
                if not torch.cuda.is_available():
                    logger.warning("CUDA requested but not available, falling back to CPU")
                    device_hint = "cpu"
                else:
                    self.device = "cuda"
                    # Prewarm CUDA
                    torch.zeros(1).cuda()
                    torch.cuda.synchronize()
            else:
                self.device = "cpu"

            # Load engine
            import importlib
            engine = importlib.import_module("artifacts.context_pack.flowtorch.flowtorch_engine")

            # Check for engine function
            if hasattr(engine, "run_prefill_decode"):
                engine_fn = engine.run_prefill_decode
            elif hasattr(engine, "generate"):
                engine_fn = engine.generate
            else:
                logger.warning("No engine function found, skipping prewarm")
                return False

            # Run dummy inference to warm caches
            logger.info(f"Prewarming {self.device} with dummy inference...")

            try:
                _ = engine_fn(
                    prompts=["test"],
                    max_new_tokens=1,
                    device_hint=self.device,
                    seed_envelope=None,
                    extras={}
                )
            except Exception as e:
                # Some engines don't support all params, try simpler call
                try:
                    _ = engine_fn(["test"], max_new_tokens=1)
                except Exception as e2:
                    logger.warning(f"Prewarm inference failed: {e2}")
                    return False

            elapsed = time.time() - start
            self.prewarm_time = elapsed
            self.prewarmed = True

            logger.info(f"Prewarm complete in {elapsed:.2f}s on {self.device}")
            return True

        except Exception as e:
            logger.error(f"Prewarm failed: {e}")
            return False

    def get_status(self) -> dict:
        """Get prewarm status."""
        return {
            "prewarmed": self.prewarmed,
            "prewarm_time_seconds": self.prewarm_time,
            "device": self.device,
        }


# Global instances
_concurrency_guard = ConcurrencyGuard(max_concurrent=4)
_gpu_prewarmer = GPUPrewarmer()


def get_concurrency_guard() -> ConcurrencyGuard:
    """Get global concurrency guard."""
    return _concurrency_guard


def get_gpu_prewarmer() -> GPUPrewarmer:
    """Get global GPU prewarmer."""
    return _gpu_prewarmer


def prewarm_on_startup(device_hint: str = "cpu"):
    """
    Run prewarm on server startup.

    Args:
        device_hint: 'cuda' or 'cpu' (from FLOWTORCH_DEVICE env)
    """
    prewarmer = get_gpu_prewarmer()
    prewarmer.prewarm(device_hint)
