import streamlit as st
import os
import re
from bs4 import BeautifulSoup
from PIL import Image
from VerifierApp.config import (
    INPUTS_DIR, MODELS_DIR, save_uploaded_file,
    get_files_in_dir, get_folders_in_dir,
    load_model_capability_overrides, save_model_capability_overrides,
)
from VerifierApp.core.models import fetch_available_models, get_model_capabilities
from VerifierApp.localisation_module import t, ICONS


def process_and_save_input_file(uploaded_file, target_dir):
    """
    Spracuje súbor pred uložením:
    - HTML oreže o balast a referencie a uloží ako čistý text.
    - Obrázky zmenší a prekonvertuje na optimalizovaný JPEG.
    - Ostatné súbory uloží bez zmeny.
    """
    file_name = uploaded_file.name
    ext = file_name.split('.')[-1].lower()

    if ext in ['html', 'htm']:
        content = uploaded_file.getvalue().decode('utf-8', errors='ignore')
        soup = BeautifulSoup(content, 'html.parser')

        for tag in soup(["script", "style", "nav", "footer", "header", "aside"]):
            tag.decompose()

        stop_words = ['references', 'bibliography', 'acknowledgements', 'author contributions', 'data availability',
                      'conflict of interest']
        for header in soup.find_all(['h1', 'h2', 'h3', 'h4']):
            if header.text and any(word in header.text.lower() for word in stop_words):
                for sibling in header.find_next_siblings():
                    sibling.decompose()
                header.decompose()

        text = soup.get_text(separator=' ', strip=True)
        clean_text = re.sub(r'\s{2,}', ' ', text)

        new_name = file_name.rsplit('.', 1)[0] + "_cleaned.txt"
        with open(os.path.join(target_dir, new_name), 'w', encoding='utf-8') as f:
            f.write(clean_text)

    elif ext in ['png', 'jpg', 'jpeg', 'webp']:
        try:
            img = Image.open(uploaded_file)
            img = img.convert('RGB')
            img.thumbnail((1024, 1024), Image.Resampling.LANCZOS)

            new_name = file_name.rsplit('.', 1)[0] + "_optimized.jpg"
            img.save(os.path.join(target_dir, new_name), format="JPEG", quality=85)
        except Exception as e:
            st.warning(f"Nepodarilo sa optimalizovať obrázok {file_name}: {e}")
            save_uploaded_file(uploaded_file, target_dir)

    else:
        save_uploaded_file(uploaded_file, target_dir)


def render_data_management():
    st.header(t("data_header"))

    tab_inputs, tab_models, tab_llm_caps = st.tabs([
        f"{ICONS['folder']} {t('data_tab_inputs')}",
        f"{ICONS['folder']} {t('data_tab_models')}",
        f":material/tune: {t('data_tab_llm_caps')}",
    ])

    with tab_inputs:
        col1, col2 = st.columns([1, 1])
        with col1:
            st.subheader(t("data_input_subheader"))
            with st.container(border=True):
                st.write(f"**{t('data_add_input_group')}**")
                input_group_name = st.text_input(t("data_input_folder_name"))
                uploaded_inputs = st.file_uploader(t("data_upload_files"), accept_multiple_files=True,
                                                   type=["pdf", "html", "txt", "csv", "jpeg", "jpg", "png", "webp"])

                if st.button(t("data_save_inputs"), icon=ICONS["save"], use_container_width=True):
                    if input_group_name and uploaded_inputs:
                        target_dir = os.path.join(INPUTS_DIR, input_group_name)
                        os.makedirs(target_dir, exist_ok=True)
                        for file in uploaded_inputs:
                            process_and_save_input_file(file, target_dir)
                        st.success(f"Dáta pre skupinu '{input_group_name}' boli úspešne orezané, optimalizované a uložené.")
                    else:
                        st.warning(t("data_input_warning"))

        with col2:
            existing_input_dirs = get_folders_in_dir(INPUTS_DIR)
            if existing_input_dirs:
                st.write(f"**{t('data_available_inputs')}**")
                for d in existing_input_dirs:
                    with st.expander(d):
                        files = get_files_in_dir(os.path.join(INPUTS_DIR, d))
                        for f in files:
                            st.text(f" - {f}")
            else:
                st.info(t("data_no_inputs"))

    with tab_models:
        col1, col2 = st.columns([1, 1])
        with col1:
            st.subheader(t("data_models_subheader"))
            with st.container(border=True):
                st.write(f"**{t('data_add_model')}**")
                model_group_name = st.text_input(t("data_model_folder_name"))
                uploaded_models = st.file_uploader(t("data_upload_model_files"), accept_multiple_files=True)

                if st.button(t("data_save_models"), icon=ICONS["save"], use_container_width=True):
                    if model_group_name and uploaded_models:
                        target_dir = os.path.join(MODELS_DIR, model_group_name)
                        os.makedirs(target_dir, exist_ok=True)
                        for file in uploaded_models:
                            save_uploaded_file(file, target_dir)
                        st.success(t("data_models_saved", name=model_group_name))
                    else:
                        st.warning(t("data_model_warning"))

        with col2:
            existing_model_dirs = get_folders_in_dir(MODELS_DIR)
            if existing_model_dirs:
                st.write(f"**{t('data_available_models')}**")
                for d in existing_model_dirs:
                    with st.expander(d):
                        files = get_files_in_dir(os.path.join(MODELS_DIR, d))
                        for f in files:
                            st.text(f" - {f}")
            else:
                st.info(t("data_no_models"))

    with tab_llm_caps:
        _render_llm_capabilities()


def _render_llm_capabilities():
    st.subheader(t("data_llm_caps_header"))
    st.caption(t("data_llm_caps_help"))

    available_models = fetch_available_models()
    if not available_models:
        st.warning(t("data_llm_caps_no_models"))
        return

    PDF_OPTIONS = [
        t("data_llm_pdf_auto"),
        t("data_llm_pdf_extract"),
        t("data_llm_pdf_native"),
    ]
    PDF_VALUES = ["auto", "extract", "native"]

    overrides = load_model_capability_overrides()
    for model_id in available_models:
        key = f"pdf_cap_{model_id}"
        if key not in st.session_state:
            current_pdf = overrides.get(model_id, {}).get("pdf", "auto")
            st.session_state[key] = PDF_OPTIONS[PDF_VALUES.index(current_pdf)]

    def _auto_save():
        new_overrides = {
            mid: {"pdf": PDF_VALUES[PDF_OPTIONS.index(st.session_state[f"pdf_cap_{mid}"])]}
            for mid in available_models
            if f"pdf_cap_{mid}" in st.session_state
            and PDF_VALUES[PDF_OPTIONS.index(st.session_state[f"pdf_cap_{mid}"])] != "auto"
        }
        save_model_capability_overrides(new_overrides)
        st.toast(t("data_llm_caps_saved"), icon=ICONS["success"])


    hcol1, hcol2, hcol3, hcol4 = st.columns([3, 1, 1, 2])
    hcol1.markdown(f"**{t('data_llm_col_model')}**")
    hcol2.markdown(f"**{t('data_llm_col_vision')}**")
    hcol3.markdown(f"**{t('data_llm_col_pdf_auto')}**")
    hcol4.markdown(f"**{t('data_llm_col_pdf_override')}**")

    for model_id in available_models:
        caps = get_model_capabilities(model_id)

        col1, col2, col3, col4 = st.columns([3, 1, 1, 2])
        col1.write(model_id)
        col2.write(ICONS["success"] if caps["vision"] else ICONS["incorrect"])
        col3.write(ICONS["success"] if caps["pdf"] else ICONS["incorrect"])

        col4.selectbox(
            label="##",
            options=PDF_OPTIONS,
            key=f"pdf_cap_{model_id}",
            label_visibility="collapsed",
            on_change=_auto_save,
        )
