import streamlit as st
from transformers import pipeline, AutoModelForMaskedLM, AutoTokenizer
from cltk.data.fetch import FetchCorpus
import builtins
import os
import json
import requests
from git import Repo

GITHUB_RAW_URL = "https://raw.githubusercontent.com/Cicciokr/latin-ai-model/main/data.json"
GITHUB_REPO = "https://github.com/Cicciokr/latin-ai-model"
GITHUB_TOKEN = "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"
DATA_FILE = "data.json"

def upload_to_github():
    repo_dir = "temp_repo"
    if not os.path.exists(repo_dir):
        Repo.clone_from(GITHUB_REPO.replace("https://", f"https://{GITHUB_TOKEN}@"), repo_dir)
    
    repo = Repo(repo_dir)
    file_path = os.path.join(repo_dir, "data.json")

    with open(file_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=4)

    repo.git.add("data.json")
    repo.index.commit("Aggiornato data.json")
    
    # Imposta l'URL remoto con autenticazione per evitare problemi di credenziali
    origin = repo.remote(name='origin')
    origin.set_url(GITHUB_REPO.replace("https://", f"https://{GITHUB_TOKEN}@"))
    
    repo.remote().push()
    st.sidebar.success("File JSON aggiornato su GitHub!")

def load_data():
    """Carica i dati salvati (token e frasi) dal file JSON da GitHub."""
    try:
        response = requests.get(GITHUB_RAW_URL)
        if response.status_code == 200:
            return response.json()
    except requests.exceptions.RequestException:
        pass
    return {"tokens": [], "phrases": {}}

def save_data(data):
    """Salva i dati (token e frasi) nel file JSON."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    file_path = os.path.join(script_dir, DATA_FILE)
    
    with open(file_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=4)


data = load_data()

def save_token_and_phrase(token, phrase):
    if phrase not in data["phrases"]:
        data["phrases"][phrase] = token
        save_data(data)

def get_valid_predictions(sentence, max_attempts=3, top_k=5):
    """Verifica se la frase è già salvata e usa il token corrispondente."""
    data = load_data()
    print(data)
    print(sentence)
    print(sentence in data["phrases"])
    if sentence in data["phrases"]:
        return [{"token_str": data["phrases"][sentence], "score": 1.0, "sequence": sentence.replace("<mask>", data["phrases"][sentence])}]
    
    attempt = 0
    filtered_predictions = []
    while attempt < max_attempts:
        predictions = fill_mask_roberta(sentence, top_k=top_k)
        filtered_predictions = [
            pred for pred in predictions if pred["token_str"] not in punctuation_marks
        ]
        if filtered_predictions:
            break
        attempt += 1
    return filtered_predictions


# Imposta una chiave di sessione per il testo input, così possiamo aggiornarlo
if "input_text_value" not in st.session_state:
    st.session_state["input_text_value"] = "Lorem ipsum dolor sit amet, <mask> adipiscing elit."

# Frasi di esempio
examples = [
    "Asdrubal, frater Annibalis, qui secundo Punico bello <mask> ingentibus copiis ab Hispania veniens ...",
    "hanno et mago qui <mask> punico bello cornelium consulem aput liparas ceperunt ...",
    "Lorem ipsum dolor sit amet, <mask> adipiscing elit.",
    "Populus Romanus cum Macedonibus <mask> ter gessit",
    "Reliqui qui tum principes numerabantur in magistratibus <mask> cotidieque fere a nobis in contionibus audiebantur."
]

st.title("Completamento di parole in testi Latino Antico con Analisi Morfologica")

st.write("Esempi di testo (clicca sul bottone per copiare la frase nel campo di input):")

# Per ogni frase, creiamo una riga con la frase + bottone "Usa questa frase"
for i, example in enumerate(examples, start=1):
    cols = st.columns([4,1])  # la prima colonna più larga per il testo, la seconda più stretta per il bottone
    with cols[0]:
        st.write(f"Esempio {i}: {example}")
    with cols[1]:
        # Se il bottone viene premuto, aggiorna la session state
        if st.button(f"Usa {i}"):
            st.session_state["input_text_value"] = example
    

# UI per l'inserimento del token e delle frasi
st.sidebar.header("Gestione Token e Frasi")
token_input = st.sidebar.text_input("Se la predizione è errata, salva il token corretto:")
if st.sidebar.button("Salva Token e Frase"):
    if token_input:
        save_token_and_phrase(token_input, st.session_state.get("input_text_value"))
        upload_to_github()
        st.sidebar.success("Token e frase salvati con successo!")
    else:
        st.sidebar.warning("Inserisci sia un token che una frase validi.")

existing_phrases = data.get("phrases", {})
st.sidebar.subheader("Frasi salvate:")
st.sidebar.write("\n\n".join(existing_phrases.keys()) if existing_phrases else "Nessuna frase salvata.")

_original_input = builtins.input
def _always_yes(prompt=""):
    print(prompt, "Y")  # per far vedere a log che abbiamo risposto 'Y'
    return "Y"

builtins.input = _always_yes

corpus_downloader = FetchCorpus(language="lat")
corpus_downloader.import_corpus("lat_models_cltk")

try:
    from cltk import NLP
    nlp_lat = NLP(language="lat")
except ImportError:
    nlp_lat = None

tokenizer_roberta = AutoTokenizer.from_pretrained("Cicciokr/Roberta-Base-Latin-Uncased")
model_roberta = AutoModelForMaskedLM.from_pretrained("Cicciokr/Roberta-Base-Latin-Uncased")
fill_mask_roberta = pipeline("fill-mask", model=model_roberta, tokenizer=tokenizer_roberta)

punctuation_marks = {".", ",", ";", ":", "!", "?"}

input_text = st.text_area(
    label="Testo:",
    height=150,
    key="input_text_value"
)

if input_text:
    input_text_roberta = input_text.replace("[MASK]", "<mask>")
    predictions_roberta = get_valid_predictions(input_text_roberta)
    st.subheader("Risultati delle previsioni:")
    for pred in predictions_roberta:
        st.write(f" Token: {pred['token_str']}")
        st.write(f" Probabilità: {pred['score']:.4f}")
        st.write(f" Sequence: {pred['sequence']}")
        st.write("---")
    if nlp_lat is not None:
        st.subheader("Analisi Morfologica con CLTK")
        for pred in predictions_roberta:
            doc = nlp_lat(pred['token_str'])
            st.write(f"Frase: {pred['token_str']}")
            for w in doc.words:
                st.write(
                    f"- **Token**: {w.string}\n"
                    f"  - Lemma: {w.lemma}\n"
                    f"  - UPOS: {w.upos}\n"
                    f"  - Morph: {w.features}\n"
                )
            st.write("---")
    else:
        st.warning("CLTK non installato. Esegui 'pip install cltk' per abilitare l'analisi.")