import pytest
from log_parser import strip_ts, classify_line

# ── strip_ts ────────────────────────────────────────────────────────────────

def test_strip_ts_with_prefix():
    line = "[2026-06-22 14:29:49] 0.00.044.403 I srv  init: running without SSL"
    assert strip_ts(line) == "0.00.044.403 I srv  init: running without SSL"

def test_strip_ts_without_prefix():
    line = "0.00.044.403 I srv  init: running without SSL"
    assert strip_ts(line) == "0.00.044.403 I srv  init: running without SSL"

def test_strip_ts_child_line():
    line = "[2026-06-22 14:31:59] [41803] 0.00.041.673 I log_info: verbosity = 3"
    assert strip_ts(line) == "[41803] 0.00.041.673 I log_info: verbosity = 3"

def test_strip_ts_cmd_line():
    line = '[2026-06-22 14:31:59] [41803] cmd_child_to_router:state:{"state":"loading"}'
    assert strip_ts(line) == '[41803] cmd_child_to_router:state:{"state":"loading"}'

# ── classify_line ────────────────────────────────────────────────────────────

def test_classify_router_line():
    kind, port, content = classify_line("0.00.044.403 I srv  init: running without SSL")
    assert kind == "router"
    assert port is None
    assert content == "0.00.044.403 I srv  init: running without SSL"

def test_classify_child_line():
    kind, port, content = classify_line("[41803] 0.00.041.673 I log_info: verbosity = 3")
    assert kind == "child"
    assert port == 41803
    assert content == "0.00.041.673 I log_info: verbosity = 3"

def test_classify_child_cmd_line():
    payload = '{"state":"loading","payload":null}'
    kind, port, content = classify_line(f"[41803] cmd_child_to_router:state:{payload}")
    assert kind == "child"
    assert port == 41803
    assert content == f"cmd_child_to_router:state:{payload}"

def test_classify_router_proxy_line():
    line = "2.10.286.889 I srv  proxy_reques: proxying request to model Qwen3 on port 41803"
    kind, port, content = classify_line(line)
    assert kind == "router"
    assert port is None


# ── ServerState data model ───────────────────────────────────────────────────

from log_parser import ServerState, LoadingInfo, Metrics


def test_server_state_defaults():
    s = ServerState()
    assert s.state == "offline"
    assert s.model is None
    assert s.loading is None
    assert s.request_count == 0


def test_server_state_to_dict_offline():
    s = ServerState()
    d = s.to_dict()
    assert d["state"] == "offline"
    assert d["model"] is None
    assert d["loading"] is None
    assert d["metrics"]["gen_speed"] is None
    assert "timestamp" in d


def test_server_state_to_dict_loading():
    s = ServerState(
        state="loading",
        model="Qwen3",
        loading=LoadingInfo(stages=["text_model", "mmproj_model"], current="text_model", progress=0.5),
    )
    d = s.to_dict()
    assert d["state"] == "loading"
    assert d["loading"]["current"] == "text_model"
    assert d["loading"]["progress"] == 0.5


def test_server_state_to_dict_generating():
    s = ServerState(
        state="generating",
        model="Qwen3",
        metrics=Metrics(gen_speed=72.2, n_decoded=218),
        request_count=5,
    )
    d = s.to_dict()
    assert d["metrics"]["gen_speed"] == 72.2
    assert d["metrics"]["n_decoded"] == 218
    assert d["request_count"] == 5


from log_parser import LogParser

# Helper: feed a list of raw lines (with ts prefix), return final state name
def feed_lines(lines):
    p = LogParser()
    last = None
    for line in lines:
        result = p.feed(line)
        if result is not None:
            last = result
    return p.state.state, p.state


# ── Router startup ───────────────────────────────────────────────────────────

def test_router_startup():
    state, _ = feed_lines([
        "[2026-06-22 14:29:49] 0.00.076.839 I srv  llama_server: router server is listening on http://0.0.0.0:8082"
    ])
    assert state == "starting"


# ── Model spawn + loading progress ──────────────────────────────────────────

SPAWN_LINE = "[2026-06-22 14:31:59] 2.10.286.889 I srv          load: spawning server instance with name=unsloth/gemma-4-31B-it:UD-Q8_K_XL-recommended on port 41803"
LOADING_START = '[2026-06-22 14:31:59] [41803] cmd_child_to_router:state:{"state":"loading","payload":{"stages":["text_model","mmproj_model"],"current":"text_model","value":0.0}}'
LOADING_MID = '[2026-06-22 14:32:05] [41803] cmd_child_to_router:state:{"state":"loading","payload":{"stages":["text_model","mmproj_model"],"current":"text_model","value":0.5}}'
LOADING_DONE = '[2026-06-22 14:32:08] [41803] cmd_child_to_router:state:{"state":"loading","payload":{"stages":["text_model","mmproj_model"],"current":"text_model","value":1.0}}'
WARMUP_LINE = "[2026-06-22 14:32:10] [41803] 0.11.107.736 I common_init_from_params: warming up the model with an empty run - please wait ..."
READY_LINE = '[2026-06-22 14:32:13] [41803] cmd_child_to_router:state:{"state":"ready","payload":{"id":"unsloth/gemma-4-31B-it:UD-Q8_K_XL-recommended","aliases":[],"tags":[],"object":"model","created":1782127933,"owned_by":"llamacpp","meta":{}}}'
IDLE_LINE = "[2026-06-22 14:32:13] [41803] 0.13.230.405 I srv  update_slots: all slots are idle"


def test_model_loading_progress():
    p = LogParser()
    p.feed(SPAWN_LINE)
    p.feed(LOADING_START)
    assert p.state.state == "loading"
    assert p.state.loading.progress == 0.0
    assert p.state.loading.current == "text_model"

    p.feed(LOADING_MID)
    assert p.state.loading.progress == 0.5

    p.feed(LOADING_DONE)
    assert p.state.loading.progress == 1.0


def test_spawn_captures_model_name():
    p = LogParser()
    p.feed(SPAWN_LINE)
    assert p.state.model == "unsloth/gemma-4-31B-it:UD-Q8_K_XL-recommended"


def test_warmup_state():
    state, _ = feed_lines([SPAWN_LINE, LOADING_START, WARMUP_LINE])
    assert state == "warming_up"


def test_ready_transitions_to_idle():
    state, s = feed_lines([SPAWN_LINE, LOADING_START, WARMUP_LINE, READY_LINE])
    assert state == "idle"
    assert s.model == "unsloth/gemma-4-31B-it:UD-Q8_K_XL-recommended"


def test_idle_line():
    state, _ = feed_lines([SPAWN_LINE, LOADING_START, READY_LINE, IDLE_LINE])
    assert state == "idle"


# ── Request lifecycle ────────────────────────────────────────────────────────

PROXY_LINE = "[2026-06-22 14:32:13] 2.23.870.146 I srv  proxy_reques: proxying request to model unsloth/gemma-4-31B-it:UD-Q8_K_XL-recommended on port 41803"
LAUNCH_LINE = "[2026-06-22 14:34:20] [41803] 0.30.889.639 I slot launch_slot_: id  0 | task 0 | processing task, is_child = 0"
PROMPT_TIMING = "[2026-06-22 14:33:17] [41803] 0.56.093.911 I slot print_timing: id  0 | task 0 | prompt processing, n_tokens =    276, progress = 0.99, t =   4.90 s / 56.38 tokens per second"
GEN_TIMING = "[2026-06-22 14:33:20] [41803] 0.59.444.220 I slot print_timing: id  0 | task 0 | n_decoded =    156, tg =  51.94 t/s, tg_3s =  51.94 t/s"
RELEASE_LINE = "[2026-06-22 14:33:24] [41803] 1.02.743.373 I slot      release: id  0 | task 0 | stop processing: n_tokens = 606, truncated = 0"


def test_proxy_sets_waiting():
    state, _ = feed_lines([READY_LINE, IDLE_LINE, PROXY_LINE])
    assert state == "waiting"


def test_launch_sets_processing_prompt():
    state, _ = feed_lines([IDLE_LINE, PROXY_LINE, LAUNCH_LINE])
    assert state == "processing_prompt"


def test_prompt_timing_updates_metrics():
    _, s = feed_lines([IDLE_LINE, PROXY_LINE, LAUNCH_LINE, PROMPT_TIMING])
    assert s.state == "processing_prompt"
    assert abs(s.metrics.prompt_speed - 56.38) < 0.01
    assert s.metrics.prompt_tokens == 276


def test_gen_timing_sets_generating():
    _, s = feed_lines([IDLE_LINE, PROXY_LINE, LAUNCH_LINE, PROMPT_TIMING, GEN_TIMING])
    assert s.state == "generating"
    assert abs(s.metrics.gen_speed - 51.94) < 0.01
    assert s.metrics.n_decoded == 156


# Short prompts never emit a live "prompt processing" line — only the end-of-slot
# summary fires, after generation has already started.
PROMPT_SUMMARY = "[2026-06-22 14:34:24] [41983] 0.34.388.097 I slot print_timing: id  0 | task 0 | prompt eval time =     425.83 ms /    11 tokens (   38.71 ms per token,    25.83 tokens per second)"


def test_prompt_summary_captures_speed_during_generating():
    _, s = feed_lines([IDLE_LINE, PROXY_LINE, LAUNCH_LINE, GEN_TIMING, PROMPT_SUMMARY])
    assert s.state == "generating"  # state name unchanged
    assert abs(s.metrics.prompt_speed - 25.83) < 0.01
    assert s.metrics.prompt_tokens == 11
    assert abs(s.metrics.gen_speed - 51.94) < 0.01  # gen metrics preserved


def test_release_increments_request_count():
    _, s = feed_lines([IDLE_LINE, PROXY_LINE, LAUNCH_LINE, GEN_TIMING, RELEASE_LINE])
    assert s.state == "idle"
    assert s.request_count == 1
    assert s.metrics.gen_speed is None  # cleared on release


def test_release_clears_metrics():
    _, s = feed_lines([IDLE_LINE, PROXY_LINE, LAUNCH_LINE, GEN_TIMING, RELEASE_LINE])
    assert s.metrics.prompt_speed is None
    assert s.metrics.gen_speed is None
    assert s.metrics.n_decoded is None


# ── Sleeping ─────────────────────────────────────────────────────────────────

SLEEPING_LINE = '[2026-06-22 15:23:11] [36417] cmd_child_to_router:state:{"state":"sleeping","payload":null}'


def test_sleeping_state():
    state, _ = feed_lines([IDLE_LINE, SLEEPING_LINE])
    assert state == "sleeping"


# ── Unloading ────────────────────────────────────────────────────────────────

UNLOAD_LINE = "[2026-06-22 14:32:19] 2.30.337.683 I srv    unload_lru: models_max limit reached, removing LRU name=foo"


def test_unload_lru():
    state, _ = feed_lines([IDLE_LINE, UNLOAD_LINE])
    assert state == "unloading"


# ── Without ts prefix (single-instance mode) ────────────────────────────────

def test_no_ts_prefix():
    p = LogParser()
    p.feed("0.00.075.237 I srv  llama_server: router server is listening on http://0.0.0.0:8082")
    assert p.state.state == "starting"


def test_child_no_ts_prefix():
    p = LogParser()
    p.feed('[41803] cmd_child_to_router:state:{"state":"loading","payload":{"stages":["text_model"],"current":"text_model","value":0.3}}')
    assert p.state.state == "loading"
    assert p.state.loading.progress == 0.3