#!/usr/bin/env python3
"""
app.py — Production Serving Harness & Playground for JMullings/Spectral_Embeddings
=================================================================================
• Auto-downloads and unpacks EMBEDDINGS_CORE.npz from the Hugging Face Hub.
• Multiscale 16,384-D feature hashing with calibrated Riemannian metric decoding.
• Exposes Ollama-compatible REST endpoints (/api/embed, /api/tags, /api/metric).
• Binds to 0.0.0.0:7860 (Hugging Face Spaces default) or environment $PORT.
=================================================================================
"""

import os
import sys
import json
import time
import math
import hashlib
import re
import threading
from pathlib import Path
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse

import numpy as np

# =====================================================================
# 1. Configuration & Weights Acquisition
# =====================================================================

REPO_ID = "JMullings/Spectral_Embeddings"
FILENAME = "EMBEDDINGS_CORE.npz"
PORT = int(os.environ.get("PORT", 7860))
HOST = "0.0.0.0"
BASE_DIR = Path(__file__).resolve().parent
LOCAL_WEIGHTS = BASE_DIR / FILENAME

if not LOCAL_WEIGHTS.exists():
    print(f"[*] Downloading {FILENAME} from huggingface.co/{REPO_ID}...")
    try:
        from huggingface_hub import hf_hub_download
        downloaded = hf_hub_download(
            repo_id=REPO_ID,
            filename=FILENAME,
            local_dir=str(BASE_DIR),
            token=os.environ.get("HF_TOKEN")
        )
        LOCAL_WEIGHTS = Path(downloaded)
        print("[+] Download complete.")
    except Exception as e:
        print(f"[!] Warning: Could not download from Hub: {e}. Checking fallback.")

if not LOCAL_WEIGHTS.exists():
    raise FileNotFoundError(f"Missing model file {LOCAL_WEIGHTS}. Ensure {FILENAME} exists.")


# =====================================================================
# 2. Model Unpacking & Multiscale Feature Encoder
# =====================================================================

print(f"[*] Loading {LOCAL_WEIGHTS.name}...")
data = np.load(LOCAL_WEIGHTS, allow_pickle=True)
keys = list(data.files)
print(f"[+] Archive keys: {keys}")

# Extract W matrix (reconstructing sparse col_idx if present)
if "W_q" in keys:
    W_raw = np.asarray(data["W_q"], dtype=np.float32)
    if "col_idx" in keys and "H" in keys:
        H = int(data["H"])
        k = int(data["k"]) if "k" in keys else W_raw.shape[0]
        col_idx = np.asarray(data["col_idx"], dtype=np.int64)
        W = np.zeros((k, H), dtype=np.float32)
        if "W_scale" in keys:
            W[:, col_idx] = W_raw * np.asarray(data["W_scale"], dtype=np.float32)[:, None]
        else:
            W[:, col_idx] = W_raw
    else:
        W = W_raw
elif "W" in keys:
    W = np.asarray(data["W"], dtype=np.float32)
elif "W0" in keys:
    W = np.asarray(data["W0"], dtype=np.float32)
else:
    matrix_keys = [k for k in keys if len(data[k].shape) == 2]
    W = np.asarray(data[matrix_keys[0]], dtype=np.float32)

b = np.asarray(data["b"] if "b" in keys else data.get("b0", np.zeros(W.shape[0], dtype=np.float32)), dtype=np.float32)

embed_dim, feature_dim = W.shape
hash_seed = int(data["hash_seed"]) if "hash_seed" in keys else 42

print(f"[+] Projection Matrix: ({embed_dim} x {feature_dim}) = {W.size:,} parameters")
print(f"[+] Hash Seed: {hash_seed}")

# ── 16,384-D Multiscale Feature Extractor ──
class TextEncoder:
    def __init__(self, H: int = 16384, seed: int = 42):
        self.H = H
        self.seed = seed
        self._cache = {}

    def _hash_token(self, tok: str) -> int:
        idx = self._cache.get(tok)
        if idx is None:
            h = hashlib.blake2b(f"{self.seed}:{tok}".encode("utf-8"), digest_size=8).digest()
            idx = int.from_bytes(h, "little") % self.H
            if len(self._cache) > 200_000:
                self._cache.clear()
            self._cache[tok] = idx
        return idx

    def transform(self, text: str) -> np.ndarray:
        x = np.zeros(self.H, dtype=np.float32)
        cleaned = text.lower().strip()
        words = re.findall(r"[a-z0-9]+", cleaned)
        for w in words:
            x[self._hash_token(f"w:{w}")] += 2.0
        for j in range(len(words) - 1):
            x[self._hash_token(f"bi:{words[j]}_{words[j+1]}")] += 3.0
        padded = f" {cleaned} "
        L = len(padded)
        for n in (3, 4):
            for j in range(L - n + 1):
                x[self._hash_token(f"ch:{padded[j:j+n]}")] += 1.0
        norm = float(np.linalg.norm(x))
        if norm > 1e-12:
            x /= norm
        return x

encoder = TextEncoder(H=feature_dim, seed=hash_seed)

# ── Calibrated Riemannian Metric (M_diag) ──
def build_unified_metric(dim: int) -> np.ndarray:
    gamma = 1.50
    omega_star = 0.78
    raw_m = np.zeros(dim, dtype=np.float64)
    for k in range(dim):
        gamma_k = 0.5 + 0.35 * (k + 1)
        center_k = 0.10 + 0.15 * k
        sech2 = 1.0 - np.tanh(gamma_k * (omega_star - center_k)) ** 2
        raw_m[k] = gamma_k * sech2
    if raw_m.max() > 1e-12:
        raw_m /= raw_m.max()
    m_diag = (1.0 + gamma * raw_m) * (dim / max(float((1.0 + gamma * raw_m).sum()), 1e-12))
    return m_diag.astype(np.float32)

M_diag = build_unified_metric(embed_dim)

def embed(text: str) -> np.ndarray:
    x = encoder.transform(text)
    z = W @ x + b
    norm = float(np.linalg.norm(z))
    return (z / max(norm, 1e-12)).astype(np.float32)

def unified_similarity(a: np.ndarray, b_vec: np.ndarray) -> float:
    diff = a - b_vec
    dist = np.sqrt(np.sum(M_diag * (diff ** 2)))
    return float(np.exp(-dist))


# =====================================================================
# 3. Startup Self-Test Verification
# =====================================================================

print("\n--- Running Startup Verification Probe ---")
v1 = embed("How do I reset my password?")
v2 = embed("Steps to recover your account credentials.")
v3 = embed("The weather forecast calls for rain.")

raw_cos_p = float(np.dot(v1, v2))
raw_cos_c = float(np.dot(v1, v3))
cal_sim_p = unified_similarity(v1, v2)
cal_sim_c = unified_similarity(v1, v3)

print(f"• Raw Cosine        : Paraphrase = {raw_cos_p:.4f} | Contradiction = {raw_cos_c:.4f} (Margin: {raw_cos_p - raw_cos_c:+.4f})")
print(f"• Calibrated Metric : Paraphrase = {cal_sim_p:.4f} | Contradiction = {cal_sim_c:.4f} (Margin: {cal_sim_p - cal_sim_c:+.4f})")
if cal_sim_p > cal_sim_c:
    print(">>> [PASS] Semantic Discrimination Probe Verified.")
else:
    print(">>> [WARN] Discrimination margin is zero or negative.")
print("------------------------------------------\n")


# =====================================================================
# 4. Multi-Threaded HTTP API Server (Ollama Drop-In)
# =====================================================================

class Handler(BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):
        # Silence routine health checks in stdout
        if args and isinstance(args[0], str) and "/api/health" in args[0]:
            return
        sys.stderr.write(f"[{self.log_date_time_string()}] {self.address_string()} {fmt % args}\n")

    def _send(self, code: int, body: bytes, ctype="application/json"):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "*")
        self.end_headers()
        try:
            self.wfile.write(body)
        except (BrokenPipeError, ConnectionResetError):
            pass

    def do_OPTIONS(self):
        self._send(204, b"")

    def do_GET(self):
        p = urlparse(self.path).path
        if p in ("/", "/index.html"):
            html_file = BASE_DIR / "index.html"
            content = html_file.read_bytes() if html_file.exists() else b"<h1>Spectral Embeddings Server Active</h1>"
            self._send(200, content, "text/html; charset=utf-8")
            return

        if p == "/api/health":
            self._send(200, json.dumps({
                "status": "ok",
                "model": REPO_ID,
                "dim": embed_dim,
                "feature_dim": feature_dim,
                "parameters": int(W.size),
                "metric": "unified_nn",
            }).encode())
            return

        if p == "/api/metric":
            self._send(200, json.dumps({
                "name": "unified_nn",
                "dim": embed_dim,
                "M_diag": M_diag.tolist(),
                "equation": "exp(-sqrt(sum_k M_kk (a_k-b_k)^2))"
            }).encode())
            return

        if p == "/api/tags":
            self._send(200, json.dumps({
                "models": [{
                    "name": "spectral-embed-v2-online",
                    "model": "spectral-embed-v2-online",
                    "size": int(W.size * 4),
                    "details": {
                        "family": "spectral",
                        "dimension": embed_dim,
                        "parameters": int(W.size)
                    }
                }]
            }).encode())
            return

        self._send(404, b'{"error":"endpoint not found"}')

    def do_POST(self):
        p = urlparse(self.path).path
        length = int(self.headers.get("Content-Length", 0))
        raw_body = self.rfile.read(length) if length > 0 else b"{}"

        try:
            body = json.loads(raw_body.decode("utf-8")) if raw_body else {}
        except Exception:
            self._send(400, b'{"error":"invalid json body"}')
            return

        if p in ("/api/embed", "/api/embeddings"):
            inp = body.get("input", body.get("prompt", ""))
            texts = [inp] if isinstance(inp, str) else [str(x) for x in inp]
            
            t0 = time.perf_counter()
            embs = [embed(t).tolist() for t in texts]
            ms = (time.perf_counter() - t0) * 1000.0

            self._send(200, json.dumps({
                "embeddings": embs,
                "dim": embed_dim,
                "model": "spectral-embed-v2-online",
                "latency_ms": round(ms, 2),
                "per_sentence_ms": round(ms / max(len(texts), 1), 3),
            }).encode())
            return

        self._send(404, b'{"error":"endpoint not found"}')


def main():
    print(f"[*] Starting Multi-Threaded Spectral Harness on http://{HOST}:{PORT}")
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n[!] Shutting down server gracefully...")
        server.server_close()


if __name__ == "__main__":
    main()