"""Tiiny SDK CLI."""

import argparse
import json
import os
import statistics
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional

from .client import TiinyClient, TiinyDevice


# Config file management
CONFIG_DIR = Path.home() / ".tiiny"
CONFIG_FILE = CONFIG_DIR / "config.json"

# Default prompt token sizes for benchmark
DEFAULT_TOKEN_SIZES = [
    256,
    512,
    1024,
    2048,
    4096,
    8192,
    16384,
    32768,
    65536,
]


def _ensure_config_dir() -> None:
    """Ensure the config directory exists."""
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)


def _load_config() -> Dict[str, Any]:
    """Load config from file, return empty dict if not exists."""
    if CONFIG_FILE.exists():
        try:
            with open(CONFIG_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except (json.JSONDecodeError, OSError):
            return {}
    return {}


def _save_config(config: Dict[str, Any]) -> None:
    """Save config to file."""
    _ensure_config_dir()
    with open(CONFIG_FILE, "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2)


def _get_config_value(key: str, default: Optional[str] = None) -> Optional[str]:
    """Get a value from config file."""
    config = _load_config()
    return config.get(key, default)


def _extract_delta_text(chunk: Dict[str, Any]) -> str:
    """Extract text content from a streaming chunk."""
    # Some providers put plain text at top-level
    top_level_text = chunk.get("text")
    if isinstance(top_level_text, str) and top_level_text:
        return top_level_text

    choices = chunk.get("choices") or []
    if not choices:
        return ""
    first_choice = choices[0] if isinstance(choices[0], dict) else {}

    # OpenAI-style delta
    delta = first_choice.get("delta") or {}
    content = delta.get("content")
    if isinstance(content, str) and content:
        return content

    # Some APIs stream in choices[].text
    text_field = first_choice.get("text")
    if isinstance(text_field, str) and text_field:
        return text_field

    # Some APIs send full message chunks
    message = first_choice.get("message") or {}
    message_content = message.get("content") if isinstance(message, dict) else None
    if isinstance(message_content, str) and message_content:
        return message_content

    return ""


def _extract_delta_reasoning(chunk: Dict[str, Any]) -> str:
    """Extract reasoning content from a streaming chunk.
    
    Used for models like DeepSeek R1 that output reasoning_content.
    """
    choices = chunk.get("choices") or []
    if not choices:
        return ""
    delta = choices[0].get("delta") or {}
    reasoning = delta.get("reasoning_content")
    return reasoning if isinstance(reasoning, str) else ""


def _extract_finish_reason(chunk: Dict[str, Any]) -> Optional[str]:
    """Extract finish_reason from a streaming chunk.
    
    Returns 'stop', 'length', or None if not finished.
    """
    choices = chunk.get("choices") or []
    if not choices:
        return None
    return choices[0].get("finish_reason")


def _extract_usage_tokens(chunk: Dict[str, Any]) -> Optional[int]:
    """Extract completion_tokens from usage field."""
    usage = chunk.get("usage")
    if not isinstance(usage, dict):
        return None
    completion_tokens = usage.get("completion_tokens")
    return completion_tokens if isinstance(completion_tokens, int) else None


def _extract_usage_info(chunk: Dict[str, Any]) -> Dict[str, int]:
    """Extract usage info from a streaming chunk if present."""
    usage = chunk.get("usage")
    if not isinstance(usage, dict):
        return {}

    result: Dict[str, int] = {}
    prompt_tokens = usage.get("prompt_tokens")
    completion_tokens = usage.get("completion_tokens")
    total_tokens = usage.get("total_tokens")

    if isinstance(prompt_tokens, int):
        result["prompt_tokens"] = prompt_tokens
    if isinstance(completion_tokens, int):
        result["completion_tokens"] = completion_tokens
    if isinstance(total_tokens, int):
        result["total_tokens"] = total_tokens
    return result


def _extract_timings(chunk: Dict[str, Any]) -> Optional[Dict[str, float]]:
    """Extract timing information from the final streaming chunk.
    
    Returns dict with:
    - prompt_tokens: number of prompt tokens
    - prompt_per_second: prefill speed (tokens/s)
    - predicted_per_second: decoding speed (tokens/s)
    - prompt_ms: prefill time in ms
    - predicted_ms: decoding time in ms
    """
    timings = chunk.get("timings")
    if not isinstance(timings, dict):
        return None

    def _pick_numeric(source: Dict[str, Any], keys: List[str]) -> Optional[float]:
        for key in keys:
            value = source.get(key)
            if isinstance(value, (int, float)):
                return float(value)
        return None

    def _pick_int(source: Dict[str, Any], keys: List[str]) -> Optional[int]:
        for key in keys:
            value = source.get(key)
            if isinstance(value, int):
                return value
            if isinstance(value, float):
                return int(value)
        return None
    
    usage = chunk.get("usage")
    prompt_tokens = None
    completion_tokens = None
    if isinstance(usage, dict):
        prompt_tokens = _pick_int(usage, ["prompt_tokens", "prompt tokens"])
        completion_tokens = _pick_int(usage, ["completion_tokens", "completion tokens"])

    if prompt_tokens is None:
        prompt_tokens = _pick_int(timings, ["prompt_n", "prompt n"])
    if completion_tokens is None:
        completion_tokens = _pick_int(timings, ["predicted_n", "predicted n", "decode_n", "decode n"])
    
    result = {}
    
    if prompt_tokens is not None:
        result["prompt_tokens"] = prompt_tokens
    if completion_tokens is not None:
        result["completion_tokens"] = completion_tokens
    
    # Prefill speed
    prompt_per_second = _pick_numeric(
        timings,
        ["prompt_per_second", "prompt per second", "prefill_per_second", "prefill per second"],
    )
    if prompt_per_second is not None:
        result["prompt_per_second"] = prompt_per_second
    
    # Decoding speed
    predicted_per_second = _pick_numeric(
        timings,
        ["predicted_per_second", "predicted per second", "decoding_per_second", "decoding per second"],
    )
    if predicted_per_second is not None:
        result["predicted_per_second"] = predicted_per_second
    
    # Times in ms
    prompt_ms = _pick_numeric(timings, ["prompt_ms", "prompt ms", "promptms", "prefill_ms", "prefill ms"])
    if prompt_ms is not None:
        result["prompt_ms"] = prompt_ms
    
    predicted_ms = _pick_numeric(timings, ["predicted_ms", "predicted ms", "predictedms", "decoding_ms", "decoding ms"])
    if predicted_ms is not None:
        result["predicted_ms"] = predicted_ms
    
    return result if result else None


def _estimate_tokens(text: str) -> int:
    """Estimate token count from text (rough approximation)."""
    if not text:
        return 0
    cjk_count = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff")
    ascii_word_count = len([token for token in text.split() if token.isascii()])
    other_char_count = max(0, len(text) - cjk_count)
    estimated = cjk_count + int(ascii_word_count * 1.2) + int(other_char_count / 4)
    return max(1, estimated)


def _parse_token_sizes(raw_values: Optional[List[str]]) -> Optional[List[int]]:
    """Parse token sizes from list of strings.

    Supports both space-separated and comma-separated forms:
    - --token-sizes 256 1024 2048
    - --token-sizes 256,1024,2048
    """
    if not raw_values:
        return None

    parsed: List[int] = []
    for value in raw_values:
        parts = [part.strip() for part in value.split(",") if part.strip()]
        for part in parts:
            token_size = int(part)
            if token_size <= 0:
                raise ValueError("token sizes must be positive integers")
            parsed.append(token_size)

    unique_sorted = sorted(set(parsed))
    return unique_sorted or None


def _build_prompt_for_tokens(target_tokens: int, topic: str) -> str:
    """Build a prompt that is approximately target_tokens long.
    
    Uses ~4 characters per token as rough estimate for English text.
    """
    effective_tokens = target_tokens
    if target_tokens >= 65536:
        effective_tokens = 50000

    # Approximate characters needed (4 chars per token for English)
    target_chars = effective_tokens * 4
    
    header = (
        f"Project Brief: {topic}\n"
        "Audience: engineering, product, QA\n"
        "Scope: summarize requirements, constraints, and decision points.\n\n"
    )

    blocks = [
        "Meeting Notes (2026-02-15)\n"
        "- Goal: improve latency, keep accuracy stable.\n"
        "- Risk: memory pressure at long context sizes.\n"
        "- Action: add streaming metrics to CLI.\n\n",
        "Requirements\n"
        "1) Support streaming chat and benchmarking.\n"
        "2) Report decode speed from API timings.\n"
        "3) Provide long-context coverage up to 64K tokens.\n\n",
        "Constraints\n"
        "- Must handle IPv6 addresses.\n"
        "- CLI should not store API keys in logs.\n\n",
        "Implementation Notes\n"
        "- Use OpenAI-compatible /v1/chat/completions.\n"
        "- Enable stream_options.include_usage.\n"
        "- Expect timings in the chunk after finish_reason.\n\n",
        "Telemetry Sample\n"
        "{\n"
        "  \"prompt_tokens\": 4096,\n"
        "  \"completion_tokens\": 32,\n"
        "  \"prompt_per_second\": 210.5,\n"
        "  \"predicted_per_second\": 28.7\n"
        "}\n\n",
        "FAQ\n"
        "Changelog\n"
        "- v0.1: baseline SDK\n"
        "- v0.2: CLI + benchmarks\n"
        "- v0.3: long-context test suite\n\n",
        "Appendix: Log excerpt\n"
        "[INFO] model start accepted\n"
        "[INFO] progress 100%\n"
        "[INFO] streaming complete\n\n",
    ]

    prompt = header
    idx = 0
    while len(prompt) < target_chars:
        prompt += blocks[idx % len(blocks)]
        idx += 1

    context_text = prompt[:target_chars]
    instruction = (
        "\n\nNow answer strictly in plain text with exactly 32 tokens. "
        "Do not include reasoning or analysis headers."
    )
    return context_text + instruction


def _resolve_api_key(device: TiinyDevice, api_key: Optional[str], master_password: Optional[str]) -> Optional[str]:
    if api_key:
        return api_key
    if master_password:
        return device.get_api_key(master_password)
    return None


def _make_client(args: argparse.Namespace, api_key: Optional[str], device: TiinyDevice) -> TiinyClient:
    base_url = args.base_url or device.get_url()
    return TiinyClient(api_key=api_key, base_url=base_url, timeout=args.timeout)


def _print_error(message: str) -> int:
    print(f"Error: {message}", file=sys.stderr)
    return 1


def _run_login(args: argparse.Namespace) -> int:
    """Run login command to save credentials to config file."""
    config = _load_config()

    # Default device IP
    DEFAULT_IP = "fd80:7:7:7::1"

    # Interactive mode - enter configuration terminal
    print("=" * 50)
    print("Tiiny Login Configuration")
    print("=" * 50)
    print("Configure your device credentials.")
    print("Port is fixed at 80 (no input needed).")
    print()

    # Device IP - use existing config or default
    current_ip = config.get("device_ip") or DEFAULT_IP
    ip_prompt = f"Device IP [{current_ip}]: "
    new_ip = input(ip_prompt).strip()
    # Save the IP (either new input or the current/default value)
    config["device_ip"] = new_ip if new_ip else current_ip

    # Master Password - use existing config or require input
    current_mp = config.get("master_password", "")
    if current_mp:
        mp_prompt = "Master Password [***]: "
    else:
        mp_prompt = "Master Password: "
    new_mp = input(mp_prompt).strip()
    if new_mp:
        config["master_password"] = new_mp
    elif current_mp:
        # Keep existing password
        config["master_password"] = current_mp
    else:
        print("Error: Master Password is required.")
        return 1

    # Always set port to 80
    config["device_port"] = 80

    _save_config(config)
    print()
    print("=" * 50)
    print("Configuration saved successfully!")
    print("=" * 50)
    print(f"Config file: {CONFIG_FILE}")
    print(f"  device_ip: {config.get('device_ip', 'not set')}")
    print(f"  device_port: 80")
    print(f"  master_password: ***")
    return 0


def _run_logout(args: argparse.Namespace) -> int:
    """Run logout command to remove config file."""
    if CONFIG_FILE.exists():
        CONFIG_FILE.unlink()
        print(f"Removed {CONFIG_FILE}")
    else:
        print("No config file found.")
    return 0


def _run_config_show(args: argparse.Namespace) -> int:
    """Show current config."""
    config = _load_config()
    if not config:
        print(f"No config file found at {CONFIG_FILE}")
        print("Run 'tiiny login' to create one.")
        return 0

    print(f"Config file: {CONFIG_FILE}")
    print("Values:")
    for key, value in config.items():
        if key in ("master_password", "api_key"):
            print(f"  {key}: ***")
        else:
            print(f"  {key}: {value}")
    return 0


def _run_chat(args: argparse.Namespace) -> int:
    device = TiinyDevice(device_ip=args.device_ip, device_port=args.device_port)
    api_key = _resolve_api_key(device, args.api_key, args.master_password)

    if not args.skip_start:
        if not api_key and not args.master_password:
            return _print_error("run mode starts the model by default, please provide --api-key or --master-password, or use --skip-start")
        started = device.start_model(args.model, api_key=api_key, master_password=args.master_password)
        if not started:
            return _print_error(f"Failed to start model: {args.model}")
        print(f"Model started: {args.model}")

    client = _make_client(args, api_key, device)
    history: List[Dict[str, str]] = []

    if args.system_prompt:
        history.append({"role": "system", "content": args.system_prompt})

    print("Entering chat mode, type /exit or /quit to exit.")
    while True:
        try:
            user_text = input("you> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye")
            return 0

        if not user_text:
            continue
        if user_text in {"/exit", "/quit"}:
            print("Bye")
            return 0
        if user_text == "/clear":
            history = [{"role": "system", "content": args.system_prompt}] if args.system_prompt else []
            print("Context cleared.")
            continue

        history.append({"role": "user", "content": user_text})
        print("assistant> ", end="", flush=True)
        assistant_text_parts: List[str] = []
        reasoning_text_parts: List[str] = []
        in_reasoning = False
        timings_data: Optional[Dict[str, float]] = None
        usage_info: Dict[str, int] = {}

        try:
            saw_finish = False
            stream = client.chat.completions.create(
                model=args.model,
                messages=history,
                stream=True,
                stream_options={"include_usage": True},
                temperature=args.temperature,
                max_tokens=args.max_tokens,
            )
            for chunk in stream:
                if not isinstance(chunk, dict):
                    continue

                # Keep latest timings/usage from stream chunks
                chunk_timings = _extract_timings(chunk)
                if chunk_timings:
                    timings_data = chunk_timings
                chunk_usage = _extract_usage_info(chunk)
                if chunk_usage:
                    usage_info.update(chunk_usage)
                
                # Handle reasoning content (e.g., DeepSeek R1)
                reasoning_delta = _extract_delta_reasoning(chunk)
                if reasoning_delta:
                    if not in_reasoning:
                        in_reasoning = True
                        print("\n[Reasoning] ", end="", flush=True)
                    reasoning_text_parts.append(reasoning_delta)
                    print(reasoning_delta, end="", flush=True)
                
                # Handle regular content
                delta = _extract_delta_text(chunk)
                if delta:
                    if in_reasoning:
                        # End reasoning section, start regular content
                        in_reasoning = False
                        print("\n[Response] ", end="", flush=True)
                    assistant_text_parts.append(delta)
                    print(delta, end="", flush=True)

                # Check for finish_reason after consuming current content
                finish_reason = _extract_finish_reason(chunk)
                if finish_reason:
                    saw_finish = True

                # Some servers send usage/timings in the next chunk after finish_reason
                if saw_finish and timings_data is not None:
                    break
            print()
            
            # Show speed info after response
            prompt_tokens = usage_info.get("prompt_tokens")
            completion_tokens = usage_info.get("completion_tokens")
            prefill = timings_data.get("prompt_per_second") if timings_data else None
            decoding = timings_data.get("predicted_per_second") if timings_data else None

            if decoding is not None:
                if completion_tokens is not None:
                    print(f"[Speed: decoding {int(completion_tokens)} tok @ {decoding:.2f} tok/s]")
                else:
                    print(f"[Speed: decoding {decoding:.2f} tok/s]")
            else:
                print("[Speed: n/a, API stream timings missing]")
        except KeyboardInterrupt:
            print("\nResponse interrupted.")
        except Exception as exc:  # noqa: BLE001
            print(f"\n[request failed] {exc}")

        assistant_text = "".join(assistant_text_parts).strip()
        if assistant_text:
            history.append({"role": "assistant", "content": assistant_text})


@dataclass
class BenchResult:
    """Benchmark result for a single test case."""
    target_tokens: int  # Target prompt token count
    actual_prompt_tokens: Optional[int]  # Actual prompt tokens from API
    completion_tokens: int
    prefill_speed: Optional[float]  # tokens/s from API timings
    decoding_speed: Optional[float]  # tokens/s from API timings
    prefill_time_ms: Optional[float]  # Prefill time in ms
    decoding_time_ms: Optional[float]  # Decoding time in ms
    total_time: float  # Total time measured locally


@dataclass
class BenchSummary:
    """Summary of benchmark results for a specific token size."""
    target_tokens: int
    results: List[BenchResult] = field(default_factory=list)

    @property
    def avg_prefill_speed(self) -> Optional[float]:
        speeds = [r.prefill_speed for r in self.results if r.prefill_speed is not None]
        return statistics.mean(speeds) if speeds else None

    @property
    def avg_decoding_speed(self) -> Optional[float]:
        speeds = [r.decoding_speed for r in self.results if r.decoding_speed is not None]
        return statistics.mean(speeds) if speeds else None

    @property
    def avg_total_time(self) -> float:
        return statistics.mean([r.total_time for r in self.results])


def _run_one_bench(
    client: TiinyClient,
    model: str,
    prompt: str,
    target_tokens: int,
    temperature: float,
    max_tokens: Optional[int],
    show_stream: bool = True,
) -> BenchResult:
    """Run a single benchmark test."""
    start_time = time.perf_counter()
    first_token_time: Optional[float] = None
    output_parts: List[str] = []
    reasoning_parts: List[str] = []
    timings_data: Optional[Dict[str, float]] = None
    usage_info: Dict[str, int] = {}

    if show_stream:
        print(f"\n[Benchmark] Target: {target_tokens} tokens")
        print("[Input Prompt]")
        print(prompt)
        print("[Streaming] ", end="", flush=True)

    stream: Iterator[Dict[str, Any]] = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        stream=True,
        stream_options={"include_usage": True},
        temperature=temperature,
        max_tokens=max_tokens,
    )

    saw_finish = False
    for chunk in stream:
        if not isinstance(chunk, dict):
            continue

        # Extract timing info if available (some APIs send it in final chunk with content)
        chunk_timings = _extract_timings(chunk)
        if chunk_timings:
            timings_data = chunk_timings
        chunk_usage = _extract_usage_info(chunk)
        if chunk_usage:
            usage_info.update(chunk_usage)

        reasoning_delta = _extract_delta_reasoning(chunk)
        if reasoning_delta:
            if first_token_time is None:
                first_token_time = time.perf_counter()
                if show_stream:
                    print(f"[TTFT: {first_token_time - start_time:.3f}s] ", end="", flush=True)
            reasoning_parts.append(reasoning_delta)
            if show_stream:
                print(reasoning_delta, end="", flush=True)
        
        delta = _extract_delta_text(chunk)
        if delta:
            if first_token_time is None:
                first_token_time = time.perf_counter()
                if show_stream:
                    print(f"[TTFT: {first_token_time - start_time:.3f}s] ", end="", flush=True)
            output_parts.append(delta)
            if show_stream:
                print(delta, end="", flush=True)

        # Check for finish_reason after consuming current content
        finish_reason = _extract_finish_reason(chunk)
        if finish_reason:
            saw_finish = True

        # Some servers send usage/timings in the next chunk after finish_reason
        if saw_finish and timings_data is not None:
            break

    end_time = time.perf_counter()
    total_time = end_time - start_time

    # Extract metrics from timings data
    actual_prompt_tokens = None
    completion_tokens = usage_info.get("completion_tokens", 0)
    prefill_speed = None
    decoding_speed = None
    prefill_time_ms = None
    decoding_time_ms = None

    if timings_data:
        actual_prompt_tokens = timings_data.get("prompt_tokens")
        completion_tokens = int(timings_data.get("completion_tokens", completion_tokens))
        prefill_speed = timings_data.get("prompt_per_second")
        decoding_speed = timings_data.get("predicted_per_second")
        prefill_time_ms = timings_data.get("prompt_ms")
        decoding_time_ms = timings_data.get("predicted_ms")

    if actual_prompt_tokens is None:
        actual_prompt_tokens = usage_info.get("prompt_tokens")

    if show_stream:
        print(f"\n[Done] Output: {len(''.join(reasoning_parts + output_parts))} chars, {completion_tokens} tokens")
        if prefill_speed is not None:
            if prefill_time_ms is not None:
                print(f"  Prefill: {prefill_time_ms:.1f}ms @ {prefill_speed:.1f} tok/s")
            else:
                print(f"  Prefill: {prefill_speed:.1f} tok/s")
        else:
            print("  Prefill: n/a (API timings missing)")
        if decoding_speed is not None:
            if decoding_time_ms is not None:
                print(f"  Decoding: {decoding_time_ms:.1f}ms @ {decoding_speed:.2f} tok/s")
            else:
                print(f"  Decoding: {decoding_speed:.2f} tok/s")
        else:
            print("  Decoding: n/a (API timings missing)")
        print(f"  Total: {total_time:.3f}s")

    return BenchResult(
        target_tokens=target_tokens,
        actual_prompt_tokens=actual_prompt_tokens,
        completion_tokens=completion_tokens,
        prefill_speed=prefill_speed,
        decoding_speed=decoding_speed,
        prefill_time_ms=prefill_time_ms,
        decoding_time_ms=decoding_time_ms,
        total_time=total_time,
    )


def _run_bench(args: argparse.Namespace) -> int:
    """Run benchmark with fixed token sizes."""
    device = TiinyDevice(device_ip=args.device_ip, device_port=args.device_port)
    api_key = _resolve_api_key(device, args.api_key, args.master_password)

    if not args.skip_start:
        if not api_key and not args.master_password:
            return _print_error("bench mode starts the model by default, please provide --api-key or --master-password, or use --skip-start")
        started = device.start_model(args.model, api_key=api_key, master_password=args.master_password)
        if not started:
            return _print_error(f"Failed to start model: {args.model}")
        print(f"Model started: {args.model}")

    client = _make_client(args, api_key, device)
    
    # Use fixed token sizes or custom sizes from args
    try:
        token_sizes = _parse_token_sizes(args.token_sizes) or DEFAULT_TOKEN_SIZES
    except ValueError as exc:
        return _print_error(f"invalid --token-sizes: {exc}")

    print("=" * 80)
    print("Tiiny Benchmark - Streaming Inference Performance Test")
    print("=" * 80)
    print(f"Model: {args.model}")
    print(f"Prompt token sizes: {token_sizes}")
    print(f"Rounds per size: {args.rounds}")
    print(f"Max output tokens: {args.max_tokens}")
    print(f"Topic: {args.topic}")
    print("=" * 80)

    # Group results by token size
    summaries: Dict[int, BenchSummary] = {ts: BenchSummary(target_tokens=ts) for ts in token_sizes}
    all_results: List[BenchResult] = []
    total_cases = len(token_sizes) * args.rounds
    current_case = 0

    for target_tokens in token_sizes:
        for round_idx in range(args.rounds):
            current_case += 1
            print(f"\n[{current_case}/{total_cases}] Testing {target_tokens} tokens (round {round_idx + 1}/{args.rounds})")
            prompt = _build_prompt_for_tokens(target_tokens, args.topic)
            
            try:
                result = _run_one_bench(
                    client=client,
                    model=args.model,
                    prompt=prompt,
                    target_tokens=target_tokens,
                    temperature=args.temperature,
                    max_tokens=args.max_tokens,
                    show_stream=not args.quiet,
                )
            except Exception as exc:  # noqa: BLE001
                return _print_error(f"Benchmark failed ({target_tokens} tokens): {exc}")

            all_results.append(result)
            summaries[target_tokens].results.append(result)

    if not all_results:
        return _print_error("No benchmark results obtained")

    # Print detailed results table
    print("\n" + "=" * 80)
    print("DETAILED RESULTS")
    print("=" * 80)
    print(f"{'Target':<10} {'Output':<10} {'Prefill':<15} {'Decoding':<15} {'Total(s)':<10}")
    print(f"{'(tokens)':<10} {'(tokens)':<10} {'(tok/s)':<15} {'(tok/s)':<15} {'':<10}")
    print("-" * 80)

    for result in all_results:
        prefill = f"{result.prefill_speed:.1f}" if result.prefill_speed is not None else "n/a"
        decoding = f"{result.decoding_speed:.2f}" if result.decoding_speed is not None else "n/a"
        print(
            f"{result.target_tokens:<10} {result.completion_tokens:<10} "
            f"{prefill:<15} {decoding:<15} {result.total_time:<10.3f}"
        )

    # Print summary by token size
    print("\n" + "=" * 80)
    print("SUMMARY BY CONTEXT LENGTH")
    print("=" * 80)
    print(f"{'Tokens':<12} {'Avg Prefill':<18} {'Avg Decoding':<18} {'Avg Total(s)':<15}")
    print("-" * 80)

    for target_tokens in token_sizes:
        summary = summaries[target_tokens]
        if summary.results:
            avg_prefill = summary.avg_prefill_speed
            avg_decoding = summary.avg_decoding_speed
            avg_total = summary.avg_total_time

            prefill_str = f"{avg_prefill:.1f}" if avg_prefill is not None else "n/a"
            decoding_str = f"{avg_decoding:.2f}" if avg_decoding is not None else "n/a"
            print(
                f"{target_tokens:<12} {prefill_str:<18} {decoding_str:<18} {avg_total:<15.3f}"
            )

    # Print overall summary
    print("\n" + "=" * 80)
    print("OVERALL SUMMARY")
    print("=" * 80)

    prefill_values = [r.prefill_speed for r in all_results if r.prefill_speed is not None]
    decoding_values = [r.decoding_speed for r in all_results if r.decoding_speed is not None]
    total_values = [r.total_time for r in all_results]

    print(f"Total test cases: {len(all_results)}")

    if prefill_values:
        print(f"\nPrefill Speed (from API timings):")
        print(f"  Average: {statistics.mean(prefill_values):.1f} tok/s")
        print(f"  Median:  {statistics.median(prefill_values):.1f} tok/s")
        print(f"  Min:     {min(prefill_values):.1f} tok/s")
        print(f"  Max:     {max(prefill_values):.1f} tok/s")

    if decoding_values:
        print(f"\nDecoding Speed (from API timings):")
        print(f"  Average: {statistics.mean(decoding_values):.2f} tok/s")
        print(f"  Median:  {statistics.median(decoding_values):.2f} tok/s")
        print(f"  Min:     {min(decoding_values):.2f} tok/s")
        print(f"  Max:     {max(decoding_values):.2f} tok/s")

    print(f"\nTotal Time:")
    print(f"  Average: {statistics.mean(total_values):.3f}s")
    print(f"  Median:  {statistics.median(total_values):.3f}s")

    print("\n" + "=" * 80)
    print("Benchmark completed successfully!")
    print("=" * 80)

    return 0


def _run_models(args: argparse.Namespace) -> int:
    device = TiinyDevice(device_ip=args.device_ip, device_port=args.device_port)
    api_key = _resolve_api_key(device, args.api_key, args.master_password)

    if args.running:
        client = _make_client(args, api_key, device)
        models = client.get_running_models()
        title = "Running models"
    else:
        models = device.get_models(api_key=api_key, master_password=args.master_password)
        title = "Available models"

    print(f"{title} ({len(models)}):")
    for model in models:
        print(f"- {model}")
    return 0


def _get_default_from_config(key: str, env_default: str) -> str:
    """Get default value from config file, fallback to env variable."""
    config_value = _get_config_value(key)
    if config_value:
        return config_value
    return os.getenv(env_default, "")


def _get_default_device_ip() -> str:
    """Get default device IP from config or env."""
    config_value = _get_config_value("device_ip")
    if config_value:
        return config_value
    return os.getenv("TIINY_DEVICE_IP", "fd80:7:7:7::1")


def _get_default_device_port() -> int:
    """Get default device port from config or env."""
    config_value = _get_config_value("device_port")
    if config_value:
        try:
            return int(config_value)
        except (ValueError, TypeError):
            pass
    return int(os.getenv("TIINY_DEVICE_PORT", "80"))


def _get_default_master_password() -> str:
    """Get default master password from config or env."""
    config_value = _get_config_value("master_password")
    if config_value:
        return config_value
    return os.getenv("TIINY_MASTER_PASSWORD", "")


def _get_default_api_key() -> str:
    """Get default API key from config or env."""
    config_value = _get_config_value("api_key")
    if config_value:
        return config_value
    return os.getenv("TIINY_API_KEY", "")


def _get_default_base_url() -> str:
    """Get default base URL from config or env."""
    config_value = _get_config_value("base_url")
    if config_value:
        return config_value
    return os.getenv("TIINY_BASE_URL", "")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="tiiny", description="Tiiny SDK command line tool")
    parser.add_argument("--device-ip", default=_get_default_device_ip(), help="Tiiny device IP")
    parser.add_argument("--device-port", type=int, default=_get_default_device_port(), help="Tiiny device port")
    parser.add_argument("--api-key", default=_get_default_api_key(), help="API key")
    parser.add_argument("--master-password", default=_get_default_master_password(), help="Master password")
    parser.add_argument("--base-url", default=_get_default_base_url(), help="OpenAI-compatible base URL, defaults to device URL")
    parser.add_argument("--timeout", type=float, default=float(os.getenv("TIINY_TIMEOUT", "3600")), help="Request timeout in seconds")

    subparsers = parser.add_subparsers(dest="command", required=True)

    # Login command
    login_parser = subparsers.add_parser("login", help="Save credentials to ~/.tiiny/config.json")
    login_parser.set_defaults(func=_run_login)

    # Logout command
    logout_parser = subparsers.add_parser("logout", help="Remove saved credentials")
    logout_parser.set_defaults(func=_run_logout)

    # Config command
    config_parser = subparsers.add_parser("config", help="Show current configuration")
    config_parser.set_defaults(func=_run_config_show)

    # Run command
    run_parser = subparsers.add_parser("run", help="Start model and open interactive terminal chat")
    run_parser.add_argument("model", help="Model ID")
    run_parser.add_argument("--skip-start", action="store_true", help="Do not call model start before chat")
    run_parser.add_argument("--system-prompt", default="", help="System prompt used at chat start")
    run_parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature")
    run_parser.add_argument("--max-tokens", type=int, default=None, help="Max tokens per reply")
    run_parser.set_defaults(func=_run_chat)

    # Bench command
    bench_parser = subparsers.add_parser("bench", help="Benchmark streaming inference speed")
    bench_parser.add_argument("model", help="Model ID")
    bench_parser.add_argument("--skip-start", action="store_true", help="Do not call model start before benchmark")
    bench_parser.add_argument("--token-sizes", type=str, nargs="+", default=None,
                              help="Fixed prompt token sizes, e.g. '256 1024 2048' or '256,1024,2048'")
    bench_parser.add_argument("--rounds", type=int, default=1, help="Rounds per token size")
    bench_parser.add_argument("--topic", default="large language model inference performance", help="Benchmark prompt topic")
    bench_parser.add_argument("--temperature", type=float, default=0.2, help="Sampling temperature")
    bench_parser.add_argument("--max-tokens", type=int, default=32, help="Max output tokens for each case (default: 32)")
    bench_parser.add_argument("--quiet", action="store_true", help="Suppress streaming output logs")
    bench_parser.set_defaults(func=_run_bench)

    # Models command
    models_parser = subparsers.add_parser("models", help="List models")
    models_parser.add_argument("--running", action="store_true", help="List currently running models via OpenAI API")
    models_parser.set_defaults(func=_run_models)

    return parser


def main(argv: Optional[List[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())
