"""Minimal loader for the Fairleap v1 CLM Sahabat-AI 8B adapter.

Four things about this model will look like broken weights if you meet them
cold, so they are handled here rather than left as surprises:

1. **The base model's own chat template silently destroys tool calls.**
   `llama3-8b-cpt-sahabatai-v1-instruct` ships the plain Llama-3 template,
   which renders every message as `content | trim` and ignores both
   `tool_calls` and the `tools` argument. An assistant tool-call turn carries
   `content: ""`, so it renders as an *empty* assistant reply -- no error
   anywhere. The tool-capable template is in `chat_template.jinja` beside the
   weights, and `from_pretrained` on **this directory** picks it up. Load the
   tokenizer from the base repo instead and tool use breaks with no warning.
2. **This repository holds a LoRA adapter only.** There are no base weights and
   no `config.json`, so `AutoModelForCausalLM.from_pretrained(".")` fails and
   PEFT falls back to treating the local path as a Hub repo id, reporting
   `HFValidationError: Repo id must be in the form 'repo_name'...`, which reads
   as a path bug.
3. **The base's generation config disagrees with its own template.** It names
   `<|end_of_text|>` (128001) as eos while every turn ends with `<|eot_id|>`
   (128009), so generation never stops. `chat()` passes eos explicitly.
4. **Tool calls come back as one line of JSON**, keyed `parameters` rather than
   the `arguments` most OpenAI-shaped parsers expect:
   `{"name": "predict_earnings", "parameters": {...}}`. `parse_tool_call()`
   recovers it. **`fairleap-api` must parse this shape** -- it differs from the
   XML block the Qwen adapter emits.

    python load_model.py            # loads and runs one grounded question
"""

import json
from pathlib import Path

REPO_ID = "fairleap-ai/fairleap-v1-clm-sahabatai-8b-adapter"
BASE_MODEL = "GoToCompany/llama3-8b-cpt-sahabatai-v1-instruct"
ADAPTER_DIR = str(Path(__file__).parent)

MAX_SEQ = 4096
EOT_TOKEN = "<|eot_id|>"

# The only callable tool. Advice topics are prompt blocks, not tools -- calling
# an LLM endpoint from an LLM is an extra hop and a second inference bill.
PREDICT_EARNINGS_TOOL = {
    "type": "function",
    "function": {
        "name": "predict_earnings",
        "description": (
            "Prediksi penghasilan harian dan jam kerja untuk rentang tanggal di masa "
            "depan menggunakan model XGBoost. Riwayat harian driver ditambahkan "
            "otomatis oleh sistem, jangan sertakan riwayat dalam argumen."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "start": {"type": "string", "description": "Tanggal awal, YYYY-MM-DD."},
                "end": {"type": "string", "description": "Tanggal akhir, YYYY-MM-DD."},
                "wellness_score": {"type": "integer", "minimum": 1, "maximum": 100},
            },
            "required": ["start", "end", "wellness_score"],
        },
    },
}

PERSONA = """Kamu adalah Fairleap AI Assistant, asisten untuk mitra driver Gojek/GOTO di Indonesia.
Jawab hangat, ringkas, dan konkret dalam Bahasa Indonesia. Jangan pernah mengarang angka."""


def build_system_prompt(today, city, vehicle, risk, wellness_score,
                        period, totals, daily=None):
    """Assemble the stuffed driver context the model was trained to read.

    daily is an optional list of (date, earnings, orders) rows. Training used
    seven-day summaries plus up to fourteen daily rows; the model reads these
    figures literally, so anything absent here it should decline to state.

    Include the *previous* period's total whenever you have it. Held-out
    evaluation found the model appends "Dibanding 7 hari sebelumnya (RpX)..."
    to half of all earnings answers, inventing X when it was never supplied.
    Supplying the real figure is the cheapest mitigation available to a caller.
    """
    lines = [
        PERSONA,
        "",
        "=== KONTEKS DRIVER (disediakan sistem, bukan hasil tool) ===",
        f"Tanggal hari ini: {today}",
        f"Kota: {city}",
        f"Kendaraan: {vehicle}",
        f"Toleransi risiko: {risk}",
        f"Skor kesehatan terakhir: {wellness_score}/100",
        f"Ringkasan 7 hari terakhir ({period}):",
    ]
    lines += [f"- {k}: {v}" for k, v in totals.items()]
    if daily:
        lines.append("Rincian harian:")
        # Indonesian thousands separator is a dot, so convert the number alone --
        # doing it on the whole line turns the comma before "order" into a dot.
        lines += [f"- {d}: Rp{e:,}".replace(",", ".") + f", {o} order"
                  for d, e, o in daily]
    return "\n".join(lines)


def get_tokenizer(processor):
    """Unwrap a processor if one is handed back. See note 1 above."""
    return getattr(processor, "tokenizer", processor)


def load(adapter_dir=ADAPTER_DIR, load_in_4bit=True):
    """Return (model, tokenizer). Prefers unsloth, falls back to peft."""
    # Without this the loaders report `Unrecognized model in <dir>. Should have
    # a model_type key in its config.json`, which sends you looking for a
    # corrupt checkpoint instead of a wrong path. ADAPTER_DIR defaults to this
    # file's own directory, so moving load_model.py away from the weights
    # breaks it.
    if not Path(adapter_dir, "adapter_config.json").is_file():
        raise SystemExit(
            f"no adapter_config.json in {adapter_dir!r} -- pass adapter_dir= "
            f"pointing at the directory holding adapter_model.safetensors"
        )

    try:
        from unsloth import FastLanguageModel

        model, processor = FastLanguageModel.from_pretrained(
            model_name=adapter_dir, max_seq_length=MAX_SEQ,
            dtype=None, load_in_4bit=load_in_4bit)
        FastLanguageModel.for_inference(model)
        return model, get_tokenizer(processor)
    except ImportError:
        pass

    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

    # From the adapter directory, never the base repo: that is where the
    # tool-capable chat_template.jinja lives.
    tokenizer = AutoTokenizer.from_pretrained(adapter_dir)
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL, torch_dtype="auto", device_map="auto")
    model = PeftModel.from_pretrained(base, adapter_dir)
    model.eval()
    return model, get_tokenizer(tokenizer)


def chat(model, tokenizer, messages, tools=None, max_new_tokens=400,
         temperature=0.7, top_p=0.9):
    """One turn. Pass tools only when a forecast is plausibly needed.

    On the distribution it was trained for the model never calls the tool
    unbidden -- held-out evaluation recorded 40 calls on exactly the 40
    conversations that offered one, and zero unsolicited. Offer it on a turn
    that does not need it, though, and it will still reach for a forecast. The
    caller decides, not the model.
    """
    text = tokenizer.apply_chat_template(
        messages, tools=tools, tokenize=False, add_generation_prompt=True)

    encoded = tokenizer(text, return_tensors="pt")
    encoded = {k: v.to(model.device) for k, v in encoded.items()
               if k in ("input_ids", "attention_mask")}

    import torch

    # The base generation config names <|end_of_text|>, but the template ends
    # every turn with <|eot_id|>. Without this override generation runs on past
    # the reply -- and past a tool call into the result it was about to request.
    eot = tokenizer.convert_tokens_to_ids(EOT_TOKEN)

    with torch.no_grad():
        out = model.generate(
            **encoded,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            eos_token_id=eot,
            pad_token_id=tokenizer.pad_token_id or eot,
        )
    return tokenizer.decode(out[0][encoded["input_ids"].shape[-1]:],
                            skip_special_tokens=True)


def parse_tool_call(text):
    """Recover a tool call. The model emits one line of JSON keyed `parameters`.

    Not `arguments`: an OpenAI-shaped parser looking for that key finds nothing
    and silently reports "no tool call" on a perfectly good one.
    """
    import re

    match = re.search(
        r'"name"\s*:\s*"([a-z_]+)".*?"(?:parameters|arguments)"\s*:\s*(\{.*?\})',
        text, re.S)
    if not match:
        return None
    try:
        return match.group(1), json.loads(match.group(2))
    except json.JSONDecodeError:
        return match.group(1), {}


if __name__ == "__main__":
    system = build_system_prompt(
        today="2026-08-24", city="Bekasi", vehicle="motor", risk="sedang",
        wellness_score=62, period="2026-08-18 s/d 2026-08-24",
        totals={
            "Total penghasilan": "Rp1.482.000",
            "Total order": "88",
            "Hari kerja": "6 dari 7 hari",
            "Rata-rata per hari kerja": "Rp247.000",
            "Hari terbaik": "2026-08-21 (Rp281.000, 16 order)",
            # Supply this and the model quotes it instead of inventing one.
            "Total 7 hari sebelumnya": "Rp1.301.000",
        },
    )

    model, tokenizer = load()

    question = "berapa penghasilan saya minggu ini?"
    reply = chat(model, tokenizer,
                 [{"role": "system", "content": system},
                  {"role": "user", "content": question}])
    print(f">>> {question}\n{reply.strip()}\n")

    question = "kira-kira minggu depan saya bisa dapat berapa?"
    reply = chat(model, tokenizer,
                 [{"role": "system", "content": system},
                  {"role": "user", "content": question}],
                 tools=[PREDICT_EARNINGS_TOOL])
    print(f">>> {question}\n{reply.strip()}")
    print(f"\nparsed tool call: {json.dumps(parse_tool_call(reply), ensure_ascii=False)}")