from flask import Flask, request, jsonify, render_template_string
from sentence_transformers import SentenceTransformer, util
import logging
import sys
import signal

# 初始化 Flask 应用
app = Flask(__name__)

# 配置日志，级别设为 INFO
logging.basicConfig(level=logging.INFO)
app.logger = logging.getLogger("CodeSearchAPI")

# 预定义代码片段
CODE_SNIPPETS = [
    "System.out.println(\"Hello, World!\");",
"public static int sum(int a, int b) { return a + b; }",
"import java.util.Random; public static int generateRandomNumber() { return new Random().nextInt(); }",
"public static boolean isEven(int number) { return number % 2 == 0; }",
"public static int stringLength(String str) { return str.length(); }",
"import java.time.LocalDate; public static LocalDate getCurrentDate() { return LocalDate.now(); }",
"import java.io.File; public static boolean fileExists(String path) { return new File(path).exists(); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static String readFileContent(String path) throws Exception { return new String(Files.readAllBytes(Paths.get(path))); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static void writeToFile(String path, String content) throws Exception { Files.write(Paths.get(path), content.getBytes()); }",
"import java.time.LocalTime; public static LocalTime getCurrentTime() { return LocalTime.now(); }",
"public static String toUpperCase(String str) { return str.toUpperCase(); }",
"public static String toLowerCase(String str) { return str.toLowerCase(); }",
"public static String reverseString(String str) { return new StringBuilder(str).reverse().toString(); }",
"public static int countListElements(List<?> list) { return list.size(); }",
"public static int findMax(List<Integer> list) { return Collections.max(list); }",
"public static int findMin(List<Integer> list) { return Collections.min(list); }",
"public static List<Integer> sortList(List<Integer> list) { Collections.sort(list); return list; }",
"public static List<?> mergeLists(List<?> list1, List<?> list2) { List<Object> mergedList = new ArrayList<>(list1); mergedList.addAll(list2); return mergedList; }",
"public static void removeElement(List<?> list, Object element) { list.remove(element); }",
"public static boolean isListEmpty(List<?> list) { return list.isEmpty(); }",
"public static int countCharOccurrences(String str, char ch) { return (int) str.chars().filter(c -> c == ch).count(); }",
"public static boolean containsSubstring(String str, String sub) { return str.contains(sub); }",
"public static String numberToString(int number) { return Integer.toString(number); }",
"public static int stringToNumber(String str) { return Integer.parseInt(str); }",
"public static boolean isNumeric(String str) { return str.matches(\"\\\\d+\"); }",
"public static int getElementIndex(List<?> list, Object element) { return list.indexOf(element); }",
"public static void clearList(List<?> list) { list.clear(); }",
"public static void reverseList(List<?> list) { Collections.reverse(list); }",
"public static List<?> removeDuplicates(List<?> list) { return new ArrayList<>(new HashSet<>(list)); }",
"public static boolean isInList(List<?> list, Object element) { return list.contains(element); }",
"public static Map<Object, Object> createDictionary() { return new HashMap<>(); }",
"public static void addToDictionary(Map<Object, Object> dict, Object key, Object value) { dict.put(key, value); }",
"public static void removeFromDictionary(Map<Object, Object> dict, Object key) { dict.remove(key); }",
"public static Set<Object> getDictionaryKeys(Map<Object, Object> dict) { return dict.keySet(); }",
"public static Collection<Object> getDictionaryValues(Map<Object, Object> dict) { return dict.values(); }",
"public static Map<Object, Object> mergeDictionaries(Map<Object, Object> dict1, Map<Object, Object> dict2) { Map<Object, Object> mergedDict = new HashMap<>(dict1); mergedDict.putAll(dict2); return mergedDict; }",
"public static boolean isDictionaryEmpty(Map<Object, Object> dict) { return dict.isEmpty(); }",
"public static Object getDictionaryValue(Map<Object, Object> dict, Object key) { return dict.get(key); }",
"public static boolean keyExistsInDictionary(Map<Object, Object> dict, Object key) { return dict.containsKey(key); }",
"public static void clearDictionary(Map<Object, Object> dict) { dict.clear(); }",
"import java.io.BufferedReader; import java.io.FileReader; public static int countFileLines(String path) throws Exception { return (int) new BufferedReader(new FileReader(path)).lines().count(); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static void writeListToFile(String path, List<?> list) throws Exception { Files.write(Paths.get(path), list.toString().getBytes()); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static List<String> readListFromFile(String path) throws Exception { return Files.readAllLines(Paths.get(path)); }",
"import java.io.BufferedReader; import java.io.FileReader; public static int countFileWords(String path) throws Exception { return new BufferedReader(new FileReader(path)).lines().mapToInt(line -> line.split(\"\\\\s+\").length).sum(); }",
"public static boolean isLeapYear(int year) { return (year % 4 == 0 && year % 100 != 0) || (year % 400 == 0); }",
"import java.time.format.DateTimeFormatter; public static String formatTime(LocalTime time, String pattern) { return time.format(DateTimeFormatter.ofPattern(pattern)); }",
"import java.time.LocalDate; import java.time.temporal.ChronoUnit; public static long daysBetween(LocalDate date1, LocalDate date2) { return ChronoUnit.DAYS.between(date1, date2); }",
"import java.nio.file.Paths; public static String getCurrentWorkingDirectory() { return Paths.get(\"\").toAbsolutePath().toString(); }",
"import java.io.File; public static List<String> listFilesInDirectory(String path) { return Arrays.asList(new File(path).list()); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static void createDirectory(String path) throws Exception { Files.createDirectory(Paths.get(path)); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static void deleteDirectory(String path) throws Exception { Files.delete(Paths.get(path)); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static boolean isFile(String path) { return Files.isRegularFile(Paths.get(path)); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static boolean isDirectory(String path) { return Files.isDirectory(Paths.get(path)); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static long getFileSize(String path) throws Exception { return Files.size(Paths.get(path)); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static void renameFile(String oldPath, String newPath) throws Exception { Files.move(Paths.get(oldPath), Paths.get(newPath)); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static void copyFile(String sourcePath, String destinationPath) throws Exception { Files.copy(Paths.get(sourcePath), Paths.get(destinationPath)); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static void moveFile(String sourcePath, String destinationPath) throws Exception { Files.move(Paths.get(sourcePath), Paths.get(destinationPath)); }",
"import java.nio.file.Files; import java.nio.file.Paths; public static void deleteFile(String path) throws Exception { Files.delete(Paths.get(path)); }",
"public static String getEnvVariable(String name) { return System.getenv(name); }",
"public static void setEnvVariable(String name, String value) { System.setProperty(name, value); }",
"import java.awt.Desktop; import java.net.URI; public static void openWebLink(String url) throws Exception { Desktop.getDesktop().browse(new URI(url)); }",
"import java.net.HttpURLConnection; import java.net.URL; public static String sendGetRequest(String url) throws Exception { HttpURLConnection connection = (HttpURLConnection) new URL(url).openConnection(); connection.setRequestMethod(\"GET\"); return new String(connection.getInputStream().readAllBytes()); }",
"import com.google.gson.JsonParser; public static Object parseJson(String json) { return JsonParser.parseString(json); }",
"import com.google.gson.Gson; import java.nio.file.Files; import java.nio.file.Paths; public static void writeJsonToFile(String path, Object obj) throws Exception { Files.write(Paths.get(path), new Gson().toJson(obj).getBytes()); }",
"import com.google.gson.Gson; import java.nio.file.Files; import java.nio.file.Paths; public static Object readJsonFromFile(String path) throws Exception { return new Gson().fromJson(new String(Files.readAllBytes(Paths.get(path))), Object.class); }",
"public static String listToString(List<?> list) { return list.toString(); }",
"public static List<String> stringToList(String str) { return Arrays.asList(str.split(\",\")); }",
"public static String joinWithComma(List<?> list) { return String.join(\",\", list.stream().map(Object::toString).toArray(String[]::new)); }",
"public static String joinWithNewline(List<?> list) { return String.join(\"\\n\", list.stream().map(Object::toString).toArray(String[]::new)); }",
"public static String[] splitBySpace(String str) { return str.split(\"\\\\s+\"); }",
"public static String[] splitByDelimiter(String str, String delimiter) { return str.split(delimiter); }",
"public static String[] splitIntoChars(String str) { return str.split(\"\"); }",
"public static String replaceInString(String str, String target, String replacement) { return str.replace(target, replacement); }",
"public static String removeSpaces(String str) { return str.replaceAll(\"\\\\s\", \"\"); }",
"public static String removePunctuation(String str) { return str.replaceAll(\"[^a-zA-Z0-9]\", \"\"); }",
"public static boolean isStringEmpty(String str) { return str.isEmpty(); }",
"public static boolean isPalindrome(String str) { return str.equals(new StringBuilder(str).reverse().toString()); }",
"import com.opencsv.CSVWriter; import java.io.FileWriter; public static void writeToCsv(String path, List<String[]> data) throws Exception { CSVWriter writer = new CSVWriter(new FileWriter(path)); writer.writeAll(data); writer.close(); }",
"import com.opencsv.CSVReader; import java.io.FileReader; public static List<String[]> readFromCsv(String path) throws Exception { CSVReader reader = new CSVReader(new FileReader(path)); return reader.readAll(); }",
"import com.opencsv.CSVReader; import java.io.FileReader; public static int countCsvLines(String path) throws Exception { return readFromCsv(path).size(); }",
"public static void shuffleList(List<?> list) { Collections.shuffle(list); }",
"public static Object getRandomElement(List<?> list) { return list.get(new Random().nextInt(list.size())); }",
"public static List<?> getRandomElements(List<?> list, int count) { Collections.shuffle(list); return list.subList(0, count); }",
"public static int rollDice() { return new Random().nextInt(6) + 1; }",
"public static String flipCoin() { return new Random().nextBoolean() ? \"Heads\" : \"Tails\"; }",
"import java.util.Random; public static String generateRandomPassword(int length) { String chars = \"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789\"; StringBuilder password = new StringBuilder(); for (int i = 0; i < length; i++) { password.append(chars.charAt(new Random().nextInt(chars.length()))); } return password.toString(); }",
"import java.util.Random; public static String generateRandomColor() { Random random = new Random(); return String.format(\"#%06x\", random.nextInt(0xFFFFFF + 1)); }",
"import java.util.UUID; public static String generateUniqueId() { return UUID.randomUUID().toString(); }",
"public class MyClass {}",
"MyClass myObject = new MyClass();",
"public class MyClass { public void myMethod() {} }",
"public class MyClass { public String myAttribute; }",
"public class ChildClass extends MyClass {}",
"public class ChildClass extends MyClass { @Override public void myMethod() {} }",
"public class MyClass { public static void myClassMethod() {} }",
"public class MyClass { public static void myStaticMethod() {} }",
"public static boolean isInstanceOf(Object obj, Class<?> clazz) { return clazz.isInstance(obj); }",
"public static Object getAttribute(Object obj, String attribute) throws Exception { return obj.getClass().getDeclaredField(attribute).get(obj); }",
"public static void setAttribute(Object obj, String attribute, Object value) throws Exception { obj.getClass().getDeclaredField(attribute).set(obj, value); }",
"public static void deleteAttribute(Object obj, String attribute) throws Exception { obj.getClass().getDeclaredField(attribute).set(obj, null); }"
"try { int result = 10 / 0; } catch (Exception e) { System.out.println(\"Exception caught\"); }",
"class CustomException extends Exception { public CustomException(String message) { super(message); } }",
"try { int result = 10 / 0; } catch (Exception e) { System.out.println(e.getMessage()); }",
"import java.util.logging.Logger; Logger logger = Logger.getLogger(MyClass.class.getName()); logger.severe(\"Error occurred\");",
"long startTime = System.currentTimeMillis();",
"long endTime = System.currentTimeMillis(); long elapsedTime = endTime - startTime;",
"for (int i = 0; i <= 100; i++) { System.out.print(\"\\rProgress: \" + i + \"%\"); Thread.sleep(50); }",
"Thread.sleep(1000);",
"Runnable r = () -> System.out.println(\"Lambda expression\"); r.run();",
"List<Integer> numbers = Arrays.asList(1, 2, 3); List<Integer> squares = numbers.stream().map(x -> x * x).collect(Collectors.toList());",
"List<Integer> numbers = Arrays.asList(1, 2, 3); List<Integer> evens = numbers.stream().filter(x -> x % 2 == 0).collect(Collectors.toList());",
"import java.util.stream.Collectors; List<Integer> numbers = Arrays.asList(1, 2, 3); int sum = numbers.stream().reduce(0, Integer::sum);",
"List<Integer> squares = IntStream.range(0, 10).map(x -> x * x).boxed().collect(Collectors.toList());",
"Map<String, Integer> map = IntStream.range(0, 10).boxed().collect(Collectors.toMap(i -> \"key\" + i, i -> i));",
"Set<Integer> set = IntStream.range(0, 10).boxed().collect(Collectors.toSet());",
"Set<Integer> set1 = new HashSet<>(Arrays.asList(1, 2, 3)); Set<Integer> set2 = new HashSet<>(Arrays.asList(2, 3, 4)); set1.retainAll(set2);",
"Set<Integer> set1 = new HashSet<>(Arrays.asList(1, 2, 3)); Set<Integer> set2 = new HashSet<>(Arrays.asList(2, 3, 4)); set1.addAll(set2);",
"Set<Integer> set1 = new HashSet<>(Arrays.asList(1, 2, 3)); Set<Integer> set2 = new HashSet<>(Arrays.asList(2, 3, 4)); set1.removeAll(set2);",
"List<Integer> list = Arrays.asList(1, null, 2, null); list.removeIf(Objects::isNull);",
"try (BufferedReader br = new BufferedReader(new FileReader(\"file.txt\"))) { br.readLine(); } catch (IOException e) { System.out.println(\"File cannot be opened\"); }",
"if (variable instanceof Integer) { System.out.println(\"Variable is an Integer\"); }",
"boolean bool = Boolean.parseBoolean(\"true\");",
"if (condition) { System.out.println(\"Condition is true\"); }",
"while (condition) { System.out.println(\"Looping\"); }",
"for (int i : list) { System.out.println(i); }",
"for (Map.Entry<String, Integer> entry : map.entrySet()) { System.out.println(entry.getKey() + \": \" + entry.getValue()); }",
"for (char c : str.toCharArray()) { System.out.println(c); }",
"for (int i = 0; i < 10; i++) { if (i == 5) break; System.out.println(i); }",
"for (int i = 0; i < 10; i++) { if (i == 5) continue; System.out.println(i); }",
"public void myFunction() { System.out.println(\"Function executed\"); }",
"public void myFunction(String param = \"default\") { System.out.println(param); }",
"public int[] returnMultipleValues() { return new int[]{1, 2}; }",
"public void myFunction(int... numbers) { for (int num : numbers) System.out.println(num); }",
"public void myFunction(Map<String, Object> kwargs) { System.out.println(kwargs); }",
"long startTime = System.nanoTime(); myFunction(); long endTime = System.nanoTime(); long duration = (endTime - startTime);",
"public static <T> Function<T, T> memoize(Function<T, T> function) { return new Function<T, T>() { private final Map<T, T> cache = new HashMap<>(); public T apply(T input) { return cache.computeIfAbsent(input, function); } }; }",
"public class MyGenerator implements Iterator<Integer> { private int current = 0; public boolean hasNext() { return current < 10; } public Integer next() { return current++; } }",
"public class MyGenerator { public Iterator<Integer> iterator() { return new Iterator<Integer>() { private int current = 0; public boolean hasNext() { return current < 10; } public Integer next() { return current++; } }; } }",
"MyGenerator generator = new MyGenerator(); while (generator.hasNext()) { System.out.println(generator.next()); }",
"for (int i = 0; i < list.size(); i++) { System.out.println(i + \": \" + list.get(i)); }",
"List<Integer> list1 = Arrays.asList(1, 2, 3); List<Integer> list2 = Arrays.asList(4, 5, 6); List<Integer> zipped = IntStream.range(0, Math.min(list1.size(), list2.size())).mapToObj(i -> list1.get(i) + list2.get(i)).collect(Collectors.toList());",
"List<String> keys = Arrays.asList(\"a\", \"b\"); List<Integer> values = Arrays.asList(1, 2); Map<String, Integer> map = IntStream.range(0, keys.size()).boxed().collect(Collectors.toMap(keys::get, values::get));",
"boolean areEqual = list1.equals(list2);",
"boolean areEqual = map1.equals(map2);",
"boolean areEqual = set1.equals(set2);",
"Set<Integer> set = new HashSet<>(Arrays.asList(1, 2, 2, 3));",
"set.clear();",
"boolean isEmpty = set.isEmpty();",
"set.add(4);",
"set.remove(3);",
"boolean contains = set.contains(2);",
"int size = set.size();",
"boolean hasIntersection = !Collections.disjoint(set1, set2);",
"boolean isSubset = new HashSet<>(list1).containsAll(list2);",
"boolean isSubstring = str.contains(sub);",
"char firstChar = str.charAt(0);",
"char lastChar = str.charAt(str.length() - 1);",
"boolean isTextFile = Files.probeContentType(Paths.get(\"file.txt\")).startsWith(\"text\");",
"boolean isImageFile = Files.probeContentType(Paths.get(\"file.jpg\")).startsWith(\"image\");",
"double rounded = Math.round(3.14159);",
"double ceil = Math.ceil(3.14159);",
"double floor = Math.floor(3.14159);",
"double formatted = Double.parseDouble(String.format(\"%.2f\", 3.14159));",
"String randomString = new Random().ints(48, 122).filter(i -> (i <= 57 || i >= 65) && (i <= 90 || i >= 97)).limit(10).collect(StringBuilder::new, StringBuilder::appendCodePoint, StringBuilder::append).toString();",
"boolean pathExists = Files.exists(Paths.get(\"path\"));",
"Files.walk(Paths.get(\"path\")).forEach(System.out::println);",
"String extension = FilenameUtils.getExtension(\"file.txt\");",
"String fileName = FilenameUtils.getName(\"path/to/file.txt\");",
"String fullPath = Paths.get(\"path/to/file.txt\").toAbsolutePath().toString();",
"String version = System.getProperty(\"java.version\");",
"String os = System.getProperty(\"os.name\");",
"int cores = Runtime.getRuntime().availableProcessors();",
"long memory = Runtime.getRuntime().totalMemory();",
"File[] roots = File.listRoots(); for (File root : roots) { System.out.println(root.getUsableSpace()); }",
"String ip = InetAddress.getLocalHost().getHostAddress();",
"boolean isConnected = InetAddress.getByName(\"www.google.com\").isReachable(5000);",
"URL url = new URL(\"http://example.com/file.txt\"); Files.copy(url.openStream(), Paths.get(\"file.txt\"));",
"String response = new HttpPost(\"http://example.com/upload\").execute().returnContent().asString();",
"String response = new HttpPost(\"http://example.com\").execute().returnContent().asString();",
"String response = new HttpPost(\"http://example.com\").addParameter(\"key\", \"value\").execute().returnContent().asString();",
"String response = new HttpPost(\"http://example.com\").addHeader(\"key\", \"value\").execute().returnContent().asString();",
"Document doc = Jsoup.connect(\"http://example.com\").get();",
"String title = doc.title();",
"Elements links = doc.select(\"a[href]\");",
"Elements images = doc.select(\"img[src]\"); for (Element img : images) { URL imgUrl = new URL(img.attr(\"src\")); Files.copy(imgUrl.openStream(), Paths.get(imgUrl.getFile())); }",
"Map<String, Integer> wordCount = new HashMap<>(); for (String word : doc.text().split(\"\\\\s+\")) { wordCount.put(word, wordCount.getOrDefault(word, 0) + 1); }",
"String response = new HttpPost(\"http://example.com/login\").addParameter(\"username\", \"user\").addParameter(\"password\", \"pass\").execute().returnContent().asString();",
"String text = Jsoup.parse(html).text();",
"Pattern emailPattern = Pattern.compile(\"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\\\.[a-zA-Z]{2,}\"); Matcher matcher = emailPattern.matcher(text); while (matcher.find()) { System.out.println(matcher.group()); }",
"Pattern phonePattern = Pattern.compile(\"\\\\+?\\\\d{10,13}\"); Matcher matcher = phonePattern.matcher(text); while (matcher.find()) { System.out.println(matcher.group()); }",
"Pattern numberPattern = Pattern.compile(\"\\\\d+\"); Matcher matcher = numberPattern.matcher(text); while (matcher.find()) { System.out.println(matcher.group()); }",
"String replaced = text.replaceAll(\"pattern\", \"replacement\");",
"boolean matches = text.matches(\"pattern\");",
"String stripped = Jsoup.parse(html).text();",
"String encoded = StringEscapeUtils.escapeHtml4(html);",
"String decoded = StringEscapeUtils.unescapeHtml4(html);",
"JFrame frame = new JFrame(\"Simple GUI\"); frame.setSize(300, 200); frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE); frame.setVisible(true);"
"JButton button = new JButton(\"Click Me\"); frame.add(button);",
"button.addActionListener(e -> System.out.println(\"Button clicked\"));",
"JOptionPane.showMessageDialog(frame, \"This is a message\");",
"JTextField textField = new JTextField(); String input = textField.getText();",
"frame.setTitle(\"New Title\");",
"frame.setSize(400, 300);",
"frame.setLocationRelativeTo(null);",
"JMenuBar menuBar = new JMenuBar(); JMenu menu = new JMenu(\"File\"); menuBar.add(menu); frame.setJMenuBar(menuBar);",
"JComboBox<String> comboBox = new JComboBox<>(new String[]{\"Option 1\", \"Option 2\"}); frame.add(comboBox);",
"JRadioButton radioButton = new JRadioButton(\"Select Me\"); frame.add(radioButton);",
"JCheckBox checkBox = new JCheckBox(\"Check Me\"); frame.add(checkBox);",
"JLabel label = new JLabel(new ImageIcon(\"image.jpg\")); frame.add(label);",
"Clip clip = AudioSystem.getClip(); clip.open(AudioSystem.getAudioInputStream(new File(\"audio.wav\"))); clip.start();",
"Player player = Manager.createRealizedPlayer(new File(\"video.mp4\").toURI().toURL()); player.start();",
"long currentTime = player.getTimeNanoseconds();",
"Robot robot = new Robot(); BufferedImage screenShot = robot.createScreenCapture(new Rectangle(Toolkit.getDefaultToolkit().getScreenSize()));",
"ScreenRecorder screenRecorder = new ScreenRecorder(new ScreenRecorderParameters()); screenRecorder.start(); Thread.sleep(5000); screenRecorder.stop();",
"PointerInfo pointerInfo = MouseInfo.getPointerInfo(); Point point = pointerInfo.getLocation();",
"robot.keyPress(KeyEvent.VK_A); robot.keyRelease(KeyEvent.VK_A);",
"robot.mousePress(InputEvent.BUTTON1_DOWN_MASK); robot.mouseRelease(InputEvent.BUTTON1_DOWN_MASK);",
"long timestamp = System.currentTimeMillis();",
"Date date = new Date(timestamp);",
"long timestamp = date.getTime();",
"int dayOfWeek = Calendar.getInstance().get(Calendar.DAY_OF_WEEK);",
"int daysInMonth = Calendar.getInstance().getActualMaximum(Calendar.DAY_OF_MONTH);",
"Calendar cal = Calendar.getInstance(); cal.set(Calendar.DAY_OF_YEAR, 1); Date firstDay = cal.getTime();",
"Calendar cal = Calendar.getInstance(); cal.set(Calendar.DAY_OF_YEAR, cal.getActualMaximum(Calendar.DAY_OF_YEAR)); Date lastDay = cal.getTime();",
"Calendar cal = Calendar.getInstance(); cal.set(Calendar.DAY_OF_MONTH, 1); Date firstDay = cal.getTime();",
"Calendar cal = Calendar.getInstance(); cal.set(Calendar.DAY_OF_MONTH, cal.getActualMaximum(Calendar.DAY_OF_MONTH)); Date lastDay = cal.getTime();",
"int dayOfWeek = Calendar.getInstance().get(Calendar.DAY_OF_WEEK); boolean isWeekday = dayOfWeek != Calendar.SATURDAY && dayOfWeek != Calendar.SUNDAY;",
"int dayOfWeek = Calendar.getInstance().get(Calendar.DAY_OF_WEEK); boolean isWeekend = dayOfWeek == Calendar.SATURDAY || dayOfWeek == Calendar.SUNDAY;",
"int hour = Calendar.getInstance().get(Calendar.HOUR_OF_DAY);",
"int minute = Calendar.getInstance().get(Calendar.MINUTE);",
"int second = Calendar.getInstance().get(Calendar.SECOND);",
"Thread.sleep(1000);",
"long millis = System.currentTimeMillis();",
"SimpleDateFormat sdf = new SimpleDateFormat(\"yyyy-MM-dd HH:mm:ss\"); String formattedDate = sdf.format(new Date());",
"SimpleDateFormat sdf = new SimpleDateFormat(\"yyyy-MM-dd HH:mm:ss\"); Date date = sdf.parse(\"2023-10-01 12:00:00\");",
"Thread thread = new Thread(() -> System.out.println(\"Thread running\")); thread.start();",
"Thread.sleep(1000);",
"ExecutorService executor = Executors.newFixedThreadPool(2); executor.submit(() -> System.out.println(\"Task 1\")); executor.submit(() -> System.out.println(\"Task 2\"));",
"String threadName = Thread.currentThread().getName();",
"thread.setDaemon(true);",
"ReentrantLock lock = new ReentrantLock(); lock.lock(); try { System.out.println(\"Locked\"); } finally { lock.unlock(); }",
"Process process = Runtime.getRuntime().exec(\"notepad.exe\");",
"long pid = process.pid();",
"boolean isAlive = process.isAlive();",
"ExecutorService executor = Executors.newFixedThreadPool(2); executor.submit(() -> System.out.println(\"Task 1\")); executor.submit(() -> System.out.println(\"Task 2\"));",
"BlockingQueue<String> queue = new LinkedBlockingQueue<>(); queue.put(\"Message\"); String message = queue.take();",
"PipedInputStream in = new PipedInputStream(); PipedOutputStream out = new PipedOutputStream(in); out.write(\"Message\".getBytes());",
"Thread.sleep(1000);",
"Process process = Runtime.getRuntime().exec(\"ls\"); BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream())); String line; while ((line = reader.readLine()) != null) { System.out.println(line); }",
"int exitCode = process.waitFor();",
"boolean success = exitCode == 0;",
"String scriptPath = new File(\"\").getAbsolutePath();",
"String[] args = new String[]{\"arg1\", \"arg2\"};",
"ArgumentParser parser = ArgumentParser(); parser.addArgument(\"--arg\"); Namespace ns = parser.parseArgs(args);",
"parser.printHelp();",
"ModuleFinder finder = ModuleFinder.ofSystem(); Set<ModuleReference> modules = finder.findAll();",
"Process process = Runtime.getRuntime().exec(\"pip install package\");",
"Process process = Runtime.getRuntime().exec(\"pip uninstall package\");",
"String version = Package.getPackage(\"package\").getImplementationVersion();",
"Process process = Runtime.getRuntime().exec(\"python -m venv venv\");",
"Process process = Runtime.getRuntime().exec(\"pip list\");",
"Process process = Runtime.getRuntime().exec(\"pip install --upgrade package\");",
"Connection conn = DriverManager.getConnection(\"jdbc:sqlite:sample.db\");",
"Statement stmt = conn.createStatement(); ResultSet rs = stmt.executeQuery(\"SELECT * FROM table\");",
"stmt.executeUpdate(\"INSERT INTO table (column) VALUES ('value')\");",
"stmt.executeUpdate(\"DELETE FROM table WHERE id = 1\");",
"stmt.executeUpdate(\"UPDATE table SET column = 'new_value' WHERE id = 1\");",
"while (rs.next()) { System.out.println(rs.getString(\"column\")); }",
"PreparedStatement pstmt = conn.prepareStatement(\"SELECT * FROM table WHERE id = ?\"); pstmt.setInt(1, 1); ResultSet rs = pstmt.executeQuery();",
"conn.close();",
"stmt.executeUpdate(\"CREATE TABLE table (id INT, name TEXT)\");",
"stmt.executeUpdate(\"DROP TABLE table\");",
"DatabaseMetaData meta = conn.getMetaData(); ResultSet rs = meta.getTables(null, null, \"table\", null); boolean exists = rs.next();",
"ResultSet rs = meta.getTables(null, null, \"%\", null); while (rs.next()) { System.out.println(rs.getString(\"TABLE_NAME\")); }",
"EntityManager em = Persistence.createEntityManagerFactory(\"pu\").createEntityManager(); em.getTransaction().begin(); em.persist(entity); em.getTransaction().commit();",
"EntityManager em = Persistence.createEntityManagerFactory(\"pu\").createEntityManager(); Query query = em.createQuery(\"SELECT e FROM Entity e\"); List<Entity> entities = query.getResultList();",
"EntityManager em = Persistence.createEntityManagerFactory(\"pu\").createEntityManager(); em.getTransaction().begin(); em.remove(entity); em.getTransaction().commit();",
"EntityManager em = Persistence.createEntityManagerFactory(\"pu\").createEntityManager(); em.getTransaction().begin(); em.merge(entity); em.getTransaction().commit();",
"@Entity public class Entity { @Id @GeneratedValue private Long id; private String name; }",
"@Entity public class ChildEntity extends ParentEntity { private String childField; }",
"@Id @GeneratedValue private Long id;",
"@Column(unique = true) private String uniqueField;",
"@Column(columnDefinition = \"varchar(255) default 'default_value'\") private String field;",
"CSVWriter writer = new CSVWriter(new FileWriter(\"data.csv\")); writer.writeAll(data); writer.close();",
"Workbook workbook = new XSSFWorkbook(); Sheet sheet = workbook.createSheet(\"Sheet1\"); FileOutputStream fileOut = new FileOutputStream(\"workbook.xlsx\"); workbook.write(fileOut); fileOut.close();",
"ObjectMapper mapper = new ObjectMapper(); mapper.writeValue(new File(\"data.json\"), data);",
"Workbook workbook = WorkbookFactory.create(new File(\"workbook.xlsx\")); Sheet sheet = workbook.getSheetAt(0);",
"Workbook workbook1 = WorkbookFactory.create(new File(\"workbook1.xlsx\")); Workbook workbook2 = WorkbookFactory.create(new File(\"workbook2.xlsx\")); Sheet sheet1 = workbook1.getSheetAt(0); Sheet sheet2 = workbook2.getSheetAt(0); Sheet newSheet = workbook1.createSheet(\"Merged\");",
"Workbook workbook = new XSSFWorkbook(); Sheet sheet = workbook.createSheet(\"New Sheet\");",
"CellStyle style = workbook.createCellStyle(); style.cloneStyleFrom(cell.getCellStyle());",
"CellStyle style = workbook.createCellStyle(); style.setFillForegroundColor(IndexedColors.RED.getIndex()); style.setFillPattern(FillPatternType.SOLID_FOREGROUND);",
"Font font = workbook.createFont(); font.setBold(true); CellStyle style = workbook.createCellStyle(); style.setFont(font);",
"Row row = sheet.getRow(0); Cell cell = row.getCell(0); String cellValue = cell.getStringCellValue();",
"Row row = sheet.createRow(0); Cell cell = row.createCell(0); cell.setCellValue(\"New Value\");",
"BufferedImage img = ImageIO.read(new File(\"image.jpg\")); int width = img.getWidth(); int height = img.getHeight();",
"BufferedImage resizedImg = new BufferedImage(newWidth, newHeight, BufferedImage.TYPE_INT_RGB); Graphics2D g = resizedImg.createGraphics(); g.drawImage(img, 0, 0, newWidth, newHeight, null); g.dispose();"




]

# 全局服务状态
service_ready = False

# 优雅关闭处理
def handle_shutdown(signum, frame):
    app.logger.info("收到终止信号，开始关闭...")
    sys.exit(0)

signal.signal(signal.SIGTERM, handle_shutdown)
signal.signal(signal.SIGINT, handle_shutdown)

# 初始化模型和预计算编码
try:
    app.logger.info("开始加载模型...")
    model = SentenceTransformer(
        "flax-sentence-embeddings/st-codesearch-distilroberta-base",
        cache_folder="/model-cache"
    )
    # 预计算代码片段的编码（强制使用 CPU）
    code_emb = model.encode(CODE_SNIPPETS, convert_to_tensor=True, device="cpu")
    service_ready = True
    app.logger.info("服务初始化完成")
except Exception as e:
    app.logger.error("初始化失败: %s", str(e))
    raise

# Hugging Face 健康检查端点，必须响应根路径
@app.route('/')
def hf_health_check():
    # 如果请求接受 HTML，则返回一个简单的 HTML 页面（包含测试链接）
    if request.accept_mimetypes.accept_html:
        html = """
        <h2>CodeSearch API</h2>
        <p>服务状态：{{ status }}</p>
        <p>你可以在地址栏输入 /search?query=你的查询 来测试接口</p>
        """
        status = "ready" if service_ready else "initializing"
        return render_template_string(html, status=status)
    # 否则返回 JSON 格式的健康检查
    if service_ready:
        return jsonify({"status": "ready"}), 200
    else:
        return jsonify({"status": "initializing"}), 503

# 搜索 API 端点，同时支持 GET 和 POST 请求
@app.route('/search', methods=['GET', 'POST'])
def handle_search():
    if not service_ready:
        app.logger.info("服务未就绪")
        return jsonify({"error": "服务正在初始化"}), 503

    try:
        # 根据请求方法提取查询内容
        if request.method == 'GET':
            query = request.args.get('query', '').strip()
        else:
            data = request.get_json() or {}
            query = data.get('query', '').strip()

        if not query:
            app.logger.info("收到空的查询请求")
            return jsonify({"error": "查询不能为空"}), 400

        # 记录接收到的查询
        app.logger.info("收到查询请求: %s", query)

        # 对查询进行编码，并进行语义搜索
        query_emb = model.encode(query, convert_to_tensor=True, device="cpu")
        hits = util.semantic_search(query_emb, code_emb, top_k=1)[0]
        best = hits[0]

        result = {
            "code": CODE_SNIPPETS[best['corpus_id']],
            "score": round(float(best['score']), 4)
        }

        # 记录返回结果
        app.logger.info("返回结果: %s", result)
        return jsonify(result)

    except Exception as e:
        app.logger.error("请求处理失败: %s", str(e))
        return jsonify({"error": "服务器内部错误"}), 500

if __name__ == "__main__":
    # 本地测试用，Hugging Face Spaces 通常通过 gunicorn 启动
    app.run(host='0.0.0.0', port=7860)