"""PUBLIC INTERFACE of Serving_Module. Other modules may import ONLY this file.

resolve_trained_dir(cfg, fetch_remote)      -> Path   local trained_model/ first, else OUTPUT_HF_REPO via `fetch_remote`
prepare_model(model, trained_dir, meta)     -> model  attaches the LoRA adapter when the checkpoint is adapter-based
build_engine(model, tokenizer, cfg)         -> engine
start(engine, cfg)                          -> None   Telegram bot if TELEGRAM_BOT_TOKEN is set, else CLI chat

Weights are loaded by the orchestrator (hardware tier logic lives in Hardware_Module); `fetch_remote` is injected
for the same reason, so this module never imports another module.
"""
from __future__ import annotations

from pathlib import Path
from typing import Callable, Dict, Optional

from .internal.cli_chat_window import run_cli
from .internal.inference_engine import InferenceEngine
from .internal.telegram_gateway import TelegramGateway

PLACEHOLDER_MARK = "PUT_YOUR_TOKEN"


def _has_meta(folder: Path) -> bool:
    return (folder / "rss_meta.json").exists()


def resolve_trained_dir(cfg: Dict, fetch_remote: Optional[Callable] = None) -> Path:
    root = Path(cfg["root"])
    local = root / "trained_model"
    if _has_meta(local):
        print(f"[Serving] using local trained model: {local}")
        return local
    repo, token = cfg.get("output_hf_repo"), cfg.get("hf_token")
    if repo and fetch_remote is not None:
        try:
            print(f"[Serving] no local trained model - pulling {repo} from Hugging Face ...")
            token = None if (not token or PLACEHOLDER_MARK in token) else token
            return Path(fetch_remote(repo, root / "trained_model", token))
        except Exception as exc:
            raise FileNotFoundError(
                f"no trained model found locally ({local}) and could not download {repo}: "
                f"{type(exc).__name__}: {exc}. Run:  python RSS_AI_Engine/main.py --mode train_all") from exc
    raise FileNotFoundError(f"no trained model found in {local}. Run:  python RSS_AI_Engine/main.py --mode train_all")


def prepare_model(model, trained_dir, meta: Dict):
    """Attach the LoRA adapter for adapter-based checkpoints; fold it into the weights when that is possible."""
    if not meta.get("has_adapter"):
        return model
    from peft import PeftModel
    adapter = Path(trained_dir) / (meta.get("adapter_dir") or "adapter")
    print(f"[Serving] attaching LoRA adapter {adapter}")
    wrapped = PeftModel.from_pretrained(model, str(adapter), is_trainable=False)
    plan = getattr(model, "_rss_plan", {}) or {}
    ring = getattr(model, "_rss_ring", None)
    if ring is None and not plan.get("quant"):
        try:
            merged = wrapped.merge_and_unload()                   # faster generation: no adapter matmuls
            for attr in ("_rss_plan", "_rss_ring"):
                setattr(merged, attr, getattr(model, attr, None))
            return merged
        except Exception as exc:
            print(f"[Serving] adapter merge skipped ({type(exc).__name__}); serving with the adapter attached")
    return wrapped


def build_engine(model, tokenizer, cfg: Dict) -> InferenceEngine:
    return InferenceEngine(model, tokenizer, chat_file=cfg.get("chat_file"),
                           max_new_tokens=(cfg.get("tune") or {}).get("serve_max_new_tokens"),
                           active_vocab_file=cfg.get("active_vocab_file"))


def start(engine: InferenceEngine, cfg: Dict, telegram_session=None) -> str:
    """Blocks until the user quits. Returns which front-end ran: 'telegram' or 'cli'."""
    token = (cfg.get("telegram_token") or "").strip()
    if token and PLACEHOLDER_MARK not in token:
        gateway = TelegramGateway(token, cfg.get("telegram_chat_id", ""), engine, session=telegram_session)
        if gateway.verify():
            gateway.notify_start()
            try:
                gateway.run()
            except KeyboardInterrupt:
                print("\n[Serving] Telegram bot stopped.")
            return "telegram"
        print("[Serving] Telegram token rejected -> falling back to the CLI chat window")
    else:
        print("[Serving] TELEGRAM_BOT_TOKEN is empty -> opening the interactive CLI chat window")
    run_cli(engine)
    return "cli"