#!/usr/bin/env python3
import json
from pathlib import Path
from typing import Any, Dict, List

import faiss
import gradio as gr
import numpy as np
import torch
import torch.nn as nn

MODEL_PATH = Path("vil-encoder-v2.pt")
DATA_PATHS = [
    Path("data/train.jsonl"),
    Path("data/validation.jsonl"),
    Path("data/test.jsonl"),
]
DEVICE = "cpu"
SEQ_LEN = 64
EMBED_DIM = 32

def encode_triplet(visible: str, braille: str, hanzi: str) -> np.ndarray:
    text = f"{visible}|{braille}|{hanzi}"
    arr = np.array([ord(c) % 256 for c in text], dtype=np.float32)
    if arr.shape[0] < SEQ_LEN:
        arr = np.pad(arr, (0, SEQ_LEN - arr.shape[0]))
    else:
        arr = arr[:SEQ_LEN]
    arr /= 255.0
    return arr.astype(np.float32)

class Encoder(nn.Module):
    def __init__(self, input_dim: int = SEQ_LEN, embed_dim: int = EMBED_DIM) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, embed_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = self.net(x)
        return nn.functional.normalize(z, dim=-1)

def load_dataset() -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    for p in DATA_PATHS:
        if p.exists():
            with p.open("r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        rows.append(json.loads(line))
    return rows

def load_model() -> tuple[Encoder, Dict[str, Any]]:
    model = Encoder()
    status: Dict[str, Any] = {
        "loaded": False,
        "model_path": str(MODEL_PATH),
        "error": None,
    }

    if not MODEL_PATH.exists():
        status["error"] = f"missing model: {MODEL_PATH}"
        return model.eval(), status

    try:
        obj = torch.load(MODEL_PATH, map_location=DEVICE)
        if isinstance(obj, dict) and "model_state_dict" in obj:
            model.load_state_dict(obj["model_state_dict"], strict=True)
        elif isinstance(obj, dict):
            model.load_state_dict(obj, strict=False)
        else:
            raise RuntimeError(f"unsupported checkpoint type: {type(obj).__name__}")
        model.eval()
        status["loaded"] = True
        return model, status
    except Exception as e:
        status["error"] = str(e)
        model.eval()
        return model, status

DATASET = load_dataset()
MODEL, MODEL_STATUS = load_model()

INDEX = faiss.IndexFlatL2(EMBED_DIM)
EMBED_MATRIX = None

def model_embed(v: str, b: str, h: str) -> np.ndarray:
    vec = encode_triplet(v, b, h)
    x = torch.from_numpy(vec).unsqueeze(0)
    with torch.no_grad():
        z = MODEL(x).cpu().numpy()[0]
    return z.astype(np.float32)

def build_index() -> None:
    global EMBED_MATRIX
    if not DATASET or not MODEL_STATUS["loaded"]:
        EMBED_MATRIX = np.zeros((0, EMBED_DIM), dtype=np.float32)
        return
    vectors = []
    for row in DATASET:
        vectors.append(model_embed(row["visible"], row["braille"], row["hanzi"]))
    EMBED_MATRIX = np.stack(vectors).astype(np.float32)
    INDEX.add(EMBED_MATRIX)

build_index()

def render_sigil(v: str, b: str, h: str) -> str:
    glyphstring = f"{v}{b}{h}"
    locked = f"⊏⚙{glyphstring}⚙⊐"
    svg = f"""
    <svg width="320" height="200" xmlns="http://www.w3.org/2000/svg">
      <rect width="100%" height="100%" fill="black"/>
      <text x="50%" y="50%" dominant-baseline="middle" text-anchor="middle"
            fill="white" font-size="36" font-family="monospace">{locked}</text>
    </svg>
    """
    return svg

def nearest(v: str, b: str, h: str, k: int = 5) -> List[Dict[str, Any]]:
    if not DATASET or not MODEL_STATUS["loaded"] or INDEX.ntotal == 0:
        return []
    q = model_embed(v, b, h).reshape(1, -1)
    distances, indices = INDEX.search(q, k)
    out: List[Dict[str, Any]] = []
    for dist, idx in zip(distances[0].tolist(), indices[0].tolist()):
        if idx < 0 or idx >= len(DATASET):
            continue
        row = dict(DATASET[idx])
        row["_distance"] = float(dist)
        out.append(row)
    return out

def run_pipeline(visible: str, braille: str, hanzi: str):
    visible = (visible or "").strip()
    braille = (braille or "").strip()
    hanzi = (hanzi or "").strip()

    if not visible or not braille or not hanzi:
        return {"error": "Provide visible, braille, and hanzi."}, ""

    if not MODEL_STATUS["loaded"]:
        return {"error": "Model not loaded.", "model_status": MODEL_STATUS}, ""

    embedding = model_embed(visible, braille, hanzi).tolist()
    matches = nearest(visible, braille, hanzi, k=5)
    svg = render_sigil(visible, braille, hanzi)

    result = {
        "input": {
            "visible": visible,
            "braille": braille,
            "hanzi": hanzi,
        },
        "embedding": embedding,
        "nearest": matches,
        "glyphstring": f"{visible}{braille}{hanzi}",
        "sigil": f"⊏⚙{visible}{braille}{hanzi}⚙⊐",
    }
    return result, svg

def search_visible(query: str):
    query = (query or "").strip()
    if not query:
        return []
    return [row for row in DATASET if query in str(row.get("visible", ""))][:10]

with gr.Blocks(title="VIL Encoder — Glyphmatic Inference Engine") as demo:
    gr.Markdown("# VIL Encoder — Glyphmatic Inference Engine")

    with gr.Tab("Encode"):
        visible = gr.Textbox(label="Visible Canon", placeholder="✶")
        braille = gr.Textbox(label="Invisible Braille", placeholder="⠁")
        hanzi = gr.Textbox(label="Hanzi Context", placeholder="一")
        run_btn = gr.Button("Run")
        result_json = gr.JSON()
        sigil_svg = gr.HTML()

        run_btn.click(
            fn=run_pipeline,
            inputs=[visible, braille, hanzi],
            outputs=[result_json, sigil_svg],
        )

    with gr.Tab("Search Dataset"):
        query = gr.Textbox(label="Query Visible", placeholder="✶")
        query_btn = gr.Button("Search")
        query_out = gr.JSON()
        query_btn.click(fn=search_visible, inputs=[query], outputs=[query_out])

    with gr.Tab("System Info"):
        gr.JSON(
            {
                "device": DEVICE,
                "model_status": MODEL_STATUS,
                "dataset_rows": len(DATASET),
                "index_size": int(INDEX.ntotal),
            }
        )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)