from __future__ import annotations

import argparse
import io
import json
import math
import sys
import time
from pathlib import Path
from typing import Any

import torch

sys.path.insert(0, str(Path(__file__).parent.parent))

from model_cpu_gpt2 import CPUGPT, gpt2_small_config, smoke_config

BYTES_PER_TOKEN = 4.0
SEQ_LEN = 1024
BATCH_SIZE = 8


def _load_ckpt(path: str) -> dict:
    if path.startswith("s3://"):
        import boto3

        bucket, key = path[5:].split("/", 1)
        buf = io.BytesIO()
        boto3.client("s3").download_fileobj(bucket, key, buf)
        buf.seek(0)
        return torch.load(buf, map_location="cpu", weights_only=False)
    return torch.load(path, map_location="cpu", weights_only=False)


def _build_our_model(config_name: str, ckpt_path: str | None) -> CPUGPT:
    cfg_map = {
        "gpt2-small": gpt2_small_config,
        "smoke": smoke_config,
    }
    if config_name not in cfg_map:
        raise ValueError(
            f"Unknown config '{config_name}'. Choose from: {list(cfg_map)}"
        )
    cfg = cfg_map[config_name]()
    model = CPUGPT(cfg)
    if ckpt_path:
        ckpt = _load_ckpt(ckpt_path)
        state = ckpt.get("model", ckpt)
        state = {k.replace("_orig_mod.", ""): v for k, v in state.items()}
        model.load_state_dict(state, strict=True)
        print(f"Loaded checkpoint: {ckpt_path}", flush=True)
    else:
        print("No checkpoint supplied — using random weights.", flush=True)
    model.eval()
    return model


def _require_lm_eval():
    try:
        import lm_eval
    except ImportError:
        print(
            "\nERROR: lm-evaluation-harness is not installed.\n"
            "Install it with:\n"
            "    pip install lm-eval>=0.4.2\n",
            file=sys.stderr,
        )
        sys.exit(1)


try:
    from lm_eval.api.model import LM as _LMBase
except ImportError:
    _LMBase = object


class LMEvalWrapper(_LMBase):
    def __init__(self, model: CPUGPT):
        super().__init__()
        self.model = model
        self.model.eval()
        try:
            import tiktoken

            self._enc = tiktoken.get_encoding("gpt2")
        except ImportError:
            raise RuntimeError("tiktoken is required: pip install tiktoken")
        try:
            import tiktoken

            self._enc = tiktoken.get_encoding("gpt2")
        except ImportError:
            raise RuntimeError("tiktoken is required: pip install tiktoken")

    def _tokenize(self, text: str) -> list[int]:
        return self._enc.encode_ordinary(text)

    @torch.no_grad()
    def _logprobs_for_tokens(self, tokens: list[int]) -> list[float]:
        max_T = self.model.cfg.seq_len
        gla_chunk = self.model.cfg.gla_chunk
        log_probs: list[float] = []

        start = 0
        while start < len(tokens) - 1:
            end = min(start + max_T + 1, len(tokens))
            chunk = tokens[start:end]
            x_ids = chunk[:-1]
            targets = chunk[1:]

            T_raw = len(x_ids)
            T_padded = max(
                gla_chunk, ((T_raw + gla_chunk - 1) // gla_chunk) * gla_chunk
            )
            T_padded = min(T_padded, max_T)

            if T_padded > T_raw:
                pad_len = T_padded - T_raw
                x_ids = list(x_ids) + [0] * pad_len

            x = torch.tensor(x_ids, dtype=torch.long).unsqueeze(0)
            tgt_t = torch.tensor(targets, dtype=torch.long)

            logits = self.model(x)
            logits = logits.squeeze(0).float()
            logits = logits[:T_raw]
            lp = torch.log_softmax(logits, dim=-1)
            tgt_lp = lp[torch.arange(len(tgt_t)), tgt_t]
            log_probs.extend(tgt_lp.tolist())

            if end == len(tokens):
                break
            start = end - 1 - max_T // 2

        return log_probs

    def loglikelihood(self, requests) -> list[tuple[float, bool]]:
        results = []
        for req in requests:
            ctx_str, cont_str = req.args if hasattr(req, "args") else req
            ctx_ids = self._tokenize(ctx_str)
            cont_ids = self._tokenize(cont_str)
            tokens = ctx_ids + cont_ids
            if not tokens:
                results.append((-float("inf"), False))
                continue

            all_lp = self._logprobs_for_tokens(tokens)
            cont_lp = all_lp[-len(cont_ids) :]
            total = sum(cont_lp)

            greedy = self._is_greedy(ctx_ids, cont_ids)
            results.append((total, greedy))
        return results

    @torch.no_grad()
    def _is_greedy(self, ctx_ids: list[int], cont_ids: list[int]) -> bool:
        if not cont_ids:
            return True
        max_T = self.model.cfg.seq_len
        gla_chunk = self.model.cfg.gla_chunk
        tokens = (ctx_ids + cont_ids)[-max_T:]
        ctx_T = min(len(ctx_ids), len(tokens) - len(cont_ids))
        x_ids = tokens[:-1] if len(tokens) > 1 else tokens
        T_raw = len(x_ids)
        T_padded = max(gla_chunk, ((T_raw + gla_chunk - 1) // gla_chunk) * gla_chunk)
        T_padded = min(T_padded, max_T)
        if T_padded > T_raw:
            x_ids = list(x_ids) + [0] * (T_padded - T_raw)
        x = torch.tensor(x_ids, dtype=torch.long).unsqueeze(0)
        logits = self.model(x).squeeze(0).float()[:T_raw]
        for i, tok in enumerate(cont_ids):
            pos = ctx_T - 1 + i
            if pos >= logits.shape[0]:
                return False
            if logits[pos].argmax().item() != tok:
                return False
        return True

    @torch.no_grad()
    def generate_until(self, requests) -> list[str]:
        results = []
        for req in requests:
            ctx_str, gen_kwargs = req.args if hasattr(req, "args") else req
            until = gen_kwargs.get("until", ["\n"])
            max_new_toks = gen_kwargs.get("max_gen_toks", 50)

            ctx_ids = self._tokenize(ctx_str)
            max_T = self.model.cfg.seq_len
            if len(ctx_ids) > max_T:
                ctx_ids = ctx_ids[-max_T:]

            prompt = torch.tensor(ctx_ids, dtype=torch.long).unsqueeze(0)
            out_ids = self.model.generate(
                prompt,
                max_new_tokens=max_new_toks,
                temperature=0.0,
                top_k=1,
            )
            generated_ids = out_ids[0].tolist()
            generated_str = self._enc.decode(generated_ids)

            for stop in until:
                idx = generated_str.find(stop)
                if idx != -1:
                    generated_str = generated_str[:idx]

            results.append(generated_str)
        return results

    def loglikelihood_rolling(self, requests) -> list[tuple[float]]:
        results = []
        for req in requests:
            (text_str,) = req.args if hasattr(req, "args") else req
            tokens = self._tokenize(text_str)
            if len(tokens) < 2:
                results.append((-float("inf"),))
                continue
            all_lp = self._logprobs_for_tokens(tokens)
            results.append((sum(all_lp),))
        return results


@torch.no_grad()
def _bpb_on_tokens(model: CPUGPT, tokens: torch.Tensor) -> dict:
    model.eval()
    total_loss = 0.0
    total_toks = 0
    n_batches = (len(tokens) - 1) // (SEQ_LEN * BATCH_SIZE)

    t0 = time.perf_counter()
    for i in range(n_batches):
        s = i * SEQ_LEN * BATCH_SIZE
        chunk = tokens[s : s + SEQ_LEN * BATCH_SIZE + 1]
        if len(chunk) < SEQ_LEN * BATCH_SIZE + 1:
            break
        x = chunk[:-1].view(BATCH_SIZE, SEQ_LEN)
        y = chunk[1:].view(BATCH_SIZE, SEQ_LEN)

        loss = model(x, y)
        total_loss += loss.item() * y.numel()
        total_toks += y.numel()

    elapsed = time.perf_counter() - t0
    nll = total_loss / max(total_toks, 1)
    bpb = nll / (math.log(2) * BYTES_PER_TOKEN)
    return {"bpb": bpb, "nll": nll, "tokens": total_toks, "elapsed_s": elapsed}


def _stream_owt_tokens(n_tokens: int, seed: int = 999) -> torch.Tensor:
    try:
        from datasets import load_dataset
    except ImportError:
        raise RuntimeError(
            "datasets library is required for wikitext BPB: pip install datasets"
        )
    import tiktoken

    enc = tiktoken.get_encoding("gpt2")
    print(
        f"Streaming OWT (seed={seed}) for {n_tokens / 1e6:.1f}M tokens...", flush=True
    )
    ds = load_dataset("openwebtext", split="train", streaming=True)
    ds = ds.shuffle(seed=seed, buffer_size=10_000)
    buf: list[int] = []
    for item in ds:
        text = item.get("text", "")
        if text:
            buf.extend(enc.encode_ordinary(text))
        if len(buf) >= n_tokens:
            break
    return torch.tensor(buf[:n_tokens], dtype=torch.long)


def _eval_wikitext_bpb(model: CPUGPT, val_tokens: int = 1_000_000) -> dict:
    tokens = _stream_owt_tokens(val_tokens)
    results = _bpb_on_tokens(model, tokens)
    return results


def _build_hf_wrapper():
    try:
        from transformers import GPT2LMHeadModel, GPT2Tokenizer
    except ImportError:
        print(
            "\nERROR: transformers is not installed (needed for --baseline).\n"
            "Install with:  pip install transformers\n",
            file=sys.stderr,
        )
        sys.exit(1)

    print("Loading HuggingFace gpt2 baseline...", flush=True)
    hf_model = GPT2LMHeadModel.from_pretrained("gpt2")
    hf_model.eval()

    class _HFWrapper:
        def __init__(self, m):
            self.model = m
            self._vocab_size = m.config.vocab_size
            self._seq_len = m.config.n_positions
            import tiktoken

            self._enc = tiktoken.get_encoding("gpt2")

        def _tokenize(self, text: str) -> list[int]:
            return self._enc.encode_ordinary(text)

        @torch.no_grad()
        def _logprobs_for_tokens(self, tokens: list[int]) -> list[float]:
            max_T = self._seq_len
            log_probs: list[float] = []
            start = 0
            while start < len(tokens) - 1:
                end = min(start + max_T + 1, len(tokens))
                chunk = tokens[start:end]
                x = torch.tensor(chunk[:-1], dtype=torch.long).unsqueeze(0)
                tgt = torch.tensor(chunk[1:], dtype=torch.long)
                out = self.model(x)
                logits = out.logits.squeeze(0).float()
                lp = torch.log_softmax(logits, dim=-1)
                tgt_lp = lp[torch.arange(len(tgt)), tgt]
                log_probs.extend(tgt_lp.tolist())
                if end == len(tokens):
                    break
                start = end - 1 - max_T // 2
            return log_probs

        @torch.no_grad()
        def bpb_on_tokens(self, tokens: torch.Tensor) -> dict:
            total_loss = 0.0
            total_toks = 0
            n_batches = (len(tokens) - 1) // (1024 * BATCH_SIZE)
            t0 = time.perf_counter()
            for i in range(n_batches):
                s = i * 1024 * BATCH_SIZE
                chunk = tokens[s : s + 1024 * BATCH_SIZE + 1]
                if len(chunk) < 1024 * BATCH_SIZE + 1:
                    break
                x = chunk[:-1].view(BATCH_SIZE, 1024)
                y = chunk[1:].view(BATCH_SIZE, 1024)
                out = self.model(x)
                logits = out.logits.float()
                loss = torch.nn.functional.cross_entropy(
                    logits.view(-1, logits.size(-1)), y.reshape(-1), reduction="sum"
                )
                total_loss += loss.item()
                total_toks += y.numel()
            elapsed = time.perf_counter() - t0
            nll = total_loss / max(total_toks, 1)
            bpb = nll / (math.log(2) * BYTES_PER_TOKEN)
            return {"bpb": bpb, "nll": nll, "tokens": total_toks, "elapsed_s": elapsed}

    return _HFWrapper(hf_model)


def _run_lm_eval_tasks(wrapper, tasks: list[str], limit: int | None = None) -> dict:
    _require_lm_eval()
    try:
        import lm_eval
        from lm_eval import evaluator
    except ImportError:
        print("lm_eval import failed.", file=sys.stderr)
        sys.exit(1)

    results = evaluator.evaluate(
        lm=wrapper,
        task_dict=lm_eval.tasks.get_task_dict(tasks),
        limit=limit,
    )
    return results.get("results", {})


def _print_markdown_table(our_results: dict, baseline_results: dict | None):
    headers = ["Task", "Metric", "Our Model"]
    if baseline_results:
        headers.append("GPT-2 Baseline")

    rows = []
    for task, metrics in our_results.items():
        for metric, val in metrics.items():
            if isinstance(val, float):
                our_val = f"{val:.4f}"
                base_val = ""
                if baseline_results and task in baseline_results:
                    bv = baseline_results[task].get(metric)
                    base_val = f"{bv:.4f}" if isinstance(bv, float) else str(bv)
                row = [task, metric, our_val]
                if baseline_results:
                    row.append(base_val)
                rows.append(row)

    col_widths = [
        max(len(h), max((len(r[i]) for r in rows), default=0))
        for i, h in enumerate(headers)
    ]

    def _fmt_row(r):
        return (
            "| "
            + " | ".join(cell.ljust(col_widths[i]) for i, cell in enumerate(r))
            + " |"
        )

    sep = "| " + " | ".join("-" * w for w in col_widths) + " |"
    print(_fmt_row(headers))
    print(sep)
    for row in rows:
        print(_fmt_row(row))


def main():
    p = argparse.ArgumentParser(description="Formal eval harness for FNO+GLA model")
    p.add_argument("--ckpt", default=None, help="Checkpoint path (or s3://...)")
    p.add_argument(
        "--config", default="gpt2-small", help="Model config: gpt2-small | smoke"
    )
    p.add_argument(
        "--tasks",
        default="wikitext",
        help="Comma-separated: hellaswag,lambada,piqa,boolq,wikitext",
    )
    p.add_argument("--output", required=True, help="Output JSON path")
    p.add_argument("--baseline", action="store_true", help="Also run HF gpt2 baseline")
    p.add_argument(
        "--limit", type=int, default=None, help="Max examples per task (None=all)"
    )
    p.add_argument(
        "--val-tokens",
        type=int,
        default=1_000_000,
        help="Tokens to use for wikitext BPB (default 1M)",
    )
    args = p.parse_args()

    task_list = [t.strip() for t in args.tasks.split(",") if t.strip()]

    print(f"\n=== run_eval.py ===", flush=True)
    print(f"Config:   {args.config}", flush=True)
    print(f"Ckpt:     {args.ckpt or '(random weights)'}", flush=True)
    print(f"Tasks:    {task_list}", flush=True)
    print(f"Baseline: {args.baseline}", flush=True)
    print(flush=True)

    our_model = _build_our_model(args.config, args.ckpt)
    print(f"Model params: {our_model.param_count() / 1e6:.1f}M", flush=True)

    our_results: dict[str, Any] = {}
    base_results: dict[str, Any] = {}

    if "wikitext" in task_list:
        print("\nRunning wikitext BPB eval (streaming OWT)...", flush=True)
        wt_res = _eval_wikitext_bpb(our_model, val_tokens=args.val_tokens)
        our_results["wikitext"] = {
            "bpb": wt_res["bpb"],
            "nll_nats": wt_res["nll"],
            "tokens_evald": wt_res["tokens"],
        }
        print(f"  wikitext BPB: {wt_res['bpb']:.4f}", flush=True)

        if args.baseline:
            hf = _build_hf_wrapper()
            tokens = _stream_owt_tokens(args.val_tokens)
            bwt = hf.bpb_on_tokens(tokens)
            base_results["wikitext"] = {
                "bpb": bwt["bpb"],
                "nll_nats": bwt["nll"],
                "tokens_evald": bwt["tokens"],
            }
            print(f"  wikitext BPB (baseline): {bwt['bpb']:.4f}", flush=True)

    lm_tasks = [t for t in task_list if t != "wikitext"]
    if lm_tasks:
        _require_lm_eval()
        wrapper = LMEvalWrapper(our_model)
        print(f"\nRunning lm-eval tasks: {lm_tasks}...", flush=True)
        lm_res = _run_lm_eval_tasks(wrapper, lm_tasks, limit=args.limit)
        our_results.update(lm_res)

        if args.baseline:
            hf_wrapper = _build_hf_wrapper()

            _lm_base_cls = type(wrapper).__mro__[1]

            class _HFLMWrapper(_lm_base_cls):
                def __init__(self, hf):
                    super().__init__()
                    import tiktoken

                    self._enc = tiktoken.get_encoding("gpt2")
                    self._hf = hf

                def _tokenize(self, text: str) -> list[int]:
                    return self._enc.encode_ordinary(text)

                @torch.no_grad()
                def _logprobs_for_tokens(self, tokens):
                    return hf_wrapper._logprobs_for_tokens(tokens)

                def loglikelihood(self, requests):
                    results = []
                    for req in requests:
                        ctx_str, cont_str = req.args if hasattr(req, "args") else req
                        ctx_ids = self._tokenize(ctx_str)
                        cont_ids = self._tokenize(cont_str)
                        tokens = ctx_ids + cont_ids
                        if not tokens:
                            results.append((-float("inf"), False))
                            continue
                        all_lp = self._logprobs_for_tokens(tokens)
                        cont_lp = all_lp[-len(cont_ids) :]
                        results.append((sum(cont_lp), False))
                    return results

                def loglikelihood_rolling(self, requests):
                    results = []
                    for req in requests:
                        (text_str,) = req.args if hasattr(req, "args") else req
                        tokens = self._tokenize(text_str)
                        if len(tokens) < 2:
                            results.append((-float("inf"),))
                            continue
                        all_lp = self._logprobs_for_tokens(tokens)
                        results.append((sum(all_lp),))
                    return results

                @torch.no_grad()
                def generate_until(self, requests):
                    results = []
                    for req in requests:
                        ctx_str, gen_kwargs = req.args if hasattr(req, "args") else req
                        until = gen_kwargs.get("until", ["\n"])
                        max_new_toks = gen_kwargs.get("max_gen_toks", 50)
                        ctx_ids = self._tokenize(ctx_str)
                        x = torch.tensor(ctx_ids, dtype=torch.long).unsqueeze(0)
                        out = hf_wrapper.model.generate(
                            x, max_new_tokens=max_new_toks, do_sample=False
                        )
                        gen_ids = out[0, len(ctx_ids) :].tolist()
                        gen_str = self._enc.decode(gen_ids)
                        for stop in until:
                            idx = gen_str.find(stop)
                            if idx != -1:
                                gen_str = gen_str[:idx]
                        results.append(gen_str)
                    return results

            hf_lm_wrapper = _HFLMWrapper(hf_wrapper)
            base_lm_res = _run_lm_eval_tasks(hf_lm_wrapper, lm_tasks, limit=args.limit)
            base_results.update(base_lm_res)

    print("\n" + "=" * 60)
    print("Results")
    print("=" * 60)
    _print_markdown_table(our_results, base_results if args.baseline else None)

    output = {
        "model_path": args.ckpt or "random_weights",
        "config": args.config,
        "tasks": task_list,
        "results": our_results,
    }
    if args.baseline:
        output["baseline_results"] = base_results

    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "w") as f:
        json.dump(output, f, indent=2)
    print(f"\nResults saved to: {out_path}", flush=True)


if __name__ == "__main__":
    main()