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 = [

    "fmt.Print(\"Hello, World!\")",
    "func sum(a, b int) int { return a + b }",
    "func randInt() int { return rand.Intn(100) }",
    "func isEven(n int) bool { return n%2 == 0 }",
    "func strLen(s string) int { return len(s) }",
    "func currentDate() string { return time.Now().Format(\"2006-01-02\") }",
    "func fileExists(path string) bool { _, err := os.Stat(path); return !os.IsNotExist(err) }",
    "func readFile(path string) (string, error) { data, err := os.ReadFile(path); return string(data), err }",
    "func writeFile(path, content string) error { return os.WriteFile(path, []byte(content), 0644) }",
    "func currentTime() string { return time.Now().Format(\"15:04:05\") }",
    "func toUpper(s string) string { return strings.ToUpper(s) }",
    "func toLower(s string) string { return strings.ToLower(s) }",
    "func reverseString(s string) string { runes := []rune(s); for i, j := 0, len(runes)-1; i < j; i, j = i+1, j-1 { runes[i], runes[j] = runes[j], runes[i] }; return string(runes) }",
    "func listLen(list []int) int { return len(list) }",
    "func maxInList(list []int) int { return max(list...) }",
    "func minInList(list []int) int { return min(list...) }",
    "func sortList(list []int) []int { sort.Ints(list); return list }",
    "func mergeLists(list1, list2 []int) []int { return append(list1, list2...) }",
    "func removeFromList(list []int, elem int) []int { return append(list[:elem], list[elem+1:]...) }",
    "func isListEmpty(list []int) bool { return len(list) == 0 }",
    "func countChar(s string, c rune) int { return strings.Count(s, string(c)) }",
    "func containsSubstring(s, substr string) bool { return strings.Contains(s, substr) }",
    "func intToString(n int) string { return strconv.Itoa(n) }",
    "func stringToInt(s string) int { n, _ := strconv.Atoi(s); return n }",
    "func isNumeric(s string) bool { _, err := strconv.Atoi(s); return err == nil }",
    "func indexInList(list []int, elem int) int { for i, v := range list { if v == elem { return i } }; return -1 }",
    "func clearList(list []int) []int { return list[:0] }",
    "func reverseList(list []int) []int { for i, j := 0, len(list)-1; i < j; i, j = i+1, j-1 { list[i], list[j] = list[j], list[i] }; return list }",
    "func removeDuplicates(list []int) []int { seen := make(map[int]bool); result := []int{}; for _, v := range list { if !seen[v] { seen[v] = true; result = append(result, v) } }; return result }",
    "func isInList(list []int, elem int) bool { for _, v := range list { if v == elem { return true } }; return false }",
    "func createDict() map[string]int { return make(map[string]int) }",
    "func addToDict(dict map[string]int, key string, value int) { dict[key] = value }",
    "func deleteFromDict(dict map[string]int, key string) { delete(dict, key) }",
    "func dictKeys(dict map[string]int) []string { keys := []string{}; for k := range dict { keys = append(keys, k) }; return keys }",
    "func dictValues(dict map[string]int) []int { values := []int{}; for _, v := range dict { values = append(values, v) }; return values }",
    "func mergeDicts(dict1, dict2 map[string]int) map[string]int { for k, v := range dict2 { dict1[k] = v }; return dict1 }",
    "func isDictEmpty(dict map[string]int) bool { return len(dict) == 0 }",
    "func dictValue(dict map[string]int, key string) int { return dict[key] }",
    "func keyInDict(dict map[string]int, key string) bool { _, ok := dict[key]; return ok }",
    "func clearDict(dict map[string]int) { for k := range dict { delete(dict, k) } }",
    "func countFileLines(path string) (int, error) { data, err := os.ReadFile(path); if err != nil { return 0, err }; return len(strings.Split(string(data), \"\n\")), nil }",
    "func writeListToFile(path string, list []int) error { data := strings.Trim(strings.Join(strings.Fields(fmt.Sprint(list)), \" \"), \"[]\"); return os.WriteFile(path, []byte(data), 0644) }",
    "func readListFromFile(path string) ([]int, error) { data, err := os.ReadFile(path); if err != nil { return nil, err }; var list []int; for _, v := range strings.Fields(string(data)) { n, _ := strconv.Atoi(v); list = append(list, n) }; return list, nil }",
    "func countFileWords(path string) (int, error) { data, err := os.ReadFile(path); if err != nil { return 0, err }; return len(strings.Fields(string(data))), nil }",
    "func isLeapYear(year int) bool { return year%4 == 0 && (year%100 != 0 || year%400 == 0) }",
    "func formatTime(t time.Time, layout string) string { return t.Format(layout) }",
    "func daysBetweenDates(date1, date2 time.Time) int { return int(date2.Sub(date1).Hours() / 24) }",
    "func currentDir() string { dir, _ := os.Getwd(); return dir }",
    "func listFiles(path string) ([]string, error) { files, err := os.ReadDir(path); if err != nil { return nil, err }; var names []string; for _, file := range files { names = append(names, file.Name()) }; return names, nil }",
    "func createDir(path string) error { return os.Mkdir(path, 0755) }",
    "func deleteDir(path string) error { return os.RemoveAll(path) }",
    "func isFile(path string) bool { info, err := os.Stat(path); return err == nil && !info.IsDir() }",
    "func isDir(path string) bool { info, err := os.Stat(path); return err == nil && info.IsDir() }",
    "func fileSize(path string) (int64, error) { info, err := os.Stat(path); if err != nil { return 0, err }; return info.Size(), nil }",
    "func renameFile(oldPath, newPath string) error { return os.Rename(oldPath, newPath) }",
    "func copyFile(src, dst string) error { data, err := os.ReadFile(src); if err != nil { return err }; return os.WriteFile(dst, data, 0644) }",
    "func moveFile(src, dst string) error { return os.Rename(src, dst) }",
    "func deleteFile(path string) error { return os.Remove(path) }",
    "func getEnvVar(key string) string { return os.Getenv(key) }",
    "func setEnvVar(key, value string) error { return os.Setenv(key, value) }",
    "func openURL(url string) error { return exec.Command(\"xdg-open\", url).Start() }",
    "func sendGetRequest(url string) (string, error) { resp, err := http.Get(url); if err != nil { return \"\", err }; defer resp.Body.Close(); body, err := io.ReadAll(resp.Body); if err != nil { return \"\", err }; return string(body), nil }",
    "func parseJSON(data string) (map[string]interface{}, error) { var result map[string]interface{}; err := json.Unmarshal([]byte(data), &result); return result, err }",
    "func writeJSONToFile(path string, data map[string]interface{}) error { jsonData, err := json.Marshal(data); if err != nil { return err }; return os.WriteFile(path, jsonData, 0644) }",
    "func readJSONFromFile(path string) (map[string]interface{}, error) { data, err := os.ReadFile(path); if err != nil { return nil, err }; var result map[string]interface{}; err = json.Unmarshal(data, &result); return result, err }",
    "func listToString(list []int) string { return strings.Trim(strings.Join(strings.Fields(fmt.Sprint(list)), \" \"), \"[]\") }",
    "func stringToList(s string) []int { var list []int; for _, v := range strings.Fields(s) { n, _ := strconv.Atoi(v); list = append(list, n) }; return list }",
    "func joinListWithComma(list []int) string { return strings.Join(strings.Fields(fmt.Sprint(list)), \",\") }",
    "func joinListWithNewline(list []int) string { return strings.Join(strings.Fields(fmt.Sprint(list)), \"\n\") }",
    "func splitStringBySpace(s string) []string { return strings.Fields(s) }",
    "func splitStringByDelimiter(s, delimiter string) []string { return strings.Split(s, delimiter) }",
    "func splitStringToChars(s string) []string { return strings.Split(s, \"\") }",
    "func replaceString(s, old, new string) string { return strings.ReplaceAll(s, old, new) }",
    "func removeSpaces(s string) string { return strings.ReplaceAll(s, \" \", \"\") }",
    "func removePunctuation(s string) string { return strings.Map(func(r rune) rune { if unicode.IsPunct(r) { return -1 }; return r }, s) }",
    "func isStringEmpty(s string) bool { return len(s) == 0 }",
    "func isPalindrome(s string) bool { return s == reverseString(s) }",
    "func writeCSV(path string, data [][]string) error { file, err := os.Create(path); if err != nil { return err }; defer file.Close(); writer := csv.NewWriter(file); return writer.WriteAll(data) }",
    "func readCSV(path string) ([][]string, error) { file, err := os.Open(path); if err != nil { return nil, err }; defer file.Close(); reader := csv.NewReader(file); return reader.ReadAll() }",
    "func countCSVLines(path string) (int, error) { data, err := readCSV(path); if err != nil { return 0, err }; return len(data), nil }",
    "func shuffleList(list []int) []int { rand.Shuffle(len(list), func(i, j int) { list[i], list[j] = list[j], list[i] }); return list }",
    "func randomElement(list []int) int { return list[rand.Intn(len(list))] }",
    "func randomElements(list []int, n int) []int { rand.Shuffle(len(list), func(i, j int) { list[i], list[j] = list[j], list[i] }); return list[:n] }",
    "func rollDice() int { return rand.Intn(6) + 1 }",
    "func flipCoin() string { if rand.Intn(2) == 0 { return \"Heads\" }; return \"Tails\" }",
    "func generatePassword(length int) string { chars := \"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()\"; password := make([]byte, length); for i := range password { password[i] = chars[rand.Intn(len(chars))] }; return string(password) }",
    "func randomColor() string { return fmt.Sprintf(\"#%06x\", rand.Intn(0xffffff)) }",
    "func generateUUID() string { return uuid.New().String() }",
    "type MyClass struct{}",
    "func NewMyClass() *MyClass { return &MyClass{} }",
    "func (m *MyClass) MyMethod() {}",
    "type MyClass struct { MyAttribute int }",
    "type ChildClass struct { MyClass }",
    "func (c *ChildClass) MyMethod() {}",
    "func (m *MyClass) MyClassMethod() {}",
    "func MyStaticMethod() {}",
    "func isInstanceOf(obj interface{}, t reflect.Type) bool { return reflect.TypeOf(obj) == t }",
    "func getAttribute(obj interface{}, attr string) interface{} { return reflect.ValueOf(obj).FieldByName(attr).Interface() }",
    "func setAttribute(obj interface{}, attr string, value interface{}) { reflect.ValueOf(obj).FieldByName(attr).Set(reflect.ValueOf(value)) }",
    "func deleteAttribute(obj interface{}, attr string) { reflect.ValueOf(obj).FieldByName(attr).Set(reflect.Zero(reflect.TypeOf(obj).FieldByName(attr).Type)) }",
"func handleException() { defer func() { if r := recover(); r != nil { fmt.Println(\"Recovered:\", r) } }(); panic(\"An error occurred\") }",
    """func customError() error {
    return errors.New("custom error")
}""",
    """func getErrorInfo() {
    if err := customError(); err != nil {
        fmt.Println("Error:", err)
    }
}""",
    """func logError() {
    if err := customError(); err != nil {
        log.Println("Error:", err)
    }
}""",
    """func timer() func() {
    start := time.Now()
    return func() {
        fmt.Println(time.Since(start))
    }
}""",
    """func runtime() {
    start := time.Now()
    time.Sleep(1 * time.Second)
    fmt.Println(time.Since(start))
}""",
    """func progressBar() {
    for i := 0; i <= 100; i += 10 {
        fmt.Printf("\r[%-10s] %d%%", strings.Repeat("=", i/10), i)
        time.Sleep(100 * time.Millisecond)
    }
    fmt.Println()
}""",
    """func delay() {
    time.Sleep(1 * time.Second)
}""",
    "lambda := func(x int) int { return x * x }",
    "result := map(lambda, []int{1, 2, 3, 4})",
    "result := filter(func(x int) bool { return x > 2 }, []int{1, 2, 3, 4})",
    "result := reduce(func(a, b int) int { return a + b }, []int{1, 2, 3, 4})",
    "result := [x * x for x in [1, 2, 3, 4]]",
    "result := {x: x * x for x in [1, 2, 3, 4]}",
    "result := {x * x for x in [1, 2, 3, 4]}",
    "result := intersection(set1, set2)",
    "result := union(set1, set2)",
    "result := difference(set1, set2)",
    "result := [x for x in [1, None, 2, None, 3] if x is not None]",
    """func checkFile(filename string) {
    file, err := os.Open(filename)
    if err != nil {
        fmt.Println("Error:", err)
    } else {
        file.Close()
    }
}""",
    """func checkType(v interface{}) {
    fmt.Printf("Type: %T\n", v)
}""",
    """func strToBool(s string) bool {
    return s == "true"
}""",
    """func ifCondition(x int) {
    if x > 0 {
        fmt.Println("Positive")
    } else {
        fmt.Println("Non-positive")
    }
}""",
    """func whileLoop() {
    i := 0
    while i < 10 {
        fmt.Println(i)
        i++
    }
}""",
    """func forList() {
    for _, v := range []int{1, 2, 3} {
        fmt.Println(v)
    }
}""",
    """func forDict() {
    for k, v := range map[string]int{"a": 1, "b": 2} {
        fmt.Println(k, v)
    }
}""",
    """func forString() {
    for _, c := range "hello" {
        fmt.Println(c)
    }
}""",
    """func breakLoop() {
    for i := 0; i < 10; i++ {
        if i == 5 {
            break
        }
        fmt.Println(i)
    }
}""",
    """func continueLoop() {
    for i := 0; i < 10; i++ {
        if i == 5 {
            continue
        }
        fmt.Println(i)
    }
}""",
    """func defineFunc() {
    fmt.Println("Function defined")
}""",
    """func defaultParam(x int = 10) {
    fmt.Println(x)
}""",
    """func returnMultiple() (int, string) {
    return 1, "hello"
}""",
    """func variadicParams(nums ...int) {
    fmt.Println(nums)
}""",
    """func keywordParams(a int, b string) {
    fmt.Println(a, b)
}""",
    """func measureTime() {
    start := time.Now()
    time.Sleep(1 * time.Second)
    fmt.Println(time.Since(start))
}""",
    """func decorator(f func()) func() {
    return func() {
        fmt.Println("Before")
        f()
        fmt.Println("After")
    }
}""",
    """func cacheResult(f func() int) func() int {
    var result int
    return func() int {
        if result == 0 {
            result = f()
        }
        return result
    }
}""",
    """func createGenerator() func() int {
    i := 0
    return func() int {
        i++
        return i
    }
}""",
    """func yieldValue() chan int {
    ch := make(chan int)
    go func() {
        ch <- 1
        ch <- 2
        close(ch)
    }()
    return ch
}""",
    """func nextValue() {
    ch := yieldValue()
    fmt.Println(<-ch)
    fmt.Println(<-ch)
}""",
    """func createIterator() func() (int, bool) {
    i := 0
    return func() (int, bool) {
        if i < 3 {
            i++
            return i, true
        }
        return 0, false
    }
}""",
    """func manualIterate() {
    it := createIterator()
    for {
        v, ok := it()
        if !ok {
            break
        }
        fmt.Println(v)
    }
}""",
    """func useEnumerate() {
    for i, v := range []string{"a", "b", "c"} {
        fmt.Println(i, v)
    }
}""",
    """func useZip() {
    for a, b := range zip([]int{1, 2}, []string{"a", "b"}) {
        fmt.Println(a, b)
    }
}""",
    """func listToDict() {
    dict := make(map[int]string)
    for i, v := range []string{"a", "b"} {
        dict[i] = v
    }
    fmt.Println(dict)
}""",
    """func compareLists() {
    fmt.Println(reflect.DeepEqual([]int{1, 2}, []int{1, 2}))
}""",
    """func compareDicts() {
    fmt.Println(reflect.DeepEqual(map[string]int{"a": 1}, map[string]int{"a": 1}))
}""",
    """func compareSets() {
    fmt.Println(reflect.DeepEqual(map[int]bool{1: true}, map[int]bool{1: true}))
}""",
    """func removeDuplicates() {
    set := make(map[int]bool)
    for _, v := range []int{1, 2, 2, 3} {
        set[v] = true
    }
    fmt.Println(set)
}""",
    """func clearSet() {
    set := map[int]bool{1: true, 2: true}
    for k := range set {
        delete(set, k)
    }
    fmt.Println(set)
}""",
"func isEmptySet(s map[interface{}]struct{}) bool { return len(s) == 0 }", "func addToSet(s map[interface{}]struct{}, element interface{}) { s[element] = struct{}{} }", "func removeFromSet(s map[interface{}]struct{}, element interface{}) { delete(s, element) }", "func containsInSet(s map[interface{}]struct{}, element interface{}) bool { _, exists := s[element]; return exists }", "func getSetLength(s map[interface{}]struct{}) int { return len(s) }", "func hasIntersection(s1, s2 map[interface{}]struct{}) bool { for k := range s1 { if _, exists := s2[k]; exists { return true } } return false }", "func isSubset(s1, s2 map[interface{}]struct{}) bool { for k := range s1 { if _, exists := s2[k]; !exists { return false } } return true }", "func isSubstring(s, substr string) bool { return strings.Contains(s, substr) }", "func getFirstChar(s string) string { return string(s[0]) }", "func getLastChar(s string) string { return string(s[len(s)-1]) }", "func isTextFile(filename string) bool { return strings.HasSuffix(filename, \".txt\") }", "func isImageFile(filename string) bool { ext := strings.ToLower(filepath.Ext(filename)); return ext == \".jpg\" || ext == \".png\" || ext == \".gif\" }", "func roundNumber(num float64) float64 { return math.Round(num) }", "func ceilNumber(num float64) float64 { return math.Ceil(num) }", "func floorNumber(num float64) float64 { return math.Floor(num) }", "func formatDecimal(num float64, precision int) string { return fmt.Sprintf(\"%.*f\", precision, num) }", "func generateRandomString(length int) string { rand.Seed(time.Now().UnixNano()); b := make([]byte, length); rand.Read(b); return fmt.Sprintf(\"%x\", b)[:length] }", "func pathExists(path string) bool { _, err := os.Stat(path); return !os.IsNotExist(err) }", """func listFilesInDir(dir string) []string {
    var files []string
    filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
        if !info.IsDir() {
            files = append(files, path)
        }
        return nil
    })
    return files
}""", "func getFileExtension(filename string) string { return filepath.Ext(filename) }", "func getFileName(filename string) string { return filepath.Base(filename) }", "func getFullPath(filename string) string { absPath, _ := filepath.Abs(filename); return absPath }", "func getPythonVersion() string { return runtime.Version() }", "func getPlatformInfo() string { return runtime.GOOS }", "func getCPUCores() int { return runtime.NumCPU() }", "func getMemorySize() uint64 { var mem runtime.MemStats; runtime.ReadMemStats(&mem); return mem.Sys }", """func getDiskUsage() string {
    var stat syscall.Statfs_t
    syscall.Statfs("/", &stat)
    return fmt.Sprintf("Total: %d, Free: %d", stat.Blocks*uint64(stat.Bsize), stat.Bavail*uint64(stat.Bsize))
}""", "func getIPAddress() string { conn, _ := net.Dial(\"udp\", \"8.8.8.8:80\"); defer conn.Close(); return conn.LocalAddr().(*net.UDPAddr).IP.String() }", "func isConnected() bool { _, err := net.LookupIP(\"google.com\"); return err == nil }", """func downloadFile(url, filename string) error {
    resp, err := http.Get(url)
    if err != nil {
        return err
    }
    defer resp.Body.Close()
    out, err := os.Create(filename)
    if err != nil {
        return err
    }
    defer out.Close()
    _, err = io.Copy(out, resp.Body)
    return err
}""", """func uploadFile(filename string) string {
    return fmt.Sprintf("File %s uploaded", filename)
}""", """func sendPostRequest(url string, data map[string]string) (string, error) {
    jsonData, _ := json.Marshal(data)
    resp, err := http.Post(url, "application/json", bytes.NewBuffer(jsonData))
    if err != nil {
        return "", err
    }
    defer resp.Body.Close()
    body, _ := io.ReadAll(resp.Body)
    return string(body), nil
}""", """func sendRequestWithParams(url string, params map[string]string) (string, error) {
    req, err := http.NewRequest("GET", url, nil)
    if err != nil {
        return "", err
    }
    q := req.URL.Query()
    for key, value := range params {
        q.Add(key, value)
    }
    req.URL.RawQuery = q.Encode()
    resp, err := http.DefaultClient.Do(req)
    if err != nil {
        return "", err
    }
    defer resp.Body.Close()
    body, _ := io.ReadAll(resp.Body)
    return string(body), nil
}""", """func setRequestHeader(req *http.Request, headers map[string]string) {
    for key, value := range headers {
        req.Header.Set(key, value)
    }
}""", """func parseHTML(htmlContent string) (*goquery.Document, error) {
    return goquery.NewDocumentFromReader(strings.NewReader(htmlContent))
}""", """func extractTitle(doc *goquery.Document) string {
    return doc.Find("title").Text()
}""", """func extractLinks(doc *goquery.Document) []string {
    var links []string
    doc.Find("a").Each(func(i int, s *goquery.Selection) {
        href, exists := s.Attr("href")
        if exists {
            links = append(links, href)
        }
    })
    return links
}""", """func downloadImagesFromPage(url string) error {
    doc, err := goquery.NewDocument(url)
    if err != nil {
        return err
    }
    doc.Find("img").Each(func(i int, s *goquery.Selection) {
        src, exists := s.Attr("src")
        if exists {
            downloadFile(src, fmt.Sprintf("image%d.jpg", i))
        }
    })
    return nil
}""", """func countWordFrequency(text string) map[string]int {
    words := strings.Fields(text)
    freq := make(map[string]int)
    for _, word := range words {
        freq[word]++
    }
    return freq
}""", """func simulateLogin(url, username, password string) (string, error) {
    form := url.Values{}
    form.Add("username", username)
    form.Add("password", password)
    resp, err := http.PostForm(url, form)
    if err != nil {
        return "", err
    }
    defer resp.Body.Close()
    body, _ := io.ReadAll(resp.Body)
    return string(body), nil
}""", """func htmlToText(htmlContent string) string {
    doc, _ := goquery.NewDocumentFromReader(strings.NewReader(htmlContent))
    return doc.Text()
}""", """func extractEmails(text string) []string {
    re := regexp.MustCompile(`[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}`)
    return re.FindAllString(text, -1)
}""", """func extractPhoneNumbers(text string) []string {
    re := regexp.MustCompile(`\d{3}-\d{3}-\d{4}`)
    return re.FindAllString(text, -1)
}""", """func findAllNumbers(text string) []string {
    re := regexp.MustCompile(`\d+`)
    return re.FindAllString(text, -1)
}""", """func replaceWithRegex(text, pattern, replacement string) string {
    re := regexp.MustCompile(pattern)
    return re.ReplaceAllString(text, replacement)
}""", """func matchRegex(text, pattern string) bool {
    re := regexp.MustCompile(pattern)
    return re.MatchString(text)
}""", """func removeHTMLTags(htmlContent string) string {
    re := regexp.MustCompile(`<[^>]*>`)
    return re.ReplaceAllString(htmlContent, "")
}""", """func encodeHTMLEntities(htmlContent string) string {
    return html.EscapeString(htmlContent)
}""", """func decodeHTMLEntities(htmlContent string) string {
    return html.UnescapeString(htmlContent)
}""", """func createGUIWindow() {
    window := gtk.NewWindow(gtk.WINDOW_TOPLEVEL)
    window.SetTitle("Simple GUI Window")
    window.Connect("destroy", func() {
        gtk.MainQuit()
    })
    window.ShowAll()
    gtk.Main()
}""",


"func addButton(w fyne.Window) { w.SetContent(widget.NewButton(\"Button\", nil)) }",
    """func handleClick(b *widget.Button) {
    b.OnTapped = func() { fmt.Println("Clicked") }
}""",
    """func showDialog(w fyne.Window) {
    dialog.ShowInformation("Title", "Message", w)
}""",
    "func getText(e *widget.Entry) string { return e.Text }",
    "func setTitle(w fyne.Window) { w.SetTitle(\"Title\") }",
    "func resizeWindow(w fyne.Window) { w.Resize(fyne.NewSize(800, 600)) }",
    "func centerWindow(w fyne.Window) { w.CenterOnScreen() }",
    """func addMenu(w fyne.Window) {
    menu := fyne.NewMainMenu(fyne.NewMenu("File"))
    w.SetMainMenu(menu)
}""",
    "func createSelect() *widget.Select { return widget.NewSelect([]string{\"a\", \"b\"}, func(s string) {}) }",
    "func createRadio() *widget.RadioGroup { return widget.NewRadioGroup([]string{\"a\", \"b\"}, func(s string) {}) }",
    "func createCheck() *widget.Check { return widget.NewCheck(\"Check\", func(b bool) {}) }",
    "func showImage() *canvas.Image { return canvas.NewImageFromFile(\"image.png\") }",
    """func playAudio() {
    exec.Command("afplay", "audio.mp3").Start()
}""",
    """func playVideo() {
    exec.Command("vlc", "video.mp4").Start()
}""",
    "func getTimestamp() int64 { return time.Now().Unix() }",
    """func timestampToDate(ts int64) string {
    return time.Unix(ts, 0).Format("2006-01-02")
}""",
    """func dateToTimestamp(date string) int64 {
    t, _ := time.Parse("2006-01-02", date)
    return t.Unix()
}""",
    "func getWeekday() string { return time.Now().Weekday().String() }",
    """func getMonthDays() int {
    return time.Date(time.Now().Year(), time.Now().Month()+1, 0, 0, 0, 0, 0, time.UTC).Day()
}""",
    "func firstDayOfYear() time.Time { return time.Date(time.Now().Year(), 1, 1, 0, 0, 0, 0, time.UTC) }",
    "func lastDayOfYear() time.Time { return time.Date(time.Now().Year(), 12, 31, 0, 0, 0, 0, time.UTC) }",
    """func firstDayOfMonth(year int, month time.Month) time.Time {
    return time.Date(year, month, 1, 0, 0, 0, 0, time.UTC)
}""",
    """func lastDayOfMonth(year int, month time.Month) int {
    return time.Date(year, month+1, 0, 0, 0, 0, 0, time.UTC).Day()
}""",
    "func isWorkday() bool { wd := time.Now().Weekday(); return wd != time.Saturday && wd != time.Sunday }",
    "func isWeekend() bool { wd := time.Now().Weekday(); return wd == time.Saturday || wd == time.Sunday }",
    "func currentHour() int { return time.Now().Hour() }",
    "func currentMinute() int { return time.Now().Minute() }",
    "func currentSecond() int { return time.Now().Second() }",
    "func sleep() { time.Sleep(time.Second) }",
    "func millis() int64 { return time.Now().UnixNano() / int64(time.Millisecond) }",
    """func formatTime(t time.Time) string {
    return t.Format("2006-01-02 15:04:05")
}""",
    """func parseTime(s string) time.Time {
    t, _ := time.Parse("2006-01-02", s)
    return t
}""",
    """func startThread() {
    go func() { fmt.Println("Thread") }()
}""",
    "func threadSleep() { time.Sleep(time.Second) }",
    """func multiThread() {
    for i := 0; i < 3; i++ { go func(i int) {}(i) }
}""",
    "func threadName() string { return runtime.Version() }",
    "func setDaemon() { /* Not applicable in Go */ }",
    """func useMutex() {
    var mu sync.Mutex
    mu.Lock()
    defer mu.Unlock()
}""",
    """func createProcess() {
    cmd := exec.Command("ls")
    cmd.Start()
}""",
    "func getPID() int { return os.Getpid() }",
    """func isProcessAlive(pid int) bool {
    err := syscall.Kill(pid, syscall.Signal(0))
    return err == nil
}""",
    """func multiProcess() {
    exec.Command("ls").Start()
}""",
    """func useQueue() {
    ch := make(chan int)
    go func() { ch <- 1 }()
    <-ch
}""",
    """func usePipe() {
    r, w, _ := os.Pipe()
    w.Write([]byte("data"))
    var b [4]byte
    r.Read(b[:])
}""",
"func limitCPUUsage() { runtime.GOMAXPROCS(1) }",
    """func runShellCommand(cmd string) error {
    return exec.Command("sh", "-c", cmd).Run()
}""",
    """func getCommandOutput(cmd string) (string, error) {
    out, err := exec.Command("sh", "-c", cmd).Output()
    return string(out), err
}""",
    """func getCommandStatus(cmd string) (int, error) {
    err := exec.Command("sh", "-c", cmd).Run()
    if err != nil {
        if exitError, ok := err.(*exec.ExitError); ok {
            return exitError.ExitCode(), nil
        }
        return -1, err
    }
    return 0, nil
}""",
    """func isCommandSuccess(cmd string) (bool, error) {
    err := exec.Command("sh", "-c", cmd).Run()
    return err == nil, err
}""",
    """func getCurrentScriptPath() string {
    _, filename, _, _ := runtime.Caller(1)
    return filepath.Dir(filename)
}""",
    """func getCommandLineArgs() []string {
    return os.Args
}""",
    """func parseArgs() (string, int) {
    var name string
    var age int
    flag.StringVar(&name, "name", "", "Name of the person")
    flag.IntVar(&age, "age", 0, "Age of the person")
    flag.Parse()
    return name, age
}""",
    """func generateHelp() string {
    flag.Usage()
    return ""
}""",
    """func listPythonModules() {
    exec.Command("python", "-c", "help('modules')").Run()
}""",
    """func installPythonPackage(pkg string) error {
    return exec.Command("pip", "install", pkg).Run()
}""",
    """func uninstallPythonPackage(pkg string) error {
    return exec.Command("pip", "uninstall", "-y", pkg).Run()
}""",
    """func getPackageVersion(pkg string) (string, error) {
    out, err := exec.Command("pip", "show", pkg).Output()
    if err != nil {
        return "", err
    }
    return string(out), nil
}""",
    """func useVirtualEnv(envPath string) error {
    return exec.Command("source", filepath.Join(envPath, "bin", "activate")).Run()
}""",
    """func listInstalledPackages() (string, error) {
    out, err := exec.Command("pip", "list").Output()
    return string(out), err
}""",
    """func upgradePythonPackage(pkg string) error {
    return exec.Command("pip", "install", "--upgrade", pkg).Run()
}""",
    """func connectLocalDB(dbName string) (*sql.DB, error) {
    return sql.Open("sqlite3", dbName)
}""",
    """func executeSQLQuery(db *sql.DB, query string) (*sql.Rows, error) {
    return db.Query(query)
}""",
    """func insertRecord(db *sql.DB, query string, args ...interface{}) (sql.Result, error) {
    return db.Exec(query, args...)
}""",
    """func deleteRecord(db *sql.DB, query string, args ...interface{}) (sql.Result, error) {
    return db.Exec(query, args...)
}""",
    """func updateRecord(db *sql.DB, query string, args ...interface{}) (sql.Result, error) {
    return db.Exec(query, args...)
}""",
    """func queryMultipleRecords(db *sql.DB, query string, args ...interface{}) (*sql.Rows, error) {
    return db.Query(query, args...)
}""",
    """func useParameterizedQuery(db *sql.DB, query string, args ...interface{}) (*sql.Rows, error) {
    return db.Query(query, args...)
}""",
    """func closeDBConnection(db *sql.DB) error {
    return db.Close()
}""",
    """func createTable(db *sql.DB, query string) (sql.Result, error) {
    return db.Exec(query)
}""",
    """func dropTable(db *sql.DB, query string) (sql.Result, error) {
    return db.Exec(query)
}""",
    """func tableExists(db *sql.DB, tableName string) (bool, error) {
    var exists bool
    err := db.QueryRow("SELECT EXISTS (SELECT 1 FROM sqlite_master WHERE type='table' AND name=?)", tableName).Scan(&exists)
    return exists, err
}""",
    """func listAllTables(db *sql.DB) ([]string, error) {
    rows, err := db.Query("SELECT name FROM sqlite_master WHERE type='table'")
    if err != nil {
        return nil, err
    }
    defer rows.Close()
    var tables []string
    for rows.Next() {
        var table string
        if err := rows.Scan(&table); err != nil {
            return nil, err
        }
        tables = append(tables, table)
    }
    return tables, nil
}""",
    """func exportToCSV(data [][]string, filename string) error {
    file, err := os.Create(filename)
    if err != nil {
        return err
    }
    defer file.Close()
    writer := csv.NewWriter(file)
    return writer.WriteAll(data)
}""",
    """func exportToExcel(data [][]string, filename string) error {
    file := excelize.NewFile()
    for i, row := range data {
        for j, cell := range row {
            file.SetCellValue("Sheet1", fmt.Sprintf("%c%d", 'A'+j, i+1), cell)
        }
    }
    return file.SaveAs(filename)
}""",
    """func exportToJSON(data interface{}, filename string) error {
    file, err := os.Create(filename)
    if err != nil {
        return err
    }
    defer file.Close()
    encoder := json.NewEncoder(file)
    return encoder.Encode(data)
}""",
    """func readExcel(filename string) ([][]string, error) {
    file, err := excelize.OpenFile(filename)
    if err != nil {
        return nil, err
    }
    rows, err := file.GetRows("Sheet1")
    return rows, err
}""",
    """func mergeExcelFiles(files []string, output string) error {
    merged := excelize.NewFile()
    for _, file := range files {
        f, err := excelize.OpenFile(file)
        if err != nil {
            return err
        }
        rows, err := f.GetRows("Sheet1")
        if err != nil {
            return err
        }
        for i, row := range rows {
            for j, cell := range row {
                merged.SetCellValue("Sheet1", fmt.Sprintf("%c%d", 'A'+j, i+1), cell)
            }
        }
    }
    return merged.SaveAs(output)
}""",
    """func addSheetToExcel(filename, sheetName string) error {
    file, err := excelize.OpenFile(filename)
    if err != nil {
        return err
    }
    file.NewSheet(sheetName)
    return file.Save()
}""",
    """func copyExcelStyle(src, dst string) error {
    srcFile, err := excelize.OpenFile(src)
    if err != nil {
        return err
    }
    dstFile, err := excelize.OpenFile(dst)
    if err != nil {
        return err
    }
    styleID, err := srcFile.GetCellStyle("Sheet1", "A1")
    if err != nil {
        return err
    }
    dstFile.SetCellStyle("Sheet1", "A1", "A1", styleID)
    return dstFile.Save()
}""",
    """func setExcelCellColor(filename, cell, color string) error {
    file, err := excelize.OpenFile(filename)
    if err != nil {
        return err
    }
    styleID, err := file.NewStyle(&excelize.Style{Fill: excelize.Fill{Type: "pattern", Color: []string{color}, Pattern: 1}})
    if err != nil {
        return err
    }
    file.SetCellStyle("Sheet1", cell, cell, styleID)
    return file.Save()
}""",
    """func setExcelFontStyle(filename, cell, fontName string, fontSize int) error {
    file, err := excelize.OpenFile(filename)
    if err != nil {
        return err
    }
    styleID, err := file.NewStyle(&excelize.Style{Font: &excelize.Font{Family: fontName, Size: fontSize}})
    if err != nil {
        return err
    }
    file.SetCellStyle("Sheet1", cell, cell, styleID)
    return file.Save()
}""",
    """func readExcelCell(filename, cell string) (string, error) {
    file, err := excelize.OpenFile(filename)
    if err != nil {
        return "", err
    }
    return file.GetCellValue("Sheet1", cell)
}""",
    """func writeExcelCell(filename, cell, value string) error {
    file, err := excelize.OpenFile(filename)
    if err != nil {
        return err
    }
    file.SetCellValue("Sheet1", cell, value)
    return file.Save()
}""",
    """func getImageDimensions(filename string) (int, int, error) {
    file, err := os.Open(filename)
    if err != nil {
        return 0, 0, err
    }
    defer file.Close()
    img, _, err := image.DecodeConfig(file)
    return img.Width, img.Height, err
}""",
    """func resizeImage(filename string, width, height int) error {
    file, err := os.Open(filename)
    if err != nil {
        return err
    }
    defer file.Close()
    img, _, err := image.Decode(file)
    if err != nil {
        return err
    }
    resized := imaging.Resize(img, width, height, imaging.Lanczos)
    return imaging.Save(resized, filename)
}"""







    
]

# 全局服务状态
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)