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)


DEFAULT_CLASSES_FILE = "classes.json"  # Файл за замовчуванням
DEFAULT_SIGNATURES_FILE = "signatures.npz"  # Файл для збереження signatures


##############################################################################
# 1. Функції для роботи з класами та signatures
##############################################################################
def load_classes(json_path: str) -> 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 reload_classes_and_signatures(json_path: str, model_name: str, force_rebuild: bool) -> str:
    """
    Перезавантажує класи з нового JSON файлу та оновлює signatures
    """
    global classes_json, class_signatures
    
    try:
        new_classes = load_classes(json_path)
        if not new_classes:
            return "Помилка: Файл класів порожній або має неправильний формат"
        
        classes_json = new_classes
        print(f"Завантажено {len(classes_json)} класів з {json_path}")
        
        # Зберігаємо новий файл класів як файл за замовчуванням
        try:
            with open(DEFAULT_CLASSES_FILE, 'w', encoding='utf-8') as f:
                json.dump(classes_json, f, ensure_ascii=False, indent=2)
            print(f"Збережено як {DEFAULT_CLASSES_FILE}")
        except Exception as e:
            print(f"Помилка при збереженні файлу класів: {str(e)}")
        
        result = initialize_signatures(
            model_name=model_name,
            signatures_file=DEFAULT_SIGNATURES_FILE,
            force_rebuild=force_rebuild
        )
        return f"Класи оновлено. {result}"
        
    except Exception as e:
        return f"Помилка при оновленні класів: {str(e)}"


def initialize_signatures(model_name: str = "text-embedding-3-large", 
                        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 = {}
df = None
embeddings = None
class_signatures = None
embeddings_mean = None
embeddings_std = None

##############################################################################
# 2. Функції для роботи з даними та класифікації
##############################################################################
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)
    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-large") -> 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:
    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 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!"}
        
    emb = get_openai_embedding(text)
    
    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
        
    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)
        
        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)

##############################################################################
# 3. Головний інтерфейс
##############################################################################
def main():
    # Ініціалізуємо класи та signatures при запуску
    print("Завантаження класів...")
    
    # Спроба завантажити класи з файлу за замовчуванням
    global classes_json
    if os.path.exists(DEFAULT_CLASSES_FILE):
        classes_json = load_classes(DEFAULT_CLASSES_FILE)
        if not classes_json:
            print(f"ПОПЕРЕДЖЕННЯ: Файл {DEFAULT_CLASSES_FILE} порожній або має неправильний формат")
    else:
        print(f"ПОПЕРЕДЖЕННЯ: Файл {DEFAULT_CLASSES_FILE} не знайдено")
        classes_json = {}
    
    # Якщо є класи, ініціалізуємо signatures
    if classes_json:
        print("Ініціалізація signatures...")
        try:
            init_message = initialize_signatures(
                signatures_file=DEFAULT_SIGNATURES_FILE
            )
            print(f"Результат ініціалізації: {init_message}")
        except Exception as e:
            print(f"ПОПЕРЕДЖЕННЯ: Помилка при ініціалізації signatures: {str(e)}")
    else:
        print("Очікую завантаження класів через інтерфейс...")

    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-large",
                            label="OpenAI model"
                        )
                        json_file = gr.File(
                            label="Завантажити новий JSON з класами",
                            file_types=[".json"]
                        )
                        force_rebuild = gr.Checkbox(
                            label="Примусово перебудувати signatures",
                            value=False
                        )
                    with gr.Row():
                        build_btn = gr.Button("Оновити signatures")
                        build_out = gr.Label(label="Статус signatures")

                # Оновлений обробник для перебудови signatures
                def update_with_file(file, model_name, force):
                    if file is None:
                        return "Виберіть файл з класами"
                    try:
                        temp_path = file.name
                        return reload_classes_and_signatures(temp_path, model_name, force)
                    except Exception as e:
                        return f"Помилка при оновленні: {str(e)}"

                single_process_btn.click(
                    fn=process_single_text,
                    inputs=[text_input, threshold_slider],
                    outputs=result_text
                )
                
                build_btn.click(
                    fn=update_with_file,
                    inputs=[json_file, model_choice, force_rebuild],
                    outputs=build_out
                )
            
            # Вкладка 2: Batch Processing
            with gr.TabItem("Пакетна обробка"):
                gr.Markdown("## 1) Завантаження даних")
                with gr.Row():
                    csv_input = gr.Textbox(
                        value="messages.csv",
                        label="CSV-файл"
                    )
                    emb_input = gr.Textbox(
                        value="embeddings.npy",
                        label="Numpy Embeddings"
                    )
                    load_btn = gr.Button("Завантажити дані")
                
                load_output = gr.Label(label="Результат завантаження")
                
                gr.Markdown("## 2) Класифікація")
                with gr.Row():
                    filter_in = gr.Textbox(label="Фільтр (опціонально)")
                    batch_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.3,
                        step=0.05,
                        label="Поріг впевненості"
                    )
                
                classify_btn = gr.Button("Класифікувати")
                classify_out = gr.Dataframe(
                    label="Результат (Message / Target / Scores)"
                )
                
                gr.Markdown("## 3) Зберегти результати")
                save_btn = gr.Button("Зберегти розмічені дані")
                save_out = gr.Label()
                
                # Підключаємо обробники подій
                load_btn.click(
                    fn=load_data,
                    inputs=[csv_input, emb_input],
                    outputs=load_output
                )
                
                classify_btn.click(
                    fn=classify_rows,
                    inputs=[filter_in, batch_threshold],
                    outputs=classify_out
                )
                
                save_btn.click(
                    fn=lambda: df.to_csv("messages_with_labels.csv", index=False) if df is not None else "Дані відсутні!",
                    inputs=[],
                    outputs=save_out
                )

        gr.Markdown("""
        ### Інструкція:
        1. У вкладці "Налаштування моделі" можна:
           - Завантажити новий JSON файл з класами
           - Вибрати модель для embeddings
           - Примусово перебудувати signatures
        2. Після зміни класів натисніть "Оновити signatures"
        3. Використовуйте повзунок "Поріг впевненості" для фільтрації результатів
        4. На вкладці "Пакетна обробка" можна аналізувати багато повідомлень
        5. Результати можна зберегти в CSV файл
        """)

    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)

if __name__ == "__main__":
    main()