import json
import time

import click
import torch
from tqdm import tqdm

from astrai.inference import InferenceEngine
from astrai.model import AutoModel
from astrai.tokenize import AutoTokenizer


def processor(
    param_path: str,
    input_json_file: str,
    output_json_file: str,
    temperature: float,
    top_k: int,
    top_p: float,
    question_key: str,
    response_key: str,
    max_tokens: int,
    batch_size: int,
    num_samples: int = 1,
    cache_len: int = 2048,
    frequency_penalty: float = 0.0,
    rep_window: int = 64,
):
    print(f"Loading model from {param_path} ...")
    t0 = time.time()
    model = AutoModel.from_pretrained(param_path)
    tokenizer = AutoTokenizer.from_pretrained(param_path)
    model.to(device="cuda", dtype=torch.bfloat16)
    print(f"  model loaded in {time.time() - t0:.1f}s")

    engine = InferenceEngine(
        model=model,
        tokenizer=tokenizer,
        max_batch_size=batch_size * num_samples,
        max_seq_len=cache_len,
    )

    print(f"Reading {input_json_file} ...")
    with open(input_json_file, "r", encoding="utf-8") as f:
        input_data = [json.loads(line) for line in f]

    if input_data and "messages" in input_data[0]:
        prompts = [
            tokenizer.apply_chat_template(item["messages"], tokenize=False)
            for item in input_data
        ]
    else:
        prompts = [item[question_key] for item in input_data]
    print(f"  {len(prompts)} prompts loaded\n")

    if max_tokens is None:
        max_tokens = model.config.max_position_embeddings

    chunk_size = max(1, batch_size)

    with open(output_json_file, "w", encoding="utf-8") as f:
        pbar = tqdm(
            total=len(prompts) * num_samples,
            unit="gen",
            desc=f" Generating ({num_samples}x/prompt)",
        )
        for chunk_start in range(0, len(prompts), chunk_size):
            chunk = prompts[chunk_start : chunk_start + chunk_size]

            if num_samples > 1:
                chunk_expanded = [p for p in chunk for _ in range(num_samples)]
                resp_chunk = engine.generate(
                    prompt=chunk_expanded,
                    stream=False,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                    top_k=top_k,
                    frequency_penalty=frequency_penalty,
                    rep_window=rep_window,
                )
                resp_chunk = [
                    resp_chunk[i * num_samples : (i + 1) * num_samples]
                    for i in range(len(chunk))
                ]
            else:
                resp_chunk = engine.generate(
                    prompt=chunk,
                    stream=False,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                    top_k=top_k,
                    frequency_penalty=frequency_penalty,
                    rep_window=rep_window,
                )

            for i, prompt in enumerate(chunk):
                if input_data and "messages" in input_data[0]:
                    orig = input_data[chunk_start + i]
                    output_item = {**orig, response_key: resp_chunk[i]}
                else:
                    output_item = {
                        question_key: prompt,
                        response_key: resp_chunk[i],
                    }
                f.write(json.dumps(output_item, ensure_ascii=False) + "\n")

            pbar.update(len(chunk) * num_samples)

        pbar.close()

    elapsed = time.time() - t0
    print(
        f"\nDone! {len(prompts)} prompts x {num_samples} samples -> {output_json_file}"
    )
    print(f"Total time: {elapsed:.1f}s ({elapsed / len(prompts):.2f}s/prompt)")

    engine.shutdown()


@click.command(name="generate", help="Batch generation from a JSONL prompt file.")
@click.option(
    "--param_path",
    type=click.Path(exists=True),
    required=True,
    help="Path to the model directory.",
)
@click.option(
    "--input_json_file",
    type=click.Path(exists=True),
    required=True,
    help="Path to the input JSONL file.",
)
@click.option(
    "--output_json_file",
    type=click.Path(),
    required=True,
    help="Path to the output JSONL file.",
)
@click.option(
    "--question_key", default="question", help="Key for the question in input JSON."
)
@click.option(
    "--response_key", default="response", help="Key for the response in output JSON."
)
@click.option("--temperature", type=float, default=0.60, help="Sampling temperature.")
@click.option("--top_k", type=int, default=30, help="Top-k filtering.")
@click.option("--top_p", type=float, default=0.95, help="Top-p filtering.")
@click.option("--batch_size", type=int, default=1, help="Batch size.")
@click.option("--num_samples", type=int, default=1, help="Responses per prompt.")
@click.option("--max_tokens", type=int, default=None, help="Max tokens to generate.")
@click.option("--cache_len", type=int, default=2048, help="KV cache length.")
@click.option("--frequency_penalty", type=float, default=0.0, help="Frequency penalty.")
@click.option(
    "--rep_window", type=int, default=64, help="Window size for frequency penalty."
)
def generate_command(**kwargs):
    """Batch generation from a JSONL prompt file."""
    with torch.inference_mode():
        processor(**kwargs)


if __name__ == "__main__":
    generate_command()