"""
Embedding module using a quantized ONNX MiniLM model and a pure Python WordPiece tokenizer.
Calculates cosine similarity and caches column embeddings for performance.
"""
from __future__ import annotations

import os
import urllib.request
import numpy as np
import onnxruntime as ort
import re
from typing import Sequence

# Default Hugging Face URLs for the quantized all-MiniLM-L6-v2 model and vocab
MODEL_URL = "https://huggingface.co/onnx-community/all-MiniLM-L6-v2-ONNX/resolve/main/onnx/model_quantized.onnx"
VOCAB_URL = "https://huggingface.co/onnx-community/all-MiniLM-L6-v2-ONNX/resolve/main/vocab.txt"

class WordPieceTokenizer:
    """Pure Python implementation of a WordPiece tokenizer."""
    def __init__(self, vocab_path: str):
        self.vocab: dict[str, int] = {}
        with open(vocab_path, "r", encoding="utf-8") as f:
            for i, line in enumerate(f):
                token = line.strip()
                self.vocab[token] = i

        self.unk_token = "[UNK]"
        self.cls_token = "[CLS]"
        self.sep_token = "[SEP]"

        self.unk_id = self.vocab.get(self.unk_token, 100)
        self.cls_id = self.vocab.get(self.cls_token, 101)
        self.sep_id = self.vocab.get(self.sep_token, 102)

    def tokenize_word(self, word: str) -> list[str]:
        """Tokenize a single word into WordPiece subwords."""
        if word in self.vocab:
            return [word]

        tokens = []
        start = 0
        is_bad = False
        while start < len(word):
            end = len(word)
            cur_substr = None
            while start < end:
                substr = word[start:end]
                if start > 0:
                    substr = "##" + substr
                if substr in self.vocab:
                    cur_substr = substr
                    break
                end -= 1
            if cur_substr is None:
                is_bad = True
                break
            tokens.append(cur_substr)
            start = end

        if is_bad:
            return [self.unk_token]
        return tokens

    def encode(self, text: str, max_length: int = 128) -> dict[str, np.ndarray]:
        """Encode text into model inputs (input_ids, attention_mask, token_type_ids)."""
        text = text.lower()
        # Basic word and punctuation splitter
        words = re.findall(r"\w+|[^\w\s]", text, re.UNICODE)
        tokens = []
        for word in words:
            tokens.extend(self.tokenize_word(word))

        # Truncate
        if len(tokens) > max_length - 2:
            tokens = tokens[:max_length - 2]

        # Build token IDs
        input_ids = [self.cls_id] + [self.vocab.get(t, self.unk_id) for t in tokens] + [self.sep_id]
        attention_mask = [1] * len(input_ids)
        token_type_ids = [0] * len(input_ids)

        # Pad to max_length
        padding_len = max_length - len(input_ids)
        if padding_len > 0:
            input_ids.extend([0] * padding_len)
            attention_mask.extend([0] * padding_len)
            token_type_ids.extend([0] * padding_len)

        # Convert to numpy arrays of type int64 (as expected by ONNX model)
        return {
            "input_ids": np.array([input_ids], dtype=np.int64),
            "attention_mask": np.array([attention_mask], dtype=np.int64),
            "token_type_ids": np.array([token_type_ids], dtype=np.int64),
        }

def cosine_similarity(v1: np.ndarray, v2: np.ndarray) -> float:
    """Calculate the cosine similarity between two 1D vectors."""
    dot = np.dot(v1, v2)
    norm1 = np.linalg.norm(v1)
    norm2 = np.linalg.norm(v2)
    if norm1 == 0 or norm2 == 0:
        return 0.0
    return float(dot / (norm1 * norm2))

class EmbeddingModel:
    """ONNX-based text embedding generator with built-in WordPiece tokenization."""
    def __init__(self, cache_dir: str | None = None):
        if cache_dir is None:
            # First try zero-llm-engine/models directory, fallback to temp dir /tmp/models
            base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
            cache_dir = os.path.join(base_dir, "zero-llm-engine", "models", "onnx_cache")
            if not os.path.exists(cache_dir):
                try:
                    os.makedirs(cache_dir, exist_ok=True)
                except Exception:
                    # Fallback if the folder is read-only (e.g. some system settings or Docker environments)
                    cache_dir = os.path.join("/tmp", "onnx_cache")
                    os.makedirs(cache_dir, exist_ok=True)

        self.cache_dir = cache_dir
        self.model_path = os.path.join(cache_dir, "model_quantized.onnx")
        self.vocab_path = os.path.join(cache_dir, "vocab.txt")
        self.session: ort.InferenceSession | None = None
        self.tokenizer: WordPieceTokenizer | None = None
        
        # In-memory embedding cache: maps text string -> np.ndarray embedding
        self._embedding_cache: dict[str, np.ndarray] = {}

    def ensure_model_files(self) -> None:
        """Download model and vocabulary files if they do not exist locally."""
        if not os.path.exists(self.vocab_path):
            print(f"[Embeddings] Downloading vocabulary to {self.vocab_path}...")
            urllib.request.urlretrieve(VOCAB_URL, self.vocab_path)
        if not os.path.exists(self.model_path):
            print(f"[Embeddings] Downloading quantized ONNX model to {self.model_path}...")
            urllib.request.urlretrieve(MODEL_URL, self.model_path)

    def load_model(self) -> None:
        """Ensure files are downloaded and load the ONNX session and tokenizer."""
        if self.session is not None and self.tokenizer is not None:
            return
            
        self.ensure_model_files()
        
        # Initialize tokenization & ONNX session
        self.tokenizer = WordPieceTokenizer(self.vocab_path)
        # Using CPU execution provider by default for maximum compatibility
        self.session = ort.InferenceSession(self.model_path, providers=["CPUExecutionProvider"])

    def get_embedding(self, text: str) -> np.ndarray:
        """Generate a 1D mean-pooled normalized embedding vector for the text."""
        self.load_model()
        text_key = text.lower().strip()
        if text_key in self._embedding_cache:
            return self._embedding_cache[text_key]

        assert self.tokenizer is not None
        assert self.session is not None

        # Tokenize and format inputs
        inputs = self.tokenizer.encode(text_key)
        
        # Run ONNX inference
        outputs = self.session.run(None, inputs)
        # The first output contains the token embeddings [batch_size, seq_len, hidden_dim]
        token_embeddings = outputs[0]
        attention_mask = inputs["attention_mask"]

        # Perform mean pooling over the active tokens
        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
        input_mask_expanded = np.broadcast_to(input_mask_expanded, token_embeddings.shape)
        
        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
        sum_mask = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
        
        # Calculate mean pooled embedding (vector shape: [hidden_dim])
        mean_pooled = (sum_embeddings / sum_mask)[0]
        
        # L2 Normalize
        norm = np.linalg.norm(mean_pooled)
        if norm > 0:
            mean_pooled = mean_pooled / norm
            
        self._embedding_cache[text_key] = mean_pooled
        return mean_pooled

    def match_column(self, text: str, columns: Sequence[str], threshold: float = 0.4) -> tuple[str | None, float]:
        """Match query text to the best column name using cosine similarity."""
        if not columns:
            return None, 0.0

        query_emb = self.get_embedding(text)
        best_col = None
        best_sim = -1.0

        for col in columns:
            col_emb = self.get_embedding(col)
            sim = cosine_similarity(query_emb, col_emb)
            if sim > best_sim:
                best_sim = sim
                best_col = col

        if best_sim >= threshold:
            return best_col, best_sim
        return None, best_sim