import torch
import gradio as gr
import spacy
from navec import Navec
import pickle
import json

from model_loaders.model_loaders import load_segm_model, load_pos_model, load_gloss_model
from utils.basic_glossing import BasicGlossing
from utils.glossing_with_lemmas import GlossingWithLemmas
from utils.gloss_text import GlossText
from utils.preprocessing import clear_punctuation

from dicts.symb_vocab import symb_vocab
from dicts.char_vocab import char_vocab
from dicts.word_vocab import word_vocab
from dicts.pos_label_vocab import pos_label_vocab
from dicts.morpheme_vocab import morpheme_vocab
from dicts.gloss_vocab import gloss_vocab
from dicts.glosses_dictionary import glosses_dictionary
from utils.trie import *

DEVICE = 'cpu'

PATHS = {
    'segm_model': 'models/segm_model.pth',
    'pos_model': 'models/pos_model.pth',
    'gloss_model': 'models/gloss_model.pth',
    'navec': 'resources/navec_hudlit_v1_12B_500K_300d_100q.tar',
    'trie': 'dicts/trie.pkl',
    'stem_vocabulary': 'dicts/stem_vocabulary.json'
}

def load_json(path):
    with open(path, 'r', encoding='utf8') as f:
        return json.load(f)
    
def load_pickle(path):
    with open(path, 'rb') as f:
        return pickle.load(f)

print("Загрузка ресурсов...")
trie = load_pickle(PATHS['trie'])
segm_model = load_segm_model(PATHS['segm_model'], device=DEVICE)
pos_model = load_pos_model(PATHS['pos_model'], device=DEVICE)
gloss_model = load_gloss_model(PATHS['gloss_model'], device=DEVICE)
stem_vocabulary = load_json(PATHS['stem_vocabulary'])

navec = Navec.load(PATHS['navec'])
spacy_lemmatizer = spacy.load("ru_core_news_sm")

print("Инициализация пайплайна...")
gloss_text = GlossText(segm_model, pos_model, gloss_model,
                       symb_vocab, char_vocab, word_vocab, pos_label_vocab,
                       morpheme_vocab, gloss_vocab, glosses_dictionary,
                       stem_vocabulary, trie, navec, spacy_lemmatizer,
                       device=DEVICE)


def predict(text, translation=None, num_sentence=None):
    if not text.strip():
        return "Введите текст"
    
    sent = clear_punctuation(text.strip())
    translation = translation.strip() if translation else None
    
    results = gloss_text.gloss_sent(sent, translation, num_sentence)
    return results


iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(lines=5, label="Нивхское предложение"),
        gr.Textbox(lines=5, label="Русский перевод (необязательно)"),
        gr.Textbox(lines=1, label="Номер предложения (необязательно)")
    ],
    outputs=[gr.Textbox(label="Глоссированное предложение", lines=3)],
    title="Глоссирование нивхского текста",
    description="Введите нивхский текст и, при наличии, его перевод"
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)