import gradio as gr
from sentence_transformers import SentenceTransformer
import numpy as np

MODEL_NAME = "Matthieufromparis/bge-small-code-search-v1"

model = SentenceTransformer(MODEL_NAME)

CODE_SNIPPETS = [
    "def parse_json_config(filepath):\n    with open(filepath) as f:\n        return json.load(f)",
    "def sort_by_key(items, key, reverse=False):\n    return sorted(items, key=lambda x: x.get(key, ''), reverse=reverse)",
    "def authenticate_user(token):\n    payload = jwt.decode(token, SECRET, algorithms=['HS256'])\n    return payload.get('user_id')",
    "def fetch_api_data(url, headers=None):\n    resp = requests.get(url, headers=headers, timeout=10)\n    resp.raise_for_status()\n    return resp.json()",
    "def validate_email(email):\n    pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$'\n    return bool(re.match(pattern, email))",
    "def connect_database(host, port, user, password, dbname):\n    conn = psycopg2.connect(host=host, port=port, user=user, password=password, dbname=dbname)\n    return conn",
    "def retry_on_failure(func, max_retries=3, delay=1):\n    for attempt in range(max_retries):\n        try:\n            return func()\n        except Exception as e:\n            if attempt == max_retries - 1:\n                raise\n            time.sleep(delay)",
    "def generate_random_string(length=12):\n    chars = string.ascii_letters + string.digits\n    return ''.join(random.choice(chars) for _ in range(length))",
    "def log_error(message, level='ERROR'):\n    logger = logging.getLogger(__name__)\n    getattr(logger, level.lower())(message)",
    "def chunk_list(items, chunk_size=100):\n    for i in range(0, len(items), chunk_size):\n        yield items[i:i + chunk_size]",
]

code_embeddings = model.encode(CODE_SNIPPETS, normalize_embeddings=True)

def search_code(query: str, top_k: int = 5):
    if not query.strip():
        return "Enter a query to search code snippets."
    query_embedding = model.encode(query, normalize_embeddings=True)
    scores = np.dot(query_embedding, code_embeddings.T)
    top_indices = np.argsort(scores)[-top_k:][::-1]
    results = []
    for idx in top_indices:
        score = float(scores[idx])
        if score < 0.3:
            continue
        results.append(f"### Score: {score:.3f}\n```python\n{CODE_SNIPPETS[idx]}\n```\n---")
    if not results:
        return "No matching code snippets found. Try a different query."
    return "\n".join(results)

def compare_embeddings(text1, text2):
    if not text1.strip() or not text2.strip():
        return "Enter both texts to compare."
    emb1 = model.encode(text1, normalize_embeddings=True)
    emb2 = model.encode(text2, normalize_embeddings=True)
    similarity = float(np.dot(emb1, emb2))
    if similarity > 0.8: level = "🟢 Very Similar"
    elif similarity > 0.6: level = "🟡 Somewhat Similar"
    elif similarity > 0.4: level = "🟠 Slightly Similar"
    else: level = "🔴 Not Similar"
    return f"## {level}\n**Cosine Similarity: {similarity:.4f}**\n\n*384-dimensional embeddings*"

with gr.Blocks(title="Code Embeddings Demo — Matthieu.AI", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""# 🔍 bge-small-code-search-v1\n### Semantic Code Search Demo\nSearch code by describing what it does. Model by [Matthieu.AI](https://huggingface.co/Matthieufromparis)""")
    with gr.Tab("Code Search"):
        query_input = gr.Textbox(label="Search Query", placeholder="e.g., 'parse a JSON config file'", lines=2)
        top_k = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Results")
        search_button = gr.Button("Search Code", variant="primary")
        search_output = gr.Markdown(label="Results")
        search_button.click(fn=search_code, inputs=[query_input, top_k], outputs=search_output)
        gr.Examples(examples=[["read a JSON config file"], ["validate an email address"], ["connect to a database"], ["retry a function if it fails"]], inputs=[query_input])
    with gr.Tab("Compare Texts"):
        text1 = gr.Textbox(label="Text 1", placeholder="Enter natural language or code...", lines=3)
        text2 = gr.Textbox(label="Text 2", placeholder="Enter another text to compare...", lines=3)
        compare_button = gr.Button("Compare", variant="primary")
        compare_output = gr.Markdown(label="Similarity")
        compare_button.click(fn=compare_embeddings, inputs=[text1, text2], outputs=compare_output)

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