import gradio as gr
import torch
from sentence_transformers import SentenceTransformer
from torch.nn.functional import cosine_similarity
from spaces import GPU

# モデルの読み込み
model = SentenceTransformer("Shuu12121/CodeCloneDetection-ModernBERT-Owl")
model.eval()

# 閾値設定（安定性の高い0.9推奨）
THRESHOLD = 0.9

@GPU
def detect_clone(code1, code2):
    if not code1.strip() or not code2.strip():
        return "❌ どちらのコードも入力してください", ""

    with torch.no_grad():
        embeddings = model.encode([code1, code2], convert_to_tensor=True)
        sim_score = cosine_similarity(embeddings[0].unsqueeze(0), embeddings[1].unsqueeze(0)).item()

    result = (
        f"🟢 類似度: {sim_score:.4f}\n→ これらのコードは **クローン** と判定されます。"
        if sim_score >= THRESHOLD
        else f"🔴 類似度: {sim_score:.4f}\n→ これらのコードは **クローンではありません**。"
    )
    return result, sim_score

# Javaクローン例（ファイル読み込み）
java_code1 = """\
public static String readFileToString(File file, Charset encoding) throws IOException {
    try (BufferedReader reader = new BufferedReader(new InputStreamReader(new FileInputStream(file), encoding))) {
        StringBuilder sb = new StringBuilder();
        String line;
        while ((line = reader.readLine()) != null) {
            sb.append(line).append("\\n");
        }
        return sb.toString();
    }
}"""

java_code2 = """\
public static String readFileToString(File file, Charset encoding) throws IOException {
    byte[] bytes = java.nio.file.Files.readAllBytes(file.toPath());
    return new String(bytes, encoding);
}"""

# Gradioインターフェースの作成
demo = gr.Interface(
    fn=detect_clone,
    inputs=[
        gr.Textbox(label="コードスニペット1", lines=12, value=java_code1),
        gr.Textbox(label="コードスニペット2", lines=8, value=java_code2),
    ],
    outputs=[
        gr.Markdown(label="判定結果"),
        gr.Number(label="Cosine Similarity")
    ],
    title="Code Clone Detection with ModernBERT-Owl 🦉",
    description="Shuu12121/CodeModernBERT-Owl によって構築された Sentence-BERT モデルを使用し、コードクローンを検出します。これは Java における意味的に同等なファイル読み込み関数の例です。"
)

if __name__ == "__main__":
    demo.launch()