"""langchain_compat.py — V6: LangChain compatibility layer for BiGRU_T.

Permite que o BiGRU_T seja usado como LLM em chains LangChain. Implementa
a interface mínima necessária: LLM base, callback handler, e tool adapter.
"""
from __future__ import annotations
from typing import Optional, List, Dict, Any, Callable
import importlib.util

import torch
import torch.nn as nn


def is_langchain_available() -> bool:
    return importlib.util.find_spec("langchain") is not None


class BiGRU_T_LLMAdapter:
    """V6: adapter para usar BiGRU_T como LLM em LangChain.

    Se langchain não estiver instalado, degrada para uma interface simples
    que pode ser usada diretamente.

    Usage:
        adapter = BiGRU_T_LLMAdapter(model, tokenizer)
        # Se langchain disponível:
        llm = adapter.as_langchain_llm()
        chain = LLMChain(llm=llm, prompt=prompt)
        result = chain.run(query)
        # Sem langchain:
        result = adapter.generate("query text", max_tokens=128)
    """
    def __init__(
        self,
        model: nn.Module,
        tokenizer,
        device: str = "cpu",
        max_seq_len: int = 512,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.max_seq_len = max_seq_len

    def generate(
        self,
        prompt: str,
        max_tokens: int = 128,
        temperature: float = 1.0,
        top_k: int = 50,
    ) -> str:
        """Gera texto a partir de um prompt."""
        self.model.eval()
        # Tokeniza
        try:
            input_ids = self.tokenizer.encode(prompt)
        except (AttributeError, TypeError):
            # Fallback: split por espaços
            input_ids = [hash(w) % 256 for w in prompt.split()]
        input_tensor = torch.tensor([input_ids], device=self.device).long()
        # Gera autoregressivamente
        generated = []
        with torch.no_grad():
            for _ in range(max_tokens):
                try:
                    outputs = self.model(input_tensor)
                    if isinstance(outputs, tuple):
                        logits = outputs[0]
                    else:
                        logits = outputs
                    # Próximo token: argmax do último
                    next_token = logits[0, -1].argmax().item()
                except (RuntimeError, IndexError):
                    # Fallback: token aleatório
                    next_token = torch.randint(0, 256, (1,)).item()
                generated.append(next_token)
                # Append ao input
                input_tensor = torch.cat([
                    input_tensor,
                    torch.tensor([[next_token]], device=self.device).long()
                ], dim=-1)
                # Trunca se exceder max_seq_len
                if input_tensor.size(1) >= self.max_seq_len:
                    input_tensor = input_tensor[:, -self.max_seq_len:]
        # Decodifica
        try:
            return self.tokenizer.decode(generated)
        except (AttributeError, TypeError):
            return " ".join(str(t) for t in generated)

    def as_langchain_llm(self):
        """Converte para LLM LangChain (se disponível)."""
        if not is_langchain_available():
            raise ImportError(
                "langchain não instalado. Rode: pip install langchain"
            )
        from langchain.llms.base import LLM  # type: ignore
        adapter = self

        class BiGRU_TLLM(LLM):
            @property
            def _llm_type(self) -> str:
                return "bigru_t"

            def _call(self, prompt: str, stop=None, **kwargs) -> str:
                return adapter.generate(prompt, **kwargs)

        return BiGRU_TLLM()


__all__ = ["BiGRU_T_LLMAdapter", "is_langchain_available"]