import gradio as gr
import onnx
from huggingface_hub import HfApi
import json
import sys
import os
import io
import requests
from urllib.parse import urlparse


HfClient = HfApi();

ONNX_PREFERED = [
    "model.onnx",
    "onnx/model.onnx"    
]

ONNX_CACHE = {
    'models': {
        
    }
}

def is_url(path):
    try:
        result = urlparse(path)
        return all([result.scheme, result.netloc])
    except ValueError:
        return False


def load_model(path):
    if is_url(path):
        print(f"Downloading model from: {path}...", file=sys.stderr)
        try:
            response = requests.get(path)
            response.raise_for_status() # Check for HTTP errors
            # Load from binary stream
            return onnx.load(io.BytesIO(response.content))
        except requests.exceptions.RequestException as e:
            print(f"Error downloading model: {e}", file=sys.stderr)
            sys.exit(1)
    else:
        # Check if local file exists before loading
        if not os.path.exists(path):
            print(f"Error: File not found at {path}", file=sys.stderr)
            sys.exit(1)
        
        return onnx.load(path)

def CheckSqlOnnx(path):
    
    OnnxModel = load_model(path);
    
    initializer_names = {init.name for init in OnnxModel.graph.initializer}
    inputs = [inp.name for inp in OnnxModel.graph.input if inp.name not in initializer_names]
    outputs = [out.name for out in OnnxModel.graph.output]

    required_inputs = {"input_ids", "attention_mask"}
    required_outputs = {"token_embeddings", "sentence_embedding"}

    is_supported = (
        required_inputs.issubset(inputs) and
        required_outputs.issubset(outputs)
    )

    OnnxInouts = {
        "supported": is_supported,
        "inputs": inputs,
        "outputs": outputs
    }
    
    return OnnxInouts;

def CheckModel(repo_id: str, path: str | None = None):
    
    MODELS_CACHE = ONNX_CACHE['models'];
    
    CacheSlot = MODELS_CACHE.get(repo_id);
    
    if CacheSlot:
        return json.dumps(CacheSlot, indent = 2);
    
    
    
    model_info = HfClient.model_info(repo_id=repo_id)

    # Extract filenames from RepoSibling objects
    sibling_files = [s.rfilename for s in model_info.siblings]

    onnx_path = None

    if path:
        if path in sibling_files:
            onnx_path = path
        else:
            return f"Error: ONNX file not found: {path}"
    else:
        for p in ONNX_PREFERED:
            if p in sibling_files:
                onnx_path = p
                break

        if not onnx_path:
            onnx_path = next(
                (f for f in sibling_files if f.lower().endswith(".onnx")),
                None
            )

    if not onnx_path:
        raise "Error: No ONNX model found in repository";

    # Build Hugging Face raw file URL
    file_url = f"https://huggingface.co/{repo_id}/resolve/main/{onnx_path}"

    # Check SQL ONNX compatibility
    OnnxInfo = CheckSqlOnnx(file_url)
    
    CacheSlot = {
         'url': file_url
        ,'onnx': OnnxInfo
    }
    
    MODELS_CACHE[repo_id] = CacheSlot

    return json.dumps({**CacheSlot, 'cached': False}, indent = 2);


with gr.Blocks() as demo:
    
    gr.Markdown("""
        This sample app tests whether a given model repository can be used with SQL Server ONNX.  
        
        In some tests, it was discovered that not every ONNX model works with SQL Server `CREATE EXTERNAL MODEL`.  
        For it to work, the input parameters of the neural network must contain specific names, and the output as well.  
        
        I don’t know if this behavior will be kept in future versions of SQL Server 2025.  
        
        However, while we don’t have official documentation about this, this repository can help discover whether a given model will work with SQL Server if you plan to download and use it with ONNX.  
        
        Just input the model name below in the format `user/model-name` (check the examples).  
        
        Look at the JSON output. If `"supported"` is `true`, then you can use it with SQL Server.  
        
        Soon, I will bring a default tested list!  
        
        **IMPORTANT**: To check this, the space will attempt to download the model ONNX file. If it is large, it can take several minutes.  

    """)
    
    ModelPath = gr.Textbox(label="Model Repository", submit_btn = True);
    ModelInfoOut = gr.Textbox(label="Model Info", lines = 10)
    
    ModelPath.submit(fn=CheckModel, inputs=ModelPath, outputs=ModelInfoOut)
    
    gr.Examples([
         ["intfloat/multilingual-e5-large"]
        ,["mixedbread-ai/mxbai-embed-xsmall-v1"]
        ,["nsense/all-MiniLM-L6-v2-onnx"]
    
    ], ModelPath)

demo.launch(
    server_name = '0.0.0.0'
)