# bclm/__init__.py
"""
bclm-pytorch — Load and run BCLM language models.

    import bclm

    model = bclm.load("bclm-1-small-preview")

    chat = model.chat(system_prompt="You are a helpful assistant.")
    print(chat.send("Hello!"))
"""
from __future__ import annotations

import sys
from typing import Iterator, List, Optional, Union

import torch

from .config import (
    BCLM1ModelConfig,
    EOT_TOKEN_ID,
    PAD_TOKEN_ID,
)
from .registry import (
    register_architecture,
    get_architecture,
    list_architectures,
)
from .hf_loader import load_model
from .tokenizer import BCLMTokenizer
from .inference import ChatSession


__all__ = [
    "load",
    "BCLMModel",
    "register_architecture",
    "list_architectures",
]


# =============================================================================
# High-level model wrapper
# =============================================================================

class BCLMModel:
    """
    High-level interface to a loaded BCLM model.

    Returned by :func:`bclm.load`. Provides chat, completion, and
    streaming methods, as well as access to the underlying components.

    Examples::

        model = bclm.load("bclm-1-small-preview")

        # Chat
        chat = model.chat(system_prompt="You are helpful.")
        print(chat.send("What is 2+2?"))

        # Streaming chat
        for chunk in chat.send_stream("Tell me a story."):
            print(chunk, end="", flush=True)

        # Completion
        print(model.complete("Once upon a time"))

        # Streaming completion
        for chunk in model.complete_stream("The quick brown fox"):
            print(chunk, end="", flush=True)
    """

    def __init__(
        self,
        raw_model: torch.nn.Module,
        generator: torch.nn.Module,
        config,
        tokenizer: BCLMTokenizer,
        architecture: str,
        device: torch.device,
        dtype: torch.dtype,
    ) -> None:
        self._raw_model = raw_model
        self._generator = generator
        self._config = config
        self._tokenizer = tokenizer
        self._architecture = architecture
        self._device = device
        self._dtype = dtype

    # ── Properties ──────────────────────────────────────────

    @property
    def raw_model(self) -> torch.nn.Module:
        """The underlying ``nn.Module`` (e.g. ``BCLM1Model``)."""
        return self._raw_model

    @property
    def generator(self) -> torch.nn.Module:
        """The inference wrapper (e.g. ``BCLM1ForGeneration``)."""
        return self._generator

    @property
    def config(self):
        """Model configuration dataclass."""
        return self._config

    @property
    def tokenizer(self) -> BCLMTokenizer:
        """The tokenizer instance."""
        return self._tokenizer

    @property
    def architecture(self) -> str:
        """Architecture name from ``config.json``."""
        return self._architecture

    @property
    def device(self) -> torch.device:
        """Device the model resides on."""
        return self._device

    @property
    def dtype(self) -> torch.dtype:
        """Weight dtype."""
        return self._dtype

    @property
    def num_parameters(self) -> int:
        """Total number of scalar parameters."""
        return sum(p.numel() for p in self._raw_model.parameters())

    # ── Chat ────────────────────────────────────────────────

    def chat(
        self,
        *,
        system_prompt: Optional[str] = None,
        max_new_tokens: int = 512,
        temperature: float = 0.7,
        top_k: Optional[int] = 50,
        top_p: Optional[float] = 0.9,
    ) -> ChatSession:
        """
        Create a multi-turn chat session.

        Returns a :class:`~bclm.inference.ChatSession` that tracks
        conversation history and supports both blocking and streaming
        responses.

        Parameters
        ----------
        system_prompt : str, optional
            Persistent system instruction prepended to every request.
        max_new_tokens : int
            Default maximum tokens per response.
        temperature : float
            Sampling temperature (0 = greedy).
        top_k : int, optional
            Top-k filtering.
        top_p : float, optional
            Nucleus sampling threshold.
        """
        return ChatSession(
            self._generator,
            self._tokenizer,
            system_prompt=system_prompt,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
        )

    # ── Completion ──────────────────────────────────────────

    def _encode_prompt(self, text: str) -> torch.Tensor:
        ids = self._tokenizer.encode(text)
        return torch.tensor([ids], dtype=torch.long, device=self._device)

    def complete(
        self,
        prompt: str,
        *,
        max_tokens: int = 256,
        temperature: float = 0.8,
        top_k: Optional[int] = 50,
        top_p: Optional[float] = 0.9,
    ) -> str:
        """
        Continue a text prompt and return the generated text.

        Parameters
        ----------
        prompt : str
            The text to continue.
        max_tokens : int
            Maximum number of tokens to generate.
        temperature : float
            Sampling temperature.
        top_k : int, optional
            Top-k filtering.
        top_p : float, optional
            Nucleus sampling threshold.

        Returns
        -------
        str
            The generated continuation (prompt is **not** included).
        """
        prompt_tokens = self._encode_prompt(prompt)

        response_ids: List[int] = []
        for tid in self._generator.generate_stream(
            prompt_tokens,
            max_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            stop_token=EOT_TOKEN_ID,
        ):
            if tid == EOT_TOKEN_ID:
                break
            response_ids.append(tid)

        return self._tokenizer.decode_response(response_ids)

    def complete_stream(
        self,
        prompt: str,
        *,
        max_tokens: int = 256,
        temperature: float = 0.8,
        top_k: Optional[int] = 50,
        top_p: Optional[float] = 0.9,
    ) -> Iterator[str]:
        """
        Continue a text prompt and yield text chunks as they are generated.

        Parameters
        ----------
        prompt : str
            The text to continue.
        max_tokens : int
            Maximum number of tokens to generate.
        temperature : float
            Sampling temperature.
        top_k : int, optional
            Top-k filtering.
        top_p : float, optional
            Nucleus sampling threshold.

        Yields
        ------
        str
            Incremental text chunks.
        """
        prompt_tokens = self._encode_prompt(prompt)

        token_stream = self._generator.generate_stream(
            prompt_tokens,
            max_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            stop_token=EOT_TOKEN_ID,
        )

        yield from self._tokenizer.stream_decode(token_stream)

    # ── Repr ────────────────────────────────────────────────

    def __repr__(self) -> str:
        return (
            f"BCLMModel(\n"
            f"  architecture={self._architecture!r},\n"
            f"  parameters={self.num_parameters / 1e6:.2f}M,\n"
            f"  device={self._device},\n"
            f"  dtype={self._dtype},\n"
            f"  tokenizer={self._tokenizer!r},\n"
            f")"
        )


# =============================================================================
# Main entry point
# =============================================================================

def load(
    source: str,
    *,
    device: Optional[Union[str, torch.device]] = None,
    dtype: Optional[torch.dtype] = None,
    compile: bool = False,
    cache_dir: Optional[str] = None,
) -> BCLMModel:
    """
    Load a BCLM model and return a ready-to-use :class:`BCLMModel`.

    Parameters
    ----------
    source : str
        Where to load from. Accepts any of:

        - **Short name**: ``"bclm-1-small-preview"`` — resolved to the
          ``bclm/`` org on Hugging Face.
        - **Full HF repo**: ``"bclm/bclm-1-small-preview"``
        - **Local directory**: ``"/path/to/model"`` — must contain
          ``config.json`` and ``model.safetensors``.
        - **HTTPS URL**: ``"https://example.com/models/bclm-1/"`` — the
          files ``config.json`` and ``model.safetensors`` are fetched
          from this base URL.

    device : str or torch.device, optional
        Target device. Defaults to CUDA if available, otherwise CPU.
    dtype : torch.dtype, optional
        Weight dtype. Defaults to ``torch.bfloat16`` on CUDA (matching
        training precision) and ``torch.float32`` on CPU.
    compile : bool
        Enable ``torch.compile`` for inference. Off by default to avoid
        warmup latency.
    cache_dir : str, optional
        Cache directory for downloaded files.

    Returns
    -------
    BCLMModel
        A high-level wrapper with ``.chat()``, ``.complete()``, and
        ``.complete_stream()`` methods.

    Examples
    --------
    ::

        import bclm

        model = bclm.load("bclm-1-small-preview")

        # Quick chat
        chat = model.chat()
        print(chat.send("Hello!"))

        # Streaming completion
        for chunk in model.complete_stream("Once upon a time"):
            print(chunk, end="", flush=True)
    """
    result = load_model(
        source,
        device=device,
        dtype=dtype,
        compile=compile,
        cache_dir=cache_dir,
    )

    # ── Load tokenizer ──────────────────────────────────────

    tokenizer_spec = result["tokenizer_spec"]
    model_config = result["config"]
    vocab_size = getattr(model_config, "vocab_size", 32768)

    try:
        tokenizer = BCLMTokenizer(
            spec=tokenizer_spec,
            vocab_size=vocab_size,
        )
    except ImportError as exc:
        print(
            f"[bclm] Tokenizer unavailable: {exc}\n"
            f"[bclm] Install the tokenizer package to enable text generation.",
            file=sys.stderr,
        )
        raise
    except Exception as exc:
        print(
            f"[bclm] Failed to load tokenizer ({tokenizer_spec!r}): {exc}",
            file=sys.stderr,
        )
        raise

    return BCLMModel(
        raw_model=result["model"],
        generator=result["generator"],
        config=result["config"],
        tokenizer=tokenizer,
        architecture=result["architecture"],
        device=result["device"],
        dtype=result["dtype"],
    )