import json
from sentence_transformers import SentenceTransformer, util

MAP_FILE = "column_name_map.json"

# Load static map
try:
    with open(MAP_FILE) as f:
        STATIC_MAP = json.load(f)
except FileNotFoundError:
    STATIC_MAP = {}

model = SentenceTransformer('all-MiniLM-L6-v2')

REFERENCE_NAMES = list(set(STATIC_MAP.values()) | {
    "customer_id", "product_name", "order_date", "region", "email", "phone_number",
    "total_amount", "quantity", "price", "shipping_address", "created_at", "updated_at"
})

reference_embeddings = model.encode(REFERENCE_NAMES, convert_to_tensor=True)

def suggest_column_name(raw_name: str):
    if raw_name in STATIC_MAP:
        return STATIC_MAP[raw_name], 1.0
    raw_embedding = model.encode(raw_name, convert_to_tensor=True)
    scores = util.cos_sim(raw_embedding, reference_embeddings)[0]
    best_idx = scores.argmax().item()
    return REFERENCE_NAMES[best_idx], scores[best_idx].item()

def update_column_name(raw: str, readable: str):
    STATIC_MAP[raw] = readable
    with open(MAP_FILE, "w") as f:
        json.dump(STATIC_MAP, f, indent=2)
    return True