"""
Frox AI Morph 1.1 — Tokenizer
Fully local, no gated repos.

Improvements over Morph 1.0:
  - 64K vocab BPE (better multilingual, math, code)
  - More special tokens (reasoning, citations, memory, search)
  - DPO formatting helpers (chosen/rejected pairs)
  - Chat template for system / user / assistant / tool roles
  - Token budget estimation utility
"""
from __future__ import annotations
import os
import tempfile
from pathlib import Path
from typing import Dict, List, Optional, Sequence

# ── Special tokens ────────────────────────────────────────────────

CHAT_SPECIAL_TOKENS = [
    "<|pad|>",
    "<|begin_of_text|>",
    "<|end_of_text|>",
    "<|unk|>",
    "<|eot_id|>",
    "<|start_header_id|>",
    "<|end_header_id|>",
    "<|finetune_right_pad_id|>",
    "<|step_id|>",
    "<|python_tag|>",
]

MORPH_SPECIAL_TOKENS = [
    # Multimodal
    "<|image|>",
    "<|video|>",
    "<|audio|>",
    # Generation requests
    "<|gen_image|>",
    "<|gen_video|>",
    "<|gen_3d|>",
    "<|/gen|>",
    # Tool calling
    "<|tool_call|>",
    "<|/tool_call|>",
    "<|tool_result|>",
    "<|/tool_result|>",
    # Reasoning / CoT (NEW 1.1)
    "<|think|>",
    "<|/think|>",
    # Citations (NEW 1.1)
    "<|cite|>",
    "<|/cite|>",
    # Memory (NEW 1.1)
    "<|memory|>",
    "<|/memory|>",
    # Search context (NEW 1.1)
    "<|search_result|>",
    "<|/search_result|>",
    # Code execution (NEW 1.1)
    "<|code_output|>",
    "<|/code_output|>",
    # Fill-in-middle, code infilling (NEW 1.1 — Morph Code tier)
    "<|fim_prefix|>",
    "<|fim_middle|>",
    "<|fim_suffix|>",
    # Browser tool (NEW — see tools/web_browser.py)
    "<|browser_result|>",
    "<|/browser_result|>",
]

ALL_SPECIAL_TOKENS = CHAT_SPECIAL_TOKENS + MORPH_SPECIAL_TOKENS

# IDs (must align with vocab_size=64000 in config)
PAD_TOKEN_ID   = 0
BOS_TOKEN_ID   = 1
EOS_TOKEN_ID   = 2
UNK_TOKEN_ID   = 3
IMAGE_TOKEN_ID = len(CHAT_SPECIAL_TOKENS)       # 10
VIDEO_TOKEN_ID = len(CHAT_SPECIAL_TOKENS) + 1   # 11


# ── Builder ───────────────────────────────────────────────────────

def build_morph_tokenizer(
    tokenizer_path: Optional[str] = None,
    corpus_files: Optional[Sequence[str]] = None,
    save_path: Optional[str] = None,
    vocab_size: int = 64000,
):
    """
    Build Frox Morph 1.1 tokenizer — fully local, zero auth.

    Priority:
      1. Load from saved tokenizer_path (fastest)
      2. Train BPE from provided corpus_files
      3. Auto-download free corpus (WikiText-103, ~180MB) and train
         Falls back to synthetic corpus if download fails.

    Args:
        tokenizer_path: Previously saved tokenizer dir (HF format).
        corpus_files:   .txt / .jsonl files for BPE training.
        save_path:      Where to save after building.
        vocab_size:     BPE vocab size (64K default for 1.1).
    """
    print(f"\n🔤 Building Morph 1.1 Tokenizer (vocab_size={vocab_size:,})")

    # 1. Load saved
    if tokenizer_path and Path(tokenizer_path).exists():
        from transformers import AutoTokenizer
        tok = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=True)
        _ensure_special_tokens(tok)
        print(f"  ✓ Loaded from: {tokenizer_path} ({len(tok):,} tokens)")
        if save_path:
            _save(tok, save_path)
        return tok

    # 2. Train from provided corpus
    if corpus_files:
        valid = [p for p in corpus_files if Path(p).exists()]
        if not valid:
            raise FileNotFoundError(f"No corpus files found: {corpus_files}")
        print(f"  Training BPE from {len(valid)} file(s)...")
        tok_model = _train_bpe(valid, vocab_size)
        tok = _wrap_tokenizer(tok_model)
        print(f"  ✓ Trained ({len(tok):,} tokens)")
        if save_path:
            _save(tok, save_path)
        return tok

    # 3. Auto-download
    print("  No corpus provided — downloading WikiText-103...")
    corpus_path = _download_corpus()
    tok_model = _train_bpe([corpus_path], vocab_size)
    tok = _wrap_tokenizer(tok_model)
    print(f"  ✓ Built ({len(tok):,} tokens)")
    if save_path:
        _save(tok, save_path)
    return tok


# ── Corpus download ───────────────────────────────────────────────

def _download_corpus() -> str:
    """Download WikiText-103 train split (public domain, ~180MB)."""
    import urllib.request

    save_dir = Path(tempfile.gettempdir()) / "morph11_corpus"
    save_dir.mkdir(exist_ok=True)
    out = save_dir / "wikitext103.txt"

    if out.exists() and out.stat().st_size > 50_000_000:
        print(f"  Using cached: {out}")
        return str(out)

    # WikiText-103 raw (no auth, MIT-licensed)
    url = "https://huggingface.co/datasets/Salesforce/wikitext/resolve/main/wikitext-103-raw-v1/train-00000-of-00002.parquet"
    alt = "https://raw.githubusercontent.com/pytorch/examples/main/word_language_model/data/wikitext-2/train.txt"

    try:
        print(f"  Downloading WikiText-103...")
        # Try to use datasets library for parquet (much larger corpus)
        from datasets import load_dataset
        ds = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1",
                          split="train", trust_remote_code=False)
        with open(out, "w", encoding="utf-8") as f:
            for ex in ds:
                t = ex.get("text", "").strip()
                if t:
                    f.write(t + "\n")
        size_mb = out.stat().st_size / 1e6
        print(f"  ✓ Downloaded {size_mb:.1f} MB")
        return str(out)
    except Exception as e:
        print(f"  datasets lib failed ({e}), falling back to WikiText-2...")

    try:
        urllib.request.urlretrieve(alt, str(out))
        return str(out)
    except Exception as e:
        print(f"  Download failed ({e}) — using synthetic corpus")
        return _make_synthetic_corpus(save_dir)


def _make_synthetic_corpus(save_dir: Path) -> str:
    """Minimal synthetic corpus so training never crashes."""
    path = save_dir / "synthetic.txt"
    lines = [
        "The quick brown fox jumps over the lazy dog.",
        "Artificial intelligence is reshaping every industry.",
        "Language models learn patterns from large text datasets.",
        "Neural networks compute features through multiple layers.",
        "Frox AI Morph 1.1 is a powerful multimodal assistant.",
        "You can ask me to generate images, videos, and 3D objects.",
        "Reasoning models can solve complex mathematical problems.",
        "The transformer architecture uses self-attention mechanisms.",
        "Tokenization converts raw text into subword tokens.",
        "Byte pair encoding merges frequent character pairs iteratively.",
        "Training requires compute, data, and careful hyperparameter tuning.",
        "Supervised fine-tuning teaches models to follow instructions.",
        "Direct preference optimization improves alignment with human values.",
        "Retrieval augmented generation grounds answers in real documents.",
    ] * 10000
    path.write_text("\n".join(lines), encoding="utf-8")
    print(f"  Synthetic corpus: {len(lines):,} sentences")
    return str(path)


# ── BPE training ──────────────────────────────────────────────────

def _train_bpe(corpus_paths: List[str], vocab_size: int):
    """Train byte-level BPE from corpus files."""
    from tokenizers import Tokenizer
    from tokenizers.models import BPE
    from tokenizers.pre_tokenizers import ByteLevel
    from tokenizers.decoders import ByteLevel as ByteLevelDecoder
    from tokenizers.trainers import BpeTrainer
    from tokenizers.normalizers import NFKC

    tok = Tokenizer(BPE(unk_token="<|unk|>"))
    tok.normalizer   = NFKC()
    tok.pre_tokenizer = ByteLevel(add_prefix_space=False)
    tok.decoder       = ByteLevelDecoder()

    trainer = BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=ALL_SPECIAL_TOKENS,
        initial_alphabet=ByteLevel.alphabet(),
        show_progress=True,
        min_frequency=2,
    )
    tok.train(corpus_paths, trainer)
    return tok


def _wrap_tokenizer(tok_model):
    """Wrap tokenizers.Tokenizer into HuggingFace PreTrainedTokenizerFast."""
    from transformers import PreTrainedTokenizerFast
    tok = PreTrainedTokenizerFast(
        tokenizer_object=tok_model,
        bos_token="<|begin_of_text|>",
        eos_token="<|end_of_text|>",
        unk_token="<|unk|>",
        pad_token="<|pad|>",
        additional_special_tokens=ALL_SPECIAL_TOKENS,
    )
    tok.chat_template = MORPH_CHAT_TEMPLATE
    return tok


def _ensure_special_tokens(tok):
    """Ensure all Morph 1.1 special tokens exist in the loaded tokenizer."""
    missing = [t for t in ALL_SPECIAL_TOKENS if t not in tok.get_vocab()]
    if missing:
        tok.add_special_tokens({"additional_special_tokens": missing})
        print(f"  Added {len(missing)} missing special tokens")
    for attr, token in [("pad_token","<|pad|>"), ("bos_token","<|begin_of_text|>"),
                         ("eos_token","<|end_of_text|>")]:
        if getattr(tok, attr) is None:
            tok.add_special_tokens({attr: token})


def _save(tok, save_path: str):
    Path(save_path).mkdir(parents=True, exist_ok=True)
    tok.save_pretrained(save_path)
    print(f"  ✓ Saved to: {save_path}")


# ── Chat templates ────────────────────────────────────────────────

MORPH_CHAT_TEMPLATE = (
    "{%- set ns = namespace(system='') -%}"
    "{%- if messages[0]['role'] == 'system' -%}"
    "{%- set ns.system = messages[0]['content'] -%}"
    "{%- set messages = messages[1:] -%}"
    "{%- endif -%}"
    "<|begin_of_text|>"
    "{%- if ns.system -%}"
    "<|start_header_id|>system<|end_header_id|>\n"
    "{{ ns.system }}<|eot_id|>"
    "{%- endif -%}"
    "{%- for message in messages -%}"
    "<|start_header_id|>{{ message['role'] }}<|end_header_id|>\n"
    "{{ message['content'] }}<|eot_id|>"
    "{%- endfor -%}"
    "{%- if add_generation_prompt -%}"
    "<|start_header_id|>assistant<|end_header_id|>\n"
    "{%- endif -%}"
)


def apply_chat_template(
    messages: List[Dict[str, str]],
    tokenizer,
    add_generation_prompt: bool = True,
    system_prompt: Optional[str] = None,
) -> str:
    """Apply Morph 1.1 chat template to a message list."""
    if system_prompt and (not messages or messages[0]["role"] != "system"):
        messages = [{"role": "system", "content": system_prompt}] + list(messages)

    try:
        return tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=add_generation_prompt,
        )
    except Exception:
        # Manual fallback
        parts = ["<|begin_of_text|>"]
        for msg in messages:
            parts.append(
                f"<|start_header_id|>{msg['role']}<|end_header_id|>\n"
                f"{msg['content']}<|eot_id|>"
            )
        if add_generation_prompt:
            parts.append("<|start_header_id|>assistant<|end_header_id|>\n")
        return "".join(parts)


def format_tool_call(tool_name: str, args: dict) -> str:
    """Format a tool call in Morph 1.1 format."""
    import json
    return f"<|tool_call|>{json.dumps({'name': tool_name, 'args': args})}<|/tool_call|>"


def format_tool_result(tool_name: str, result) -> str:
    """Format a tool result for injection back to the model."""
    import json
    return f"<|tool_result|>{json.dumps({'name': tool_name, 'result': result})}<|/tool_result|>"


def format_thinking(thinking: str) -> str:
    """Format chain-of-thought reasoning block."""
    return f"<|think|>{thinking}<|/think|>"


def format_fim(prefix: str, suffix: str) -> str:
    """
    Format a fill-in-middle prompt (PSM order: prefix, suffix, then the
    model generates the middle). Used by Morph Code for code infilling —
    e.g. completing a function body given the code before and after it.
    """
    return f"<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>"


# ── DPO formatting (NEW 1.1) ──────────────────────────────────────

def format_dpo_pair(
    messages: List[Dict[str, str]],
    chosen: str,
    rejected: str,
    tokenizer,
    system_prompt: Optional[str] = None,
) -> Dict[str, str]:
    """
    Format a preference pair for DPO training.

    Returns:
        {
          "prompt": "<prompt without final assistant turn>",
          "chosen": "<prompt + chosen response>",
          "rejected": "<prompt + rejected response>",
        }
    """
    prompt = apply_chat_template(
        messages, tokenizer, add_generation_prompt=True,
        system_prompt=system_prompt,
    )
    assistant_suffix = "<|start_header_id|>assistant<|end_header_id|>\n"

    return {
        "prompt":   prompt,
        "chosen":   prompt + chosen + "<|eot_id|>",
        "rejected": prompt + rejected + "<|eot_id|>",
    }


# ── Token budget ──────────────────────────────────────────────────

def count_tokens(text: str, tokenizer) -> int:
    """Fast token count without full encode."""
    return len(tokenizer.encode(text, add_special_tokens=False))


def estimate_cost_tokens(
    messages: List[Dict[str, str]],
    max_new_tokens: int,
    tokenizer,
) -> Dict[str, int]:
    """Estimate total tokens for a request."""
    prompt_str = apply_chat_template(messages, tokenizer, add_generation_prompt=True)
    prompt_tokens = count_tokens(prompt_str, tokenizer)
    return {
        "prompt_tokens": prompt_tokens,
        "max_completion": max_new_tokens,
        "total_max": prompt_tokens + max_new_tokens,
    }