import os
import gradio as gr
import pandas as pd
import numpy as np
import json
from typing import Dict, List

from openai import OpenAI
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# OpenAI setup
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=OPENAI_API_KEY)

##############################################################################
# 1. Функції для роботи з класами та signatures
##############################################################################
classes_json_load = "classes_short.json"

def load_classes(json_path: str = classes_json_load) -> dict:
    """
    Завантаження класів та їх хінтів з JSON файлу
    """
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            classes = json.load(f)
        return classes
    except FileNotFoundError:
        print(f"Файл {json_path} не знайдено! Використовуємо пустий словник класів.")
        return {}
    except json.JSONDecodeError:
        print(f"Помилка читання JSON з файлу {json_path}! Використовуємо пустий словник класів.")
        return {}

def save_signatures(signatures: Dict[str, np.ndarray], filename: str = "signatures.npz") -> None:
    """
    Зберігає signatures у NPZ файл
    """
    if signatures:
        np.savez(filename, **signatures)

def load_signatures(filename: str = "signatures.npz") -> Dict[str, np.ndarray]:
    """
    Завантажує signatures з NPZ файлу
    """
    try:
        with np.load(filename) as data:
            return {key: data[key] for key in data.files}
    except (FileNotFoundError, IOError):
        return None

def initialize_signatures(model_name: str = "text-embedding-3-small", 
                        signatures_file: str = "signatures.npz",
                        force_rebuild: bool = False) -> str:
    """
    Ініціалізує signatures: завантажує існуючі або створює нові
    """
    global class_signatures, classes_json
    
    if not classes_json:
        return "Помилка: Не знайдено жодного класу в classes.json"
    
    print(f"Знайдено {len(classes_json)} класів")
    
    # Спробуємо завантажити існуючі signatures
    if not force_rebuild and os.path.exists(signatures_file):
        try:
            loaded_signatures = load_signatures(signatures_file)
            # Перевіряємо, чи всі класи з classes_json є в signatures
            if loaded_signatures and all(cls in loaded_signatures for cls in classes_json):
                class_signatures = loaded_signatures
                print("Успішно завантажено збережені signatures")
                return f"Завантажено існуючі signatures для {len(class_signatures)} класів"
        except Exception as e:
            print(f"Помилка при завантаженні signatures: {str(e)}")
    
    # Якщо немає файлу або примусова перебудова - створюємо нові
    try:
        class_signatures = {}
        total_classes = len(classes_json)
        print(f"Починаємо створення нових signatures для {total_classes} класів...")
        
        for idx, (cls_name, hints) in enumerate(classes_json.items(), 1):
            if not hints:
                print(f"Пропускаємо клас {cls_name} - немає хінтів")
                continue
                
            print(f"Обробка класу {cls_name} ({idx}/{total_classes})...")
            try:
                arr = embed_hints(hints, model_name=model_name)
                class_signatures[cls_name] = arr.mean(axis=0)
                print(f"Успішно створено signature для {cls_name}")
            except Exception as e:
                print(f"Помилка при створенні signature для {cls_name}: {str(e)}")
                continue
        
        if not class_signatures:
            return "Помилка: Не вдалося створити жодного signature"
        
        # Зберігаємо нові signatures
        try:
            save_signatures(class_signatures, signatures_file)
            print("Signatures збережено у файл")
        except Exception as e:
            print(f"Помилка при збереженні signatures: {str(e)}")
        
        return f"Створено та збережено нові signatures для {len(class_signatures)} класів"
    except Exception as e:
        return f"Помилка при створенні signatures: {str(e)}"


# Замість хардкоду classes_json тепер використовуємо:
classes_json = load_classes()

##############################################################################
# 2. Глобальні змінні
##############################################################################
df = None
embeddings = None
class_signatures = None
embeddings_mean = None  # Для нормалізації single text
embeddings_std = None   # Для нормалізації single text

##############################################################################
# 3. Функції для роботи з даними та класифікації
##############################################################################
def load_data(csv_path: str = "messages.csv", emb_path: str = "embeddings.npy"):
    global df, embeddings, embeddings_mean, embeddings_std
    df_local = pd.read_csv(csv_path)
    emb_local = np.load(emb_path)
    assert len(df_local) == len(emb_local), "CSV і embeddings різної довжини!"

    df_local["Target"] = "Unlabeled"
    
    # Зберігаємо параметри нормалізації
    embeddings_mean = emb_local.mean(axis=0)
    embeddings_std = emb_local.std(axis=0)
    
    # Нормалізація embeddings
    emb_local = (emb_local - embeddings_mean) / embeddings_std

    df = df_local
    embeddings = emb_local
    
    return f"Завантажено {len(df)} рядків"

def get_openai_embedding(text: str, model_name: str = "text-embedding-3-small") -> list:
    response = client.embeddings.create(
        input=text,
        model=model_name
    )
    return response.data[0].embedding

def embed_hints(hint_list: List[str], model_name: str) -> np.ndarray:
    """
    Отримує embeddings для списку хінтів з виводом прогресу
    """
    emb_list = []
    total_hints = len(hint_list)
    
    for idx, hint in enumerate(hint_list, 1):
        try:
            print(f"  Отримання embedding {idx}/{total_hints}: '{hint}'")
            emb = get_openai_embedding(hint, model_name=model_name)
            emb_list.append(emb)
        except Exception as e:
            print(f"  Помилка при отриманні embedding для '{hint}': {str(e)}")
            continue
    
    if not emb_list:
        raise ValueError("Не вдалося отримати жодного embedding")
    
    return np.array(emb_list, dtype=np.float32)

def build_class_signatures(model_name: str):
    global class_signatures
    signatures = {}
    for cls_name, hints in classes_json.items():
        if not hints:
            continue
        arr = embed_hints(hints, model_name=model_name)
        signatures[cls_name] = arr.mean(axis=0)
    class_signatures = signatures
    return "Signatures побудовано!"

def predict_classes(text_embedding: np.ndarray, 
                   signatures: Dict[str, np.ndarray],
                   threshold: float = 0.0) -> Dict[str, float]:
    """
    Повертає словник класів та їх scores для одного тексту.
    Scores - це значення dot product між embedding тексту та signature класу
    """
    results = {}
    for cls, sign in signatures.items():
        score = float(np.dot(text_embedding, sign))
        if score > threshold:
            results[cls] = score
            
    # Сортуємо за спаданням score
    results = dict(sorted(results.items(), 
                         key=lambda x: x[1], 
                         reverse=True))
    
    return results

def process_single_text(text: str, threshold: float = 0.3) -> dict:
    """
    Обробка одного тексту
    """
    if class_signatures is None:
        return {"error": "Спочатку збудуйте signatures!"}
        
    # Отримуємо embedding для тексту
    emb = get_openai_embedding(text)
    
    # Нормалізуємо embedding використовуючи збережені параметри
    if embeddings_mean is not None and embeddings_std is not None:
        emb = (emb - embeddings_mean) / embeddings_std
    
    # Отримуємо передбачення
    predictions = predict_classes(emb, class_signatures, threshold)
    
    # Форматуємо результат
    if not predictions:
        return {"message": text, "result": "Жодного класу не знайдено"}
    
    formatted_results = []
    for cls, score in sorted(predictions.items(), key=lambda x: x[1], reverse=True):
        formatted_results.append(f"{cls}: {score:.2%}")
    
    return {
        "message": text,
        "result": "\n".join(formatted_results)
    }

def classify_rows(filter_substring: str = "", threshold: float = 0.3):
    """
    Класифікація з множинними мітками
    """
    global df, embeddings, class_signatures
    
    if class_signatures is None:
        return "Спочатку збудуйте signatures!"

    if df is None or embeddings is None:
        return "Дані не завантажені! Спочатку викличте load_data."

    if filter_substring:
        filtered_idx = df[df["Message"].str.contains(filter_substring, 
                                                   case=False, 
                                                   na=False)].index
    else:
        filtered_idx = df.index
        
    # Додаємо колонки для кожного класу зі scores
    for cls in class_signatures.keys():
        df[f"Score_{cls}"] = 0.0
        
    for i in filtered_idx:
        emb_vec = embeddings[i]
        predictions = predict_classes(emb_vec, 
                                   class_signatures, 
                                   threshold=threshold)
        
        # Записуємо scores для кожного класу
        for cls, score in predictions.items():
            df.at[i, f"Score_{cls}"] = score
            
        # Визначаємо основні класи (можна встановити поріг)
        main_classes = [cls for cls, score in predictions.items() 
                       if score > threshold]
        df.at[i, "Target"] = "|".join(main_classes) if main_classes else "None"

    result_columns = ["Message", "Target"] + [f"Score_{cls}" 
                                            for cls in class_signatures.keys()]
    result_df = df.loc[filtered_idx, result_columns].copy()
    return result_df.reset_index(drop=True)

##############################################################################
# 4. Функції для UI
##############################################################################
def ui_load_data(csv_path, emb_path):
    msg = load_data(csv_path, emb_path)
    return f"{msg}"

def ui_build_signatures(model_name):
    msg = build_class_signatures(model_name)
    return msg

def ui_save_data():
    global df
    if df is None:
        return "Дані відсутні!"
    df.to_csv("messages_with_labels.csv", index=False)
    return "Файл 'messages_with_labels.csv' збережено!"

##############################################################################
# 5. Головний інтерфейс
##############################################################################
def main():
    # Ініціалізуємо класи та signatures при запуску
    print("Завантаження класів...")
    if not classes_json:
        print("КРИТИЧНА ПОМИЛКА: Не вдалося завантажити класи!")
        return
    
    print("Ініціалізація signatures...")
    try:
        init_message = initialize_signatures()
        print(f"Результат ініціалізації: {init_message}")
        if "Помилка" in init_message:
            print("ПОПЕРЕДЖЕННЯ: Проблеми з ініціалізацією signatures")
    except Exception as e:
        print(f"КРИТИЧНА ПОМИЛКА при ініціалізації signatures: {str(e)}")
        return
    
    with gr.Blocks() as demo:
        gr.Markdown("# SDC Classifier з Gradio")
        
        with gr.Tabs():
            # Вкладка 1: Single Text Testing
            with gr.TabItem("Тестування одного тексту"):
                with gr.Row():
                    with gr.Column():
                        text_input = gr.Textbox(
                            label="Введіть текст для аналізу",
                            lines=5,
                            placeholder="Введіть текст..."
                        )
                        threshold_slider = gr.Slider(
                            minimum=0.0,
                            maximum=1.0,
                            value=0.3,
                            step=0.05,
                            label="Поріг впевненості"
                        )
                        single_process_btn = gr.Button("Проаналізувати")
                    
                    with gr.Column():
                        result_text = gr.JSON(
                            label="Результати аналізу"
                        )
                
                # Модифікована панель налаштувань
                with gr.Accordion("Налаштування моделі", open=False):
                    with gr.Row():
                        model_choice = gr.Dropdown(
                            choices=["text-embedding-3-large","text-embedding-3-small"],
                            value="text-embedding-3-small",
                            label="OpenAI model"
                        )
                        force_rebuild = gr.Checkbox(
                            label="Примусово перебудувати signatures",
                            value=False
                        )
                        build_btn = gr.Button("Оновити signatures")
                        build_out = gr.Label(label="Статус signatures")
                
                # Оновлений обробник для перебудови signatures
                def rebuild_signatures(model_name, force):
                    return initialize_signatures(model_name, force_rebuild=force)
                
                single_process_btn.click(
                    fn=process_single_text,
                    inputs=[text_input, threshold_slider],
                    outputs=result_text
                )
                
                build_btn.click(
                    fn=rebuild_signatures,
                    inputs=[model_choice, force_rebuild],
                    outputs=build_out
                )
            
            # Вкладка 2: Batch Processing [залишається без змін]
            
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)

if __name__ == "__main__":
    main()