"""Live llama.cpp inference probe for the local Gemma 4 server.

Run with the server up on localhost:8080:
    .venv\\Scripts\\python.exe scripts\\llm_smoke_test.py
"""

from __future__ import annotations

import json
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from loosecanvas.llm_probe import (  # noqa: E402
    BARE_JSON_SCHEMA,
    JSON_OBJECT_ENUM_SCHEMA,
    OPENAI_WRAPPED_ENUM_SCHEMA,
    SCENEPLAN_CONTEXT,
    SCENEPLAN_SYSTEM_PROMPT,
    TOOL_MESSAGES,
    WEATHER_TOOL,
    assistant_message,
    chat,
    get_base_url,
    http_get,
    resolve_sceneplan_response_format,
    tool_result_roundtrip_messages,
    valid_enum_payload,
    valid_weather_tool_call,
    validate_sceneplan_payload,
)


def emit(line: str = "") -> None:
    sys.stdout.write(f"{line}\n")


def as_dict(value: object) -> dict[str, object]:
    return value if isinstance(value, dict) else {}


def as_list(value: object) -> list[object]:
    return value if isinstance(value, list) else []


def as_str(value: object) -> str:
    return value if isinstance(value, str) else ""


def hr(title: str) -> None:
    emit("\n" + "=" * 78)
    emit(title)
    emit("=" * 78)


def show_msg(resp: dict[str, object]) -> dict[str, object]:
    choices = as_list(resp.get("choices"))
    choice = as_dict(choices[0]) if choices else {}
    msg = as_dict(choice.get("message"))
    emit(f"  message keys     : {sorted(str(key) for key in msg)}")
    emit(f"  finish_reason    : {choice.get('finish_reason')}")
    reasoning_content = as_str(msg.get("reasoning_content"))
    if reasoning_content:
        emit(
            f"  reasoning_content: ({len(reasoning_content)} chars) {reasoning_content[:200]!r}"
        )
    emit(f"  content          : {as_str(msg.get('content'))[:400]!r}")
    tool_calls = as_list(msg.get("tool_calls"))
    if tool_calls:
        emit(f"  tool_calls       : {len(tool_calls)}")
        for index, tool_call_value in enumerate(tool_calls):
            tool_call = as_dict(tool_call_value)
            function = as_dict(tool_call.get("function"))
            emit(
                f"    [{index}] type={tool_call.get('type')!r} "
                f"name={function.get('name')!r} args={function.get('arguments')!r}"
            )
    timings = as_dict(resp.get("timings"))
    if timings:
        predicted_per_second = timings.get("predicted_per_second", 0.0)
        pps = (
            float(predicted_per_second)
            if isinstance(predicted_per_second, int | float)
            else 0.0
        )
        emit(
            f"  timings          : prompt_n={timings.get('prompt_n')} "
            f"predicted_n={timings.get('predicted_n')} pred/s={pps:.1f}"
        )
    return msg


# --------------------------------------------------------------------------- #
# 1. Health + props
# --------------------------------------------------------------------------- #
hr("1. /health and /props")
emit(f"  base_url         : {get_base_url()}")
emit(f"  health: {http_get('/health')}")

props = http_get("/props")
dgs = props.get("default_generation_settings", {})
params = dgs.get("params", {})
emit(f"  n_ctx            : {dgs.get('n_ctx')}")
emit(f"  model_path       : {props.get('model_path')}")
emit(f"  default stop     : {params.get('stop')}")
emit(f"  default samplers : {params.get('samplers')}")


# --------------------------------------------------------------------------- #
# 2. Plain completion
# --------------------------------------------------------------------------- #
hr("2. Plain completion (thinking as launched = ON)")
show_msg(chat([{"role": "user", "content": "Reply with a one-sentence greeting."}]))


# --------------------------------------------------------------------------- #
# 3. Thinking control
# --------------------------------------------------------------------------- #
hr("3. Plain completion with chat_template_kwargs={enable_thinking: false}")
show_msg(
    chat(
        [{"role": "user", "content": "Reply with a one-sentence greeting."}],
        chat_template_kwargs={"enable_thinking": False},
    )
)

hr("3b. Same request with reasoning_format='none'")
show_msg(
    chat(
        [{"role": "user", "content": "Reply with a one-sentence greeting."}],
        reasoning_format="none",
    )
)


# --------------------------------------------------------------------------- #
# 4. Negative control: bare response_format shape should fail on this build
# --------------------------------------------------------------------------- #
hr("4. Negative control — bare response_format json_schema x3")
bare_ok = 0
for index in range(3):
    resp = chat(
        [{"role": "user", "content": "Is the sky blue on a clear day? Answer."}],
        response_format=BARE_JSON_SCHEMA,
        max_tokens=20,
        chat_template_kwargs={"enable_thinking": False},
    )
    content = as_str(assistant_message(resp).get("content"))
    valid = valid_enum_payload(content)
    bare_ok += int(valid)
    emit(f"  [{index}] valid={valid} content={content!r}")
emit(f"  BARE JSON_SCHEMA ENFORCEMENT: {bare_ok}/3 valid (expected 0 on this build)")


# --------------------------------------------------------------------------- #
# 5. Verified OpenAI wrapper: enum enforcement
# --------------------------------------------------------------------------- #
hr("5. response_format OpenAI json_schema wrapper — enum enforcement x5")
wrapped_ok = 0
for index in range(5):
    resp = chat(
        [{"role": "user", "content": "Is the sky blue on a clear day? Answer."}],
        response_format=OPENAI_WRAPPED_ENUM_SCHEMA,
        max_tokens=20,
        chat_template_kwargs={"enable_thinking": False},
    )
    content = as_str(assistant_message(resp).get("content"))
    valid = valid_enum_payload(content)
    wrapped_ok += int(valid)
    emit(f"  [{index}] valid={valid} content={content!r}")
emit(f"  WRAPPED ENUM ENFORCEMENT: {wrapped_ok}/5 valid")


# --------------------------------------------------------------------------- #
# 6. Fallback json_object schema enforcement
# --------------------------------------------------------------------------- #
hr("6. response_format json_object + schema — enum enforcement x3")
json_object_ok = 0
for index in range(3):
    resp = chat(
        [{"role": "user", "content": "Is the sky blue on a clear day? Answer."}],
        response_format=JSON_OBJECT_ENUM_SCHEMA,
        max_tokens=20,
        chat_template_kwargs={"enable_thinking": False},
    )
    content = as_str(assistant_message(resp).get("content"))
    valid = valid_enum_payload(content)
    json_object_ok += int(valid)
    emit(f"  [{index}] valid={valid} content={content!r}")
emit(f"  JSON_OBJECT ENFORCEMENT: {json_object_ok}/3 valid")


# --------------------------------------------------------------------------- #
# 7. ScenePlan-like schema with the verified wrapper
# --------------------------------------------------------------------------- #
hr("7. ScenePlan-like schema via OpenAI wrapper x3")
sceneplan_response_format, sceneplan_source = resolve_sceneplan_response_format()
emit(f"  sceneplan source : {sceneplan_source}")
scene_ok = 0
started = time.time()
for index in range(3):
    resp = chat(
        [
            {"role": "system", "content": SCENEPLAN_SYSTEM_PROMPT},
            {"role": "user", "content": json.dumps(SCENEPLAN_CONTEXT)},
        ],
        response_format=sceneplan_response_format,
        max_tokens=600,
        chat_template_kwargs={"enable_thinking": False},
    )
    content = as_str(assistant_message(resp).get("content"))
    valid = validate_sceneplan_payload(content)
    scene_ok += int(valid)
    emit(f"  [{index}] valid={valid} content={content[:300]!r}")
emit(f"  SCENEPLAN SCHEMA: {scene_ok}/3 valid")
emit(f"  mean wall-clock/turn: {(time.time() - started) / 3:.2f}s")


# --------------------------------------------------------------------------- #
# 8. Tool calling — auto selection
# --------------------------------------------------------------------------- #
hr("8. Tool calling with tools + tool_choice='auto' x3")
tool_ok = 0
last_tool_msg: dict[str, object] | None = None
for index in range(3):
    resp = chat(
        TOOL_MESSAGES,
        tools=[WEATHER_TOOL],
        tool_choice="auto",
        parse_tool_calls=True,
        max_tokens=256,
        chat_template_kwargs={"enable_thinking": False},
    )
    msg = show_msg(resp)
    valid = valid_weather_tool_call(msg)
    tool_ok += int(valid)
    last_tool_msg = msg
    emit(f"  [{index}] valid_tool_call={valid}")
emit(f"  TOOL CALL GENERATION: {tool_ok}/3 valid")


# --------------------------------------------------------------------------- #
# 9. Tool roundtrip — feed tool result back to the model
# --------------------------------------------------------------------------- #
hr("9. Tool roundtrip with tool result")
if last_tool_msg and valid_weather_tool_call(last_tool_msg):
    follow_up_messages = TOOL_MESSAGES + tool_result_roundtrip_messages(last_tool_msg)
    roundtrip_resp = chat(
        follow_up_messages,
        tools=[WEATHER_TOOL],
        tool_choice="none",
        parse_tool_calls=True,
        max_tokens=128,
        chat_template_kwargs={"enable_thinking": False},
    )
    final_msg = show_msg(roundtrip_resp)
    final_content = as_str(final_msg.get("content")).lower()
    roundtrip_ok = "seattle" in final_content and (
        "17" in final_content
        or "light rain" in final_content
        or "celsius" in final_content
    )
    emit(f"  roundtrip_ok     : {roundtrip_ok}")
else:
    emit("  roundtrip skipped: no valid tool call was produced in section 8")

emit("\nDone.")