import torch
from typing import List, Dict, Any, Optional
from inspector import (
    compute_probabilities,
    compute_entropy,
    get_top_k,
    build_step_result,
)

DECODING_MODES = {
    "greedy":   "Greedy — selalu pilih token tertinggi (deterministik)",
    "sampling": "Sampling — random sample dari distribusi (variatif)",
}


def greedy_decode_with_inspection(
    model,
    tokenizer,
    device: str,
    prompt: str,
    max_new_tokens: int = 50,
    top_k: int = 10,
    temperature: float = 1.0,
    stop_at_eos: bool = True,
    decoding_mode: str = "greedy",      # ← BARU
    top_p: float = 0.9,                 # ← BARU
    repetition_penalty: float = 1.0,    # ← BARU
    seed: Optional[int] = None,         # ← BARU
    progress_callback=None,
) -> List[Dict[str, Any]]:
    """
    Run decoding step-by-step and extract full probability
    inspection data at each step.

    Args:
        model             : Loaded HuggingFace causal LM
        tokenizer         : Corresponding tokenizer
        device            : 'cpu' or 'cuda'
        prompt            : Input text string
        max_new_tokens    : Maximum number of tokens to generate
        top_k             : Number of top candidates to extract per step
        temperature       : Softmax temperature
        stop_at_eos       : Stop when EOS token is produced
        decoding_mode     : 'greedy' or 'sampling'           ← BARU
        top_p             : Nucleus sampling threshold        ← BARU
        repetition_penalty: Penalize repeated tokens          ← BARU
        seed              : Random seed (sampling mode only)  ← BARU
        progress_callback : Optional callable(step, total)
    """
    if seed is not None:
        torch.manual_seed(seed)

    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
    eos_token_id = tokenizer.eos_token_id

    results = []
    current_ids = input_ids.clone()

    with torch.no_grad():
        for step in range(1, max_new_tokens + 1):

            outputs = model(input_ids=current_ids)
            next_token_logits = outputs.logits[0, -1, :].clone()

            # ── Repetition Penalty ──────────────────────────────────
            if repetition_penalty != 1.0:
                next_token_logits = _apply_repetition_penalty(
                    next_token_logits, current_ids[0], repetition_penalty
                )

            # ── Probability Distribution ────────────────────────────
            probs = compute_probabilities(next_token_logits, temperature=temperature)

            # ── Entropy & Top-K (selalu dari full distribution) ─────
            entropy = compute_entropy(probs)
            top_candidates = get_top_k(probs, tokenizer, k=top_k)

            # ── Token Selection ─────────────────────────────────────
            if decoding_mode == "greedy":
                chosen_token_id = torch.argmax(probs).item()

            elif decoding_mode == "sampling":
                filtered_probs = _top_p_filter(probs.clone(), top_p=top_p)
                chosen_token_id = torch.multinomial(filtered_probs, num_samples=1).item()

            else:
                raise ValueError(f"Unknown decoding_mode: '{decoding_mode}'. Use 'greedy' or 'sampling'.")

            chosen_token_str = tokenizer.decode([chosen_token_id])
            input_so_far = tokenizer.decode(current_ids[0], skip_special_tokens=True)

            step_result = build_step_result(
                step=step,
                input_so_far=input_so_far,
                chosen_token=chosen_token_str,
                top_candidates=top_candidates,
                entropy=entropy,
            )
            step_result["decoding_mode"] = decoding_mode  # ← BARU
            results.append(step_result)

            if progress_callback:
                progress_callback(step, max_new_tokens)

            next_token_tensor = torch.tensor([[chosen_token_id]], device=device)
            current_ids = torch.cat([current_ids, next_token_tensor], dim=1)

            if stop_at_eos and chosen_token_id == eos_token_id:
                break

    return results


def _apply_repetition_penalty(
    logits: torch.Tensor,
    input_ids: torch.Tensor,
    penalty: float,
) -> torch.Tensor:
    """
    Turunkan probabilitas token yang sudah muncul di konteks.
    Logits positif dibagi penalty, logits negatif dikali penalty.
    """
    for token_id in set(input_ids.tolist()):
        if logits[token_id] > 0:
            logits[token_id] /= penalty
        else:
            logits[token_id] *= penalty
    return logits


def _top_p_filter(
    probs: torch.Tensor,
    top_p: float = 0.9,
) -> torch.Tensor:
    """
    Nucleus (top-p) filtering: zero out token di luar
    top-p% cumulative probability mass.
    """
    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

    sorted_indices_to_remove = cumulative_probs - sorted_probs > top_p
    sorted_probs[sorted_indices_to_remove] = 0.0

    filtered = torch.zeros_like(probs)
    filtered.scatter_(0, sorted_indices, sorted_probs)

    total = filtered.sum()
    if total > 0:
        filtered = filtered / total

    return filtered


def get_generated_text(prompt: str, results: List[Dict]) -> str:
    """Reconstruct the full generated text from step results."""
    tokens = [r["chosen_token"] for r in results]
    return prompt + "".join(tokens)