import re
from sklearn.metrics.pairwise import cosine_similarity
from embeddings_utils import load_model
import torch

# Load the SentenceTransformer model directly in this file
device = "cuda" if torch.cuda.is_available() else "cpu"
model = load_model(device)

def find_top_question(query, metadata, embeddings):
    query_embedding = model.encode(query, convert_to_tensor=True).cpu().numpy().reshape(1, -1)
    similarities = cosine_similarity(query_embedding, embeddings).flatten()
    top_index = similarities.argsort()[-1]
    top_result = metadata.iloc[top_index].copy()
    top_result['similarity_score'] = similarities[top_index]
    return top_result

def generate_detailed_prompt(question_metadata):
    return (
        f"Transform this LeetCode question into a real-world interview scenario.\n\n"
        f"**Company**: {question_metadata['company']}\n"
        f"**Question Name**: {question_metadata['questionName']}\n"
        f"**Difficulty Level**: {question_metadata['difficulty level']}\n"
        f"**Tags**: {question_metadata['Tags']}\n"
        f"**Content**: {question_metadata['Content']}\n"
        f"\nPlease create a real-world interview question based on this information. "
        f"Include sections for problem description, code template, sample input, and expected output."
    )
def extract_code_and_test_case(response):
    import re
    
    code_template = ""
    sample_test_case = ""
    expected_output = ""

    # Extract Python code block
    code_match = re.search(r'```python(.*?)```', response, re.DOTALL)
    if code_match:
        code_template = code_match.group(1).strip()
    
    # Extract test case and expected output
    test_case_match = re.search(r'Sample Input:\s*(.*?)\n', response, re.DOTALL)
    expected_output_match = re.search(r'Expected Output:\s*(.*?)\n', response, re.DOTALL)

    if test_case_match:
        sample_test_case = test_case_match.group(1).strip()
    
    if expected_output_match:
        expected_output = expected_output_match.group(1).strip()

    return code_template, sample_test_case, expected_output