"""
DeepX Embedding Server — Deploy on RTX 3060 12GB.

Provides REST API for text embedding.
Loads INT8 quantized model for efficient inference.

Usage:
  python scripts/serve_embedding.py \
    --checkpoint checkpoints/deploy/deepx_int8.pt \
    --port 8080

API:
  POST /embed
  Body: {"texts": ["text1", "text2", ...], "normalize": true}
  Response: {"embeddings": [[...], [...]], "dim": 1536}

  POST /similarity
  Body: {"query": "...", "documents": ["doc1", "doc2", ...]}
  Response: {"scores": [0.85, 0.72, ...]}

  GET /health
  Response: {"status": "ok", "model": "deepx-v0.7", "device": "cuda"}
"""
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import torch
import torch.nn.functional as F
import argparse
import logging
import time
from typing import List
from flask import Flask, request, jsonify
from transformers import AutoTokenizer
from config import DeepXConfig
from modeling.pipeline import DeepXPipeline

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Globals
pipeline = None
tokenizer = None
device = None
MAX_LEN = 2048
MAX_BATCH = 32


def dequantize_int8(state_dict):
    """Dequantize int8 back to float16."""
    dequantized = {}
    scale_keys = {k for k in state_dict if k.endswith("._scale")}
    for key, tensor in state_dict.items():
        if key in scale_keys:
            continue
        scale_key = key + "._scale"
        if tensor.dtype == torch.int8 and scale_key in state_dict:
            scale = state_dict[scale_key].float()
            dequantized[key] = (tensor.float() * scale).half()
        else:
            dequantized[key] = tensor
    return dequantized


def load_model(checkpoint_path, tokenizer_path, device_str="cuda"):
    global pipeline, tokenizer, device
    
    device = torch.device(device_str if torch.cuda.is_available() else "cpu")
    logger.info(f"Device: {device}")
    
    # Tokenizer
    logger.info(f"Loading tokenizer from {tokenizer_path}...")
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # Model
    logger.info(f"Loading model from {checkpoint_path}...")
    config = DeepXConfig()
    pipeline = DeepXPipeline(config, embed_path="pretrained/gemma4_e2b_embed.pt")
    
    ckpt = torch.load(checkpoint_path, map_location="cpu")
    sd = ckpt.get("model_state_dict", ckpt)
    
    # Dequantize if INT8
    if ckpt.get("quantized", False):
        logger.info("Dequantizing INT8 model...")
        sd = dequantize_int8(sd)
    
    pipeline.backbone.load_state_dict(sd, strict=False)
    pipeline = pipeline.to(device).half().eval()
    
    # Force pure GDN-2
    for m in pipeline.modules():
        if hasattr(m, 'path_mix_logit'):
            m._alpha_override = 0.0
    
    logger.info("Model loaded successfully!")
    logger.info(f"  VRAM: {torch.cuda.memory_allocated()/1024**2:.0f} MB" if device.type == "cuda" else "  CPU mode")


@torch.no_grad()
def encode_texts(texts: List[str], normalize: bool = True) -> torch.Tensor:
    """Encode texts to embeddings."""
    # Tokenize
    encoded = tokenizer(
        texts, padding=True, truncation=False, return_tensors="pt", max_length=MAX_LEN
    )
    
    # Skip texts too long
    input_ids = encoded["input_ids"]
    attention_mask = encoded["attention_mask"]
    
    # Trim to max actual length (save compute)
    max_len = attention_mask.sum(dim=1).max().item()
    max_len = min(max_len, MAX_LEN)
    input_ids = input_ids[:, :max_len].to(device)
    attention_mask = attention_mask[:, :max_len].to(device)
    
    # Encode in batches
    all_embs = []
    for i in range(0, len(texts), MAX_BATCH):
        batch_ids = input_ids[i:i+MAX_BATCH]
        batch_mask = attention_mask[i:i+MAX_BATCH]
        
        with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
            emb = pipeline(batch_ids, attention_mask=batch_mask, normalize=normalize)
        all_embs.append(emb.cpu())
    
    return torch.cat(all_embs)


@app.route("/health", methods=["GET"])
def health():
    return jsonify({
        "status": "ok",
        "model": "deepx-v0.7-gdn2",
        "device": str(device),
        "max_len": MAX_LEN,
        "dim": 1536,
    })


@app.route("/embed", methods=["POST"])
def embed():
    data = request.json
    texts = data.get("texts", [])
    normalize = data.get("normalize", True)
    
    if not texts:
        return jsonify({"error": "No texts provided"}), 400
    if len(texts) > 100:
        return jsonify({"error": "Max 100 texts per request"}), 400
    
    t0 = time.time()
    embeddings = encode_texts(texts, normalize=normalize)
    elapsed = time.time() - t0
    
    return jsonify({
        "embeddings": embeddings.tolist(),
        "dim": embeddings.shape[1],
        "count": len(texts),
        "time_ms": round(elapsed * 1000, 1),
    })


@app.route("/similarity", methods=["POST"])
def similarity():
    data = request.json
    query = data.get("query", "")
    documents = data.get("documents", [])
    
    if not query or not documents:
        return jsonify({"error": "Need query and documents"}), 400
    
    t0 = time.time()
    all_texts = [query] + documents
    embeddings = encode_texts(all_texts, normalize=True)
    
    query_emb = embeddings[0:1]
    doc_embs = embeddings[1:]
    scores = (query_emb @ doc_embs.T).squeeze(0).tolist()
    elapsed = time.time() - t0
    
    return jsonify({
        "scores": scores,
        "time_ms": round(elapsed * 1000, 1),
    })


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", default="checkpoints/deploy/deepx_int8.pt")
    parser.add_argument("--tokenizer", default="/mnt/d/DX/hf_cache/hub/models--google--gemma-4-E2B-it/snapshots/6b7e72c67d3c4556f42b56d5a68b4b8e864c63b4")
    parser.add_argument("--port", type=int, default=8080)
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--device", default="cuda")
    args = parser.parse_args()
    
    load_model(args.checkpoint, args.tokenizer, args.device)
    
    logger.info(f"Starting server on {args.host}:{args.port}")
    app.run(host=args.host, port=args.port, debug=False)


if __name__ == "__main__":
    main()