src/mini_transformer/apps/chainlit_app.py

from __future__ import annotations

import os
import shlex
from typing import Any

import chainlit as cl
from hydra.core.global_hydra import GlobalHydra
from omegaconf import DictConfig

from mini_transformer.inference import run_inference
from mini_transformer.model_loader import (
    MODEL_NAME_ENV,
    compose_model_config,
    ensure_models_root,
    list_model_names,
)

JOB_NAME = "chainlit_app"
MODEL_SESSION_KEY = "selected_model"
GENERATION_SESSION_KEY = "generation_overrides"
HAS_ACTION_PROMPT = hasattr(cl, "AskActionMessage") and hasattr(cl, "Action")


def _discover_action_fields() -> set[str]:
    if not HAS_ACTION_PROMPT:
        return set()
    action_cls = cl.Action
    for attr in ("model_fields", "__fields__"):
        fields = getattr(action_cls, attr, None)
        if isinstance(fields, dict):
            return set(fields.keys())
    return set()


ACTION_FIELDS = _discover_action_fields()
GENERATION_FIELD_ORDER = [
    "max_new_tokens",
    "temperature",
    "top_k",
    "top_p",
    "do_sample",
    "presence_penalty",
    "frequency_penalty",
    "no_repeat_ngram",
    "min_steps_before_eos",
]
OPTIONAL_KEYS = {"top_k", "top_p", "no_repeat_ngram"}
BOOL_KEYS = {"do_sample"}


def _parse_bool(value: str) -> bool:
    candidates = {"true", "1", "yes", "y", "on"}
    anti_candidates = {"false", "0", "no", "n", "off"}
    lower = value.lower()
    if lower in candidates:
        return True
    if lower in anti_candidates:
        return False
    raise ValueError("expected a boolean (true/false)")


def _parse_optional(value: str) -> str | None:
    lower = value.lower()
    if lower in {"none", "null", "default"}:
        return None
    return value


def _parse_generation_value(key: str, raw: str) -> Any:
    parsed = _parse_optional(raw) if key in OPTIONAL_KEYS else raw
    if parsed is None:
        return None
    if key in BOOL_KEYS:
        return _parse_bool(parsed)
    if key in {"max_new_tokens", "top_k", "no_repeat_ngram", "min_steps_before_eos"}:
        return int(parsed)
    return float(parsed)


def _get_generation_overrides() -> dict[str, Any]:
    overrides: dict[str, Any] | None = cl.user_session.get(GENERATION_SESSION_KEY)
    return dict(overrides) if overrides else {}


def _set_generation_overrides(overrides: dict[str, Any]) -> None:
    cl.user_session.set(GENERATION_SESSION_KEY, overrides)


def _format_generation_overrides(overrides: dict[str, Any]) -> str:
    if not overrides:
        return "  (none)"
    lines: list[str] = []
    for key in GENERATION_FIELD_ORDER:
        if key not in overrides:
            continue
        value = overrides[key]
        if isinstance(value, bool):
            display = "true" if value else "false"
        else:
            display = "None" if value is None else str(value)
        lines.append(f"- {key} = {display}")
    return "\n".join(lines) if lines else "  (none)"


def _render_generation_help(overrides: dict[str, Any]) -> str:
    allowed = ", ".join(GENERATION_FIELD_ORDER)
    summary = _format_generation_overrides(overrides)
    return (
        "Current generation overrides:\n"
        f"{summary}\n\n"
        "Usage: /config key=value [key=value ...]\n"
        "Set optional fields to 'none' to restore default values.\n"
        "Run `/config reset` to clear all overrides.\n"
        f"Allowed keys: {allowed}"
    )


class ModelOption:
    """Lightweight container describing a trained model configuration."""

    def __init__(self, name: str) -> None:
        self.name = name

    @property
    def job_name(self) -> str:
        return f"{JOB_NAME}_{self.name}"


def discover_model_options() -> list[ModelOption]:
    ensure_models_root()
    return [ModelOption(name) for name in list_model_names()]


def restore_model_from_session() -> ModelOption | None:
    data: dict[str, str] | None = cl.user_session.get(MODEL_SESSION_KEY)
    if not data:
        env_model = os.environ.get(MODEL_NAME_ENV)
        if env_model and env_model in list_model_names():
            return ModelOption(name=env_model)
        return None
    name = data.get("name")
    if not name:
        return None
    if name not in list_model_names():
        return None
    return ModelOption(name=name)


def persist_model_selection(option: ModelOption) -> None:
    cl.user_session.set(MODEL_SESSION_KEY, {"name": option.name})


def extract_action_value(result: Any) -> str | None:
    if result is None:
        return None
    if isinstance(result, str):
        return result
    if isinstance(result, dict):
        payload = result.get("payload")
        if isinstance(payload, dict):
            selected = payload.get("model") or payload.get("value") or payload.get("name")
            if selected:
                return str(selected)
        for key in ("value", "name", "action"):
            value = result.get(key)
            if value:
                return str(value)
        return None
    value = getattr(result, "value", None)
    if value:
        return str(value)
    name = getattr(result, "name", None)
    if name:
        return str(name)
    payload = getattr(result, "payload", None)
    if isinstance(payload, dict):
        selected = payload.get("model") or payload.get("value") or payload.get("name")
        if selected:
            return str(selected)
    return None


async def prompt_model_selection(header: str | None = None) -> ModelOption | None:
    options = discover_model_options()
    if not options:
        await cl.Message(
            content=(
                "No trained models found in `trained_models/`. "
                "Place a model folder with configs/config_inference.yaml there and restart the chat."
            )
        ).send()
        return None

    if not HAS_ACTION_PROMPT:
        chosen = options[0]
        await cl.Message(
            content=f"Defaulting to `{chosen.name}`. Action prompts are unavailable."
        ).send()
        return chosen

    fields = ACTION_FIELDS or {"name", "value", "label", "payload"}
    actions = []
    for option in options:
        kwargs: dict[str, Any] = {"name": option.name}
        if "label" in fields:
            kwargs["label"] = option.name
        if "value" in fields:
            kwargs["value"] = option.name
        if not fields or "payload" in fields:
            kwargs["payload"] = {"model": option.name}
        actions.append(cl.Action(**kwargs))
    message = header or "Select a trained model to use for inference:"
    result = await cl.AskActionMessage(content=message, actions=actions).send()
    selected_name = extract_action_value(result)
    if not selected_name:
        selected_name = options[0].name
    return next((opt for opt in options if opt.name == selected_name), options[0])


def build_config(user_input: str, model: ModelOption) -> DictConfig:
    if GlobalHydra.instance().is_initialized():
        GlobalHydra.instance().clear()
    cfg = compose_model_config(model.name, job_name=model.job_name)
    overrides = _get_generation_overrides()
    if overrides:
        for key, value in overrides.items():
            if key in cfg.generation:
                cfg.generation[key] = value
    cfg.input_text = user_input
    return cfg


async def ensure_model_selected(header: str | None = None) -> ModelOption | None:
    model = restore_model_from_session()
    if model:
        return model
    model = await prompt_model_selection(header=header)
    if model:
        persist_model_selection(model)
    return model


async def handle_model_command(content: str) -> None:
    parts = content.split(maxsplit=1)
    options = discover_model_options()
    if not options:
        await cl.Message(
            content="No trained models available to select. Add a model under `trained_models/` and try again."
        ).send()
        return

    if len(parts) == 1:
        model = await prompt_model_selection(header="Pick a new model:")
    else:
        requested = parts[1].strip()
        model = next((opt for opt in options if opt.name == requested), None)
        if not model:
            available = ", ".join(opt.name for opt in options)
            await cl.Message(
                content=f"Model `{requested}` not found. Available models: {available}."
            ).send()
            return
    if model:
        persist_model_selection(model)
        await cl.Message(
            content=f"Switched to `{model.name}`. Provide text to generate a continuation."
        ).send()


async def handle_generation_command(content: str) -> None:
    tokens = shlex.split(content)
    overrides = _get_generation_overrides()

    if len(tokens) == 1:
        await cl.Message(content=_render_generation_help(overrides)).send()
        return

    first = tokens[1].lower()
    if first == "reset":
        _set_generation_overrides({})
        await cl.Message(
            content="Generation overrides cleared. Using model defaults for all parameters."
        ).send()
        return

    updates: dict[str, Any] = {}
    errors: list[str] = []
    for token in tokens[1:]:
        if "=" not in token:
            errors.append(f"Missing '=' in '{token}'. Use key=value format.")
            continue
        key, raw_value = token.split("=", 1)
        normalized_key = key.strip().lower()
        if normalized_key not in GENERATION_FIELD_ORDER:
            errors.append(f"Unknown parameter '{key}'.")
            continue
        try:
            value = _parse_generation_value(normalized_key, raw_value.strip())
        except ValueError as exc:
            errors.append(f"Invalid value for '{key}': {exc}")
            continue
        updates[normalized_key] = value

    if updates:
        overrides.update(updates)
        _set_generation_overrides(overrides)

    parts: list[str] = []
    if updates:
        parts.append("Updated generation overrides:\n" + _format_generation_overrides(overrides))
    if errors:
        formatted_errors = "\n".join(f"- {msg}" for msg in errors)
        parts.append(f"Issues detected:\n{formatted_errors}")
    if not parts:
        parts.append(_render_generation_help(overrides))

    await cl.Message(content="\n\n".join(parts)).send()


@cl.on_chat_start
async def on_chat_start() -> None:
    model = await ensure_model_selected()
    if not model:
        return
    await cl.Message(
        content=(
            f"Using trained model `{model.name}`. Send text for generation, or type `/model` to switch models."
            " Use `/config` to tweak generation settings."
        )
    ).send()


@cl.on_message
async def on_message(message: cl.Message) -> None:
    user_input = message.content.strip()

    if user_input.lower().startswith("/model"):
        await handle_model_command(user_input)
        return
    if user_input.lower().startswith("/config"):
        await handle_generation_command(user_input)
        return

    if not user_input:
        await cl.Message(content="Please provide some text for the model.").send()
        return

    model = await ensure_model_selected(header="Select a model before requesting generations:")
    if not model:
        return

    cfg = build_config(user_input, model)
    try:
        outputs = run_inference(cfg)
    except Exception as exc:  # pragma: no cover
        await cl.Message(content=f"Inference error: {exc}").send()
        return

    response = outputs[0] if outputs else "No output generated."
    await cl.Message(content=response).send()