"""Generate a task and behavior representation from a conversation."""

import argparse
import json
import sys
import time
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForImageTextToText, AutoTokenizer


# An authored, simplified example inspired by an Excel-help WildChat conversation.
EXAMPLE_CONVERSATION = [
    {
        "role": "user",
        "content": "I need an Excel formula to count dates in A2:A100 before today plus B1.",
    },
    {
        "role": "assistant",
        "content": 'Try =COUNTIF(A2:A100,"<"&TODAY()+B1).',
    },
    {
        "role": "user",
        "content": "That didn't work. The dates come from another formula.",
    },
    {
        "role": "assistant",
        "content": "Are those results actual Excel dates, or text that looks like dates?",
    },
    {"role": "user", "content": "They're text. Some cells are empty too."},
    {
        "role": "assistant",
        "content": "Convert the nonempty text values with DATEVALUE, then count the resulting dates.",
    },
    {
        "role": "user",
        "content": "Can you show the steps using a helper column?",
    },
    {
        "role": "assistant",
        "content": 'In C2 enter =IF(A2="","",DATEVALUE(A2)), fill down to C100, then use =COUNTIFS(C2:C100,"<"&TODAY()+B1,C2:C100,">0").',
    },
]


def format_prompt(conversation: list[dict[str, str]], template: str) -> str:
    if not isinstance(conversation, list) or not conversation:
        raise ValueError("The conversation must be a nonempty list of messages.")
    for message in conversation:
        if not isinstance(message, dict):
            raise ValueError("Each message must be an object with role and content.")
        if message.get("role") not in ("user", "assistant"):
            raise ValueError("Message roles must be 'user' or 'assistant'.")
        if not isinstance(message.get("content"), str) or not message["content"].strip():
            raise ValueError("Each message must contain nonempty text.")
    history = "\n".join(f"[{m['role']}]: {m['content']}" for m in conversation)
    return template.format(conversation_history=history)


def parse_representation(text: str) -> dict[str, str]:
    representation = json.loads(text)
    if not isinstance(representation, dict) or set(representation) != {"task", "user_behavior"}:
        raise ValueError(f"Expected exactly task and user_behavior fields; received: {text}")
    if not all(isinstance(value, str) and value.strip() for value in representation.values()):
        raise ValueError(f"Both representation fields must be nonempty strings; received: {text}")
    return representation


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="microsoft/UserRepGen-4B", help="Hub ID or local model directory.")
    parser.add_argument("--conversation", type=Path, help="JSON list of role/content messages; otherwise use the built-in example.")
    parser.add_argument("--device", choices=["cpu", "cuda"], default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()

    conversation = (
        json.loads(args.conversation.read_text(encoding="utf-8"))
        if args.conversation is not None
        else EXAMPLE_CONVERSATION
    )
    if Path(args.model).is_dir():
        prompt_path = Path(args.model) / "representation_prompt.txt"
    else:
        prompt_path = Path(hf_hub_download(args.model, "representation_prompt.txt"))
    prompt = format_prompt(conversation, prompt_path.read_text(encoding="utf-8").strip())

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    inputs = tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}],
        tokenize=True,
        add_generation_prompt=True,
        enable_thinking=False,
        return_tensors="pt",
        return_dict=True,
    )
    prompt_tokens = inputs["input_ids"].shape[1]
    if prompt_tokens > 8192:
        raise ValueError(f"Prompt has {prompt_tokens} tokens; the supported maximum is 8192.")

    if args.device == "cpu":
        torch.set_num_threads(8)
    print(f"Loading {args.model} on {args.device}; prompt: {prompt_tokens} tokens.", file=sys.stderr, flush=True)
    model, loading_info = AutoModelForImageTextToText.from_pretrained(
        args.model,
        dtype=torch.float32 if args.device == "cpu" else torch.bfloat16,
        device_map=args.device,
        output_loading_info=True,
    )
    for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs"):
        if loading_info.get(key):
            raise RuntimeError(f"Checkpoint loading failed ({key}): {loading_info[key]}")
    model.eval()
    inputs = inputs.to(args.device)
    started = time.monotonic()
    print("Generating representation...", file=sys.stderr, flush=True)
    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=1024,
            do_sample=False,
            temperature=None,
            top_p=None,
            top_k=None,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )
    generated_ids = output[0, prompt_tokens:]
    text = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
    representation = parse_representation(text)
    print(json.dumps(representation, indent=2, ensure_ascii=False))
    print(
        f"Validated JSON; generated {len(generated_ids)} tokens in {time.monotonic() - started:.1f}s.",
        file=sys.stderr,
    )


if __name__ == "__main__":
    main()