import random
import time
from pathlib import Path
from typing import Generator

import ollama

from .probe import HardwareProbe

# Load the word pool from the bundled text file (one word per line).
# These are common English words that each tokenize to ~1 token in most LLMs.
_WORD_POOL_PATH = Path(__file__).parent / "wordpool.txt"
_WORD_POOL = _WORD_POOL_PATH.read_text().split()


def _build_filler_prompt(target_tokens: int) -> str:
    """Build a pseudo-random prompt that maps to approximately *target_tokens*.

    Each word in _WORD_POOL is ~1 token for most LLM tokenizers.  We add a
    short instruction prefix so the model knows to reply briefly (we only
    care about measuring prompt-eval speed, not generation quality).
    """
    prefix = (
        "Below is a long passage. Read it carefully, then reply with "
        "only the word 'done'.\n\n"
    )
    # prefix ≈ 20 tokens, so fill the rest with random words
    fill_count = max(target_tokens - 20, 100)
    words = random.choices(_WORD_POOL, k=fill_count)
    return prefix + " ".join(words)

def _format_elapsed(seconds: float) -> str:
    """Format seconds into a human-readable duration."""
    if seconds < 60:
        return f"{round(seconds, 1)}s"
    minutes = int(seconds // 60)
    secs = round(seconds % 60)
    return f"{minutes}m {secs}s"


def _format_ttft(ms: float) -> str:
    """Format TTFT — ms when under 1s, seconds when over."""
    if ms < 1000:
        return f"{round(ms)}ms"
    return f"{round(ms / 1000, 1)}s"



try:
    from llama_cpp import Llama
    LLAMA_CPP_AVAILABLE = True
except ImportError:
    LLAMA_CPP_AVAILABLE = False


class StressTester:
    def __init__(self, model_name: str = ""):
        self.model_name = model_name
        self.probe = HardwareProbe()

    def iter_ollama_test(self, depth: int) -> Generator[dict, None, None]:
        """Yield one result dict per context step via Ollama API.

        Using a generator lets the caller collect partial results
        even if the run is interrupted (KeyboardInterrupt / OOM).
        """
        steps = [2048 * (4 ** i) for i in range(depth)]

        for size in steps:
            prompt = _build_filler_prompt(size)
            wall_start = time.perf_counter()

            try:
                response = ollama.generate(
                    model=self.model_name,
                    prompt=prompt,
                    stream=False,
                    options={
                        "num_ctx": size,
                        "num_predict": 32,
                    },
                )
                wall_elapsed = time.perf_counter() - wall_start

                stats = self.probe.get_stats()

                eval_count = response.get("eval_count", 0)
                eval_duration = response.get("eval_duration", 1)
                prompt_eval_duration = response.get("prompt_eval_duration", 0)

                tps = eval_count / (eval_duration / 1e9) if eval_duration else 0
                ttft = prompt_eval_duration / 1e6

                yield {
                    "context": f"{size // 1024}k",
                    "tps": round(tps),
                    "ttft": _format_ttft(ttft),
                    "elapsed": _format_elapsed(wall_elapsed),
                    "vram": stats["vram"],
                    "gpu_temp": stats["gpu_temp"],
                    "ram": stats["ram"],
                    "cpu": stats["cpu"],
                    "status": _verdict(tps),
                }
            except MemoryError:
                wall_elapsed = time.perf_counter() - wall_start
                stats = self.probe.get_stats()
                yield {
                    "context": f"{size // 1024}k",
                    "tps": 0,
                    "ttft": "FAIL",
                    "elapsed": _format_elapsed(wall_elapsed),
                    "vram": stats["vram"],
                    "gpu_temp": stats["gpu_temp"],
                    "ram": stats["ram"],
                    "cpu": stats["cpu"],
                    "status": "💀 OOM",
                }
                return  # Stop testing — larger contexts will also OOM
            except Exception:
                wall_elapsed = time.perf_counter() - wall_start
                stats = self.probe.get_stats()
                yield {
                    "context": f"{size // 1024}k",
                    "tps": 0,
                    "ttft": "FAIL",
                    "elapsed": _format_elapsed(wall_elapsed),
                    "vram": stats["vram"],
                    "gpu_temp": stats["gpu_temp"],
                    "ram": stats["ram"],
                    "cpu": stats["cpu"],
                    "status": "💀 Crash",
                }

    def run_ollama_test(self, depth: int) -> list[dict]:
        """Convenience wrapper — collects all results into a list."""
        return list(self.iter_ollama_test(depth))
    def iter_ollama_test_single(self, size: int, timeout: int = 300) -> Generator[dict, None, None]:
        """Test a single context size via Ollama. Yields exactly one result.

        Runs the Ollama call in a thread so Ctrl+C and --timeout work.
        """
        import concurrent.futures

        prompt = _build_filler_prompt(size)
        wall_start = time.perf_counter()

        def _call():
            return ollama.generate(
                model=self.model_name,
                prompt=prompt,
                stream=False,
                options={"num_ctx": size, "num_predict": 32},
            )

        try:
            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                future = pool.submit(_call)
                response = future.result(timeout=timeout)

            wall_elapsed = time.perf_counter() - wall_start
            stats = self.probe.get_stats()

            eval_count = response.get("eval_count", 0)
            eval_duration = response.get("eval_duration", 1)
            prompt_eval_duration = response.get("prompt_eval_duration", 0)

            tps = eval_count / (eval_duration / 1e9) if eval_duration else 0
            ttft = prompt_eval_duration / 1e6

            yield {
                "context": f"{size // 1024}k",
                "tps": round(tps),
                "ttft": _format_ttft(ttft),
                "elapsed": _format_elapsed(wall_elapsed),
                "vram": stats["vram"],
                "gpu_temp": stats["gpu_temp"],
                "ram": stats["ram"],
                "cpu": stats["cpu"],
                "status": _verdict(tps),
            }
        except concurrent.futures.TimeoutError:
            wall_elapsed = time.perf_counter() - wall_start
            stats = self.probe.get_stats()
            yield {
                "context": f"{size // 1024}k",
                "tps": 0,
                "ttft": "TIMEOUT",
                "elapsed": _format_elapsed(wall_elapsed),
                "vram": stats["vram"],
                "gpu_temp": stats["gpu_temp"],
                "ram": stats["ram"],
                "cpu": stats["cpu"],
                "status": "⏱️ Timeout",
            }
        except MemoryError:
            wall_elapsed = time.perf_counter() - wall_start
            stats = self.probe.get_stats()
            yield {
                "context": f"{size // 1024}k",
                "tps": 0,
                "ttft": "FAIL",
                "elapsed": _format_elapsed(wall_elapsed),
                "vram": stats["vram"],
                "gpu_temp": stats["gpu_temp"],
                "ram": stats["ram"],
                "cpu": stats["cpu"],
                "status": "💀 OOM",
            }
        except Exception:
            wall_elapsed = time.perf_counter() - wall_start
            stats = self.probe.get_stats()
            yield {
                "context": f"{size // 1024}k",
                "tps": 0,
                "ttft": "FAIL",
                "elapsed": _format_elapsed(wall_elapsed),
                "vram": stats["vram"],
                "gpu_temp": stats["gpu_temp"],
                "ram": stats["ram"],
                "cpu": stats["cpu"],
                "status": "💀 Crash",
            }

    def iter_direct_test(
        self, model_path: str, depth: int, n_gpu: int
    ) -> Generator[dict, None, None]:
        """Yield one result dict per context step for direct .gguf loading."""
        if not LLAMA_CPP_AVAILABLE:
            raise ImportError(
                "llama-cpp-python is not installed.\n"
                "Install it with: pip install stressllm[gguf]"
            )

        steps = [2048 * (4 ** i) for i in range(depth)]

        try:
            llm = Llama(
                model_path=model_path,
                n_gpu_layers=n_gpu,
                n_ctx=max(steps),
                verbose=False,
            )
        except MemoryError:
            stats = self.probe.get_stats()
            yield {
                "context": "LOAD",
                "tps": 0,
                "ttft": "FAIL",
                "elapsed": "—",
                "vram": stats["vram"],
                "gpu_temp": stats["gpu_temp"],
                "ram": stats["ram"],
                "cpu": stats["cpu"],
                "status": "💀 OOM on load",
            }
            return
        except Exception as e:
            stats = self.probe.get_stats()
            yield {
                "context": "LOAD",
                "tps": 0,
                "ttft": "FAIL",
                "elapsed": "—",
                "vram": stats["vram"],
                "gpu_temp": stats["gpu_temp"],
                "ram": stats["ram"],
                "cpu": stats["cpu"],
                "status": f"💀 Load failed: {e}",
            }
            return

        for size in steps:
            try:
                filler = _build_filler_prompt(size)
                tokens = llm.tokenize(filler.encode("utf-8"))

                start = time.perf_counter()
                llm.eval(tokens[: min(len(tokens), size)])
                end = time.perf_counter()
                wall_elapsed = end - start

                stats = self.probe.get_stats()
                elapsed_ms = round(wall_elapsed * 1000)

                yield {
                    "context": f"{size // 1024}k",
                    "tps": "Direct",
                    "ttft": f"{elapsed_ms}ms",
                    "elapsed": _format_elapsed(wall_elapsed),
                    "vram": stats["vram"],
                    "gpu_temp": stats["gpu_temp"],
                    "ram": stats["ram"],
                    "cpu": stats["cpu"],
                    "status": "✅ Loaded",
                }
            except MemoryError:
                stats = self.probe.get_stats()
                yield {
                    "context": f"{size // 1024}k",
                    "tps": 0,
                    "ttft": "FAIL",
                    "elapsed": "—",
                    "vram": stats["vram"],
                    "gpu_temp": stats["gpu_temp"],
                    "ram": stats["ram"],
                    "cpu": stats["cpu"],
                    "status": "💀 OOM",
                }
                return
            except Exception:
                stats = self.probe.get_stats()
                yield {
                    "context": f"{size // 1024}k",
                    "tps": 0,
                    "ttft": "FAIL",
                    "elapsed": "—",
                    "vram": stats["vram"],
                    "gpu_temp": stats["gpu_temp"],
                    "ram": stats["ram"],
                    "cpu": stats["cpu"],
                    "status": "💀 Crash",
                }

    def run_direct_test(self, model_path: str, depth: int, n_gpu: int) -> list[dict]:
        """Convenience wrapper — collects all results into a list."""
        return list(self.iter_direct_test(model_path, depth, n_gpu))


def _verdict(tps: float) -> str:
    """Map tokens-per-second to a human-readable status."""
    if tps > 15:
        return "✅ Smooth"
    if tps > 5:
        return "⚠️ Slowing"
    return "💀 Cliff"
