"""Run command - Single inference and interactive chat."""

from __future__ import annotations

import asyncio
import logging
import os
import sys
from contextlib import contextmanager

import click

from ..utils.backend import resolve_backend_and_engine
from ..utils.config import init_default_config
from ..utils.console import get_console

DEFAULT_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
DEFAULT_TEMPERATURE = 0.7
DEFAULT_TOP_P = 0.9
DEFAULT_TOP_K = 40
DEFAULT_REPETITION_PENALTY = (
    1.5  # 1.0-1.3 tested insufficient for tiny models; 1.5 eliminates repetition
)


@contextmanager
def _runtime_log_mode(debug: bool):
    """Switch runtime log verbosity for inference paths."""
    if debug:
        yield
        return

    previous_root_level = logging.getLogger().level
    previous_levels: dict[str, int] = {}
    noisy_loggers = [
        "sagellm",
        "sagellm_core",
        "sagellm_backend",
        "sagellm_protocol",
        "transformers",
        "torch",
        "urllib3",
    ]
    for logger_name in noisy_loggers:
        logger = logging.getLogger(logger_name)
        previous_levels[logger_name] = logger.level
        logger.setLevel(logging.ERROR)

    logging.getLogger().setLevel(logging.ERROR)
    previous_disable = logging.root.manager.disable
    logging.disable(logging.CRITICAL)
    os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")

    try:
        yield
    finally:
        logging.disable(previous_disable)
        logging.getLogger().setLevel(previous_root_level)
        for logger_name, level in previous_levels.items():
            logging.getLogger(logger_name).setLevel(level)


@click.command()
@click.option("--prompt", "-p", type=str, required=True, help="Input prompt for inference")
@click.option("--model", "-m", type=str, default=None, help="Model name or path")
@click.option("--backend", type=str, default=None, help="Backend: cpu/cuda/ascend")
@click.option("--max-tokens", type=int, default=256, show_default=True, help="Max output tokens")
@click.option("--stream", is_flag=True, help="Enable streaming output")
# ==================== 采样参数（新增）====================
@click.option(
    "--temperature",
    "-t",
    type=float,
    default=DEFAULT_TEMPERATURE,
    show_default=True,
    help="Sampling temperature (0=greedy, 0.7=creative, 1.0=diverse)",
)
@click.option(
    "--top-p",
    type=float,
    default=DEFAULT_TOP_P,
    show_default=True,
    help="Nucleus sampling threshold (0.9 recommended)",
)
@click.option(
    "--top-k",
    type=int,
    default=DEFAULT_TOP_K,
    show_default=True,
    help="Top-k sampling (0=disabled, 40 recommended)",
)
@click.option(
    "--repetition-penalty",
    type=float,
    default=DEFAULT_REPETITION_PENALTY,
    show_default=True,
    help="Repetition penalty (>1 discourages repetition, default 1.5 tested safe for greedy mode)",
)
@click.option("--seed", type=int, default=None, help="Random seed for reproducibility")
@click.option("--debug", is_flag=True, help="Show verbose runtime logs and full traceback")
def run(
    prompt: str,
    model: str | None,
    backend: str | None,
    max_tokens: int,
    stream: bool,
    temperature: float,
    top_p: float,
    top_k: int,
    repetition_penalty: float,
    seed: int | None,
    debug: bool,
) -> None:
    """Run a single inference (like `ollama run`).

    \b
    Examples:
        sagellm run -p "What is 2+2?"
        sagellm run -p "Hello" -m Qwen/Qwen2.5-1.5B-Instruct
        sagellm run -p "Write a poem" --temperature 0.7 --top-p 0.9
        sagellm run -p "Code example" --stream --temperature 0.7
    """
    init_default_config()
    console = get_console()

    backend_kind, engine_kind = resolve_backend_and_engine(backend=backend)
    model_name = model or DEFAULT_MODEL

    console.print("\n[bold blue]🚀 sageLLM Inference[/bold blue]\n")
    console.print(f"  [cyan]Model:[/cyan] {model_name}")
    console.print(f"  [cyan]Backend:[/cyan] {backend_kind}")
    console.print(
        f"  [cyan]Prompt:[/cyan] {prompt[:50]}..."
        if len(prompt) > 50
        else f"  [cyan]Prompt:[/cyan] {prompt}"
    )
    console.print(f"  [cyan]Temperature:[/cyan] {temperature}")
    if temperature > 0:
        console.print(f"  [cyan]Top-p:[/cyan] {top_p}")
        console.print(f"  [cyan]Top-k:[/cyan] {top_k}")
    console.print(f"  [cyan]Debug:[/cyan] {'ON' if debug else 'OFF'}")
    console.print()

    try:

        async def _run() -> None:
            # Use unified LLMEngine (hardware-agnostic)
            from sagellm_core import LLMEngine, LLMEngineConfig
            from sagellm_protocol.sampling import SamplingParams
            from sagellm_protocol.types import Request

            provider_name = backend_kind

            config = LLMEngineConfig(
                model_path=model_name,
                backend_type=provider_name,
                max_new_tokens=max_tokens,
            )
            with _runtime_log_mode(debug):
                engine = LLMEngine(config)

            # 创建采样参数
            sampling_params = SamplingParams(
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                repetition_penalty=repetition_penalty,
                seed=seed,
            )

            with _runtime_log_mode(debug):
                with console.status("[bold cyan]⏳ 正在加载模型...[/bold cyan]", spinner="dots"):
                    await engine.start()

            try:
                if stream:
                    console.print("[green]📝 Output:[/green] ", end="")
                    stream_request = Request(
                        request_id="cli-stream",
                        trace_id="cli-stream",
                        model=model_name,
                        prompt="",
                        messages=[{"role": "user", "content": prompt}],
                        max_tokens=max_tokens,
                        stream=True,
                        temperature=temperature,
                        top_p=top_p,
                    )
                    async for event in engine.stream(stream_request):
                        if hasattr(event, "content") and event.event == "delta":
                            console.print(event.content, end="", style="cyan")
                    console.print()  # Newline
                else:
                    with _runtime_log_mode(debug):
                        with console.status(
                            "[bold cyan]🤔 正在思考...[/bold cyan]", spinner="dots"
                        ):
                            response = await engine.generate(
                                messages=[{"role": "user", "content": prompt}],
                                sampling_params=sampling_params,
                            )
                    console.print(f"[green]📝 Output:[/green]\n{response.output_text}\n")
                    console.print("[blue]📊 Metrics:[/blue]")
                    console.print(f"   TTFT: {response.metrics.ttft_ms:.1f} ms")
                    console.print(f"   Throughput: {response.metrics.throughput_tps:.1f} tokens/s")

            finally:
                await engine.stop()

        asyncio.run(_run())

    except ImportError as e:
        console.print(f"[red]❌ Import error: {e}[/red]")
        console.print("\n[yellow]💡 Install: pip install isagellm[/yellow]")
        sys.exit(1)
    except Exception as e:
        console.print(f"[red]❌ Error: {e}[/red]")
        if debug:
            import traceback

            console.print(f"[dim]{traceback.format_exc()}[/dim]")
        else:
            console.print("[dim]Use --debug to view full runtime logs and traceback.[/dim]")
        sys.exit(1)


@click.command()
@click.option("--model", "-m", type=str, default=None, help="Model name or path")
@click.option("--backend", type=str, default=None, help="Backend: cpu/cuda/ascend")
@click.option(
    "--max-tokens", type=int, default=256, show_default=True, help="Max output tokens per message"
)
@click.option("--temperature", "-t", type=float, default=DEFAULT_TEMPERATURE, show_default=True)
@click.option("--top-p", type=float, default=DEFAULT_TOP_P, show_default=True)
@click.option("--top-k", type=int, default=DEFAULT_TOP_K, show_default=True)
@click.option(
    "--repetition-penalty",
    type=float,
    default=DEFAULT_REPETITION_PENALTY,
    show_default=True,
)
@click.option("--debug", is_flag=True, help="Show verbose runtime logs and full traceback")
def chat(
    model: str | None,
    backend: str | None,
    max_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
    repetition_penalty: float,
    debug: bool,
) -> None:
    """Interactive chat session (like `ollama run` with interactive mode).

    \b
    Examples:
        sagellm chat                            # Start chat with default model
        sagellm chat -m Qwen/Qwen2.5-1.5B-Instruct   # Chat with specific model
        sagellm chat --backend cuda     # Force CUDA backend
    """
    init_default_config()
    console = get_console()

    backend_kind, engine_kind = resolve_backend_and_engine(backend=backend)
    model_name = model or DEFAULT_MODEL

    console.print("\n[bold blue]💬 sageLLM Interactive Chat[/bold blue]")
    console.print("=" * 60)
    console.print(f"  [cyan]Model:[/cyan] {model_name}")
    console.print(f"  [cyan]Backend:[/cyan] {backend_kind}")
    console.print(f"  [cyan]Debug:[/cyan] {'ON' if debug else 'OFF'}")
    console.print("  [dim]Type 'exit' or press Ctrl+C to quit[/dim]")
    console.print("=" * 60)
    console.print()

    try:
        from sagellm_core import LLMEngine, LLMEngineConfig
        from sagellm_protocol.sampling import SamplingParams

        provider_name = backend_kind

        config = LLMEngineConfig(
            model_path=model_name,
            backend_type=provider_name,
            max_new_tokens=max_tokens,
        )
        with _runtime_log_mode(debug):
            engine = LLMEngine(config)

        async def _chat() -> None:
            with _runtime_log_mode(debug):
                with console.status("[bold cyan]⏳ 正在加载模型...[/bold cyan]", spinner="dots"):
                    await engine.start()

            try:
                while True:
                    try:
                        # Get user input
                        user_input = console.input("[bold green]You:[/bold green] ").strip()

                        if not user_input:
                            continue

                        if user_input.lower() in ["exit", "quit", "q", "/exit", "/quit"]:
                            console.print("\n[yellow]👋 Goodbye![/yellow]")
                            break

                        # Generate response
                        console.print("[bold blue]Assistant:[/bold blue] ", end="")
                        sampling_params = SamplingParams(
                            max_tokens=max_tokens,
                            temperature=temperature,
                            top_p=top_p,
                            top_k=top_k,
                            repetition_penalty=repetition_penalty,
                        )
                        with _runtime_log_mode(debug):
                            with console.status(
                                "[bold cyan]🤔 正在思考...[/bold cyan]", spinner="dots"
                            ):
                                response = await engine.generate(
                                    messages=[{"role": "user", "content": user_input}],
                                    sampling_params=sampling_params,
                                )
                        console.print(response.output_text)
                        console.print()

                    except KeyboardInterrupt:
                        console.print("\n\n[yellow]👋 Goodbye![/yellow]")
                        break
                    except EOFError:
                        console.print("\n\n[yellow]👋 Goodbye![/yellow]")
                        break
                    except Exception as e:
                        console.print(f"[red]❌ Error: {e}[/red]")
                        console.print()

            finally:
                await engine.stop()

        asyncio.run(_chat())

    except ImportError as e:
        console.print(f"[red]❌ Import error: {e}[/red]")
        console.print("\n[yellow]💡 Install: pip install isagellm[/yellow]")
        sys.exit(1)
    except Exception as e:
        console.print(f"[red]❌ Error: {e}[/red]")
        if debug:
            import traceback

            console.print(f"[dim]{traceback.format_exc()}[/dim]")
        else:
            console.print("[dim]Use --debug to view full runtime logs and traceback.[/dim]")
        sys.exit(1)
