"""Streamlit app for Presidio with multilingual punctuation model."""
import logging
import sys
import os
import traceback

import dotenv
import pandas as pd
import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification
from annotated_text import annotated_text
from streamlit_tags import st_tags

# Nastavení logování
logging.basicConfig(level=logging.DEBUG, stream=sys.stdout)
logger = logging.getLogger("presidio-streamlit")

# Načtení proměnných prostředí
dotenv.load_dotenv()

# Konfigurace stránky
st.set_page_config(
    page_title="Presidio demo s vícejazyčnou interpunkcí",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={
        "About": "https://microsoft.github.io/presidio/",
    },
)

# Sidebar
st.sidebar.header(
    """
Analýza textu s [Microsoft Presidio](https://microsoft.github.io/presidio/) a vícejazyčným modelem interpunkce
"""
)

# Výběr modelu
model_name = "oliverguhr/fullstop-punctuation-multilang-large"
st.sidebar.info(f"Používaný model: {model_name}")

try:
    # Kontrola, zda je CUDA dostupná
    device = "cuda" if torch.cuda.is_available() else "cpu"
    st.sidebar.info(f"Používané zařízení: {device}")

    # Načtení tokenizeru a modelu
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForTokenClassification.from_pretrained(model_name).to(device)

    st.sidebar.success("Model úspěšně načten!")
except Exception as e:
    st.sidebar.error(f"Chyba při načítání modelu: {str(e)}")
    logger.error(f"Chyba při načítání modelu: {str(e)}")
    st.stop()

# Funkce pro zpracování textu
def process_text(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    
    predicted_ids = outputs.logits.argmax(dim=-1)[0]
    predicted_tokens = [model.config.id2label[int(id)] for id in predicted_ids]
    
    words = text.split()
    result = []
    for word, pred in zip(words, predicted_tokens):
        if pred != "0":
            result.append(word + pred)
        else:
            result.append(word)
    
    return " ".join(result)

# Hlavní panel
st.title("Analýza textu s vícejazyčným modelem interpunkce")

# Vstupní text
input_text = st.text_area("Zadejte text k analýze:", height=200)

if st.button("Analyzovat"):
    if input_text:
        try:
            processed_text = process_text(input_text)
            st.subheader("Výsledek analýzy:")
            st.write(processed_text)
        except Exception as e:
            st.error(f"Chyba při zpracování textu: {str(e)}")
            logger.error(f"Chyba při zpracování textu: {str(e)}")
    else:
        st.warning("Zadejte prosím nějaký text k analýze.")

# Informace o aplikaci
st.markdown("---")
st.info(
    """
    Tato aplikace používá vícejazyčný model pro obnovu interpunkce v textu.
    Model podporuje angličtinu, němčinu, francouzštinu a italštinu.
    
    Více informací o modelu najdete na [Hugging Face](https://huggingface.co/oliverguhr/fullstop-punctuation-multilang-large).
    """
)

# Footer
st.markdown(
    "[![GitHub](https://img.shields.io/badge/GitHub-Presidio-blue?style=flat&logo=github)](https://github.com/microsoft/presidio)"
    "[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)"
)