app.py

import os
import io
import json
import uuid
import re
from datetime import datetime, timezone
from typing import Optional, Tuple, List, Any

import numpy as np
import pandas as pd
import streamlit as st
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
from huggingface_hub import login, HfApi


# ==========================
# НАСТРОЙКИ ПРИЛОЖЕНИЯ
# ==========================

st.set_page_config(page_title="Поиск постдока", layout="wide")

HF_TOKEN = os.getenv("HF_TOKEN")
HF_MERGED_REPO = os.getenv("HF_MERGED_REPO")
HF_EMB_REPO = os.getenv("HF_EMB_REPO")
MODEL_NAME = os.getenv("MODEL_NAME")

HF_REQUESTS_REPO = os.getenv("HF_REQUESTS_REPO")
HF_REQUESTS_REPO_TYPE = os.getenv("HF_REQUESTS_REPO_TYPE")
HF_WRITE_TOKEN = os.getenv("HF_WRITE_TOKEN")

OA_ENRICH_REPO = os.getenv("OA_ENRICH_REPO")

SLIDER_MIN_YEAR = 2005


COLUMN_LABELS_RU_EXCEL = {
    "№": "№",
    "score": "Сходство",
    "fio": "ФИО",
    "title": "Название диссертации",
    "author_org_short": "Организация",
    "dissertation_type": "Тип",
    "protection_year": "Год",
    "registration_number": "Регистрационный номер",
    "vak_link": "Ссылка ВАК",
    "openalex_url": "OpenAlex",
    "orcid_url": "ORCID",
    "h_index": "h-index",
    "i10_index": "i10-index",
    "works_count": "Работ",
    "cited_by_count": "Цитат",
}

DISPLAY_COLUMNS_ALL = [
    "№",
    "score",
    "fio",
    "title",
    "author_org_short",
    "dissertation_type",
    "protection_year",
    "registration_number",
]


# ==========================
# UI: КОЛОНКИ ТАБЛИЦЫ + НАСТРОЙКИ ОТОБРАЖЕНИЯ
# ВАЖНО:
# - ФИО НЕ ссылка
# - Название диссертации -> ссылка на ВАК (если есть vak_link)
# - после "Год": OpenAlex (ID+ссылка), ORCID (ID+ссылка)
# - Тип по умолчанию скрыт
# - Регистрационный номер по умолчанию скрыт
# ==========================

UI_TABLE_COLUMNS = [
    "Сходство",
    "ФИО",
    "Название диссертации",
    "Организация",
    "Тип",
    "Год",
    "OpenAlex",
    "ORCID",
    "Регистрационный номер",
    "h-index",
    "i10-index",
    "Работ",
    "Цитат",
]

# По умолчанию: все включены, кроме "Организация", "Тип" и "Регистрационный номер"
DEFAULT_VISIBLE_UI = {
    c: (c not in {"Организация", "Тип", "Регистрационный номер"})
    for c in UI_TABLE_COLUMNS
}


# ==========================
# ПРОВЕРКА СЕКРЕТОВ
# ==========================

if HF_TOKEN is None:
    st.error(
        "Не найден секрет `HF_TOKEN`. "
        "Задайте его в Settings → Variables and secrets вашего Space."
    )
    st.stop()

try:
    login(token=HF_TOKEN)
except Exception:
    pass


# ==========================
# CSS: уменьшение шрифта таблицы до ~80%
# ==========================

st.markdown(
    """
    <style>
    div[data-testid="stDataFrame"] { font-size: 80% !important; }
    div[data-testid="stDataFrame"] * { font-size: 80% !important; }
    .stDataFrame { font-size: 80% !important; }
    .gdg-w, .gdg-canvas { font-size: 80% !important; }

    div[data-testid="stDataEditor"] { font-size: 80% !important; }
    div[data-testid="stDataEditor"] * { font-size: 80% !important; }
    </style>
    """,
    unsafe_allow_html=True,
)


# ==========================
# ВСПОМОГАТЕЛЬНЫЕ ФУНКЦИИ
# ==========================

def _norm_regnum(x: Any) -> str:
    s = "" if x is None else str(x).strip()
    if s.endswith(".0") and s[:-2].isdigit():
        s = s[:-2]
    return s


def _safe_text(x: Any) -> str:
    if x is None:
        return ""
    if isinstance(x, float) and np.isnan(x):
        return ""
    s = str(x).strip()
    if s.lower() in {"none", "nan", "<na>"}:
        return ""
    return s


def _safe_fragment(s: Any) -> str:
    t = _safe_text(s).replace("#", " ").replace("\n", " ").replace("\r", " ").strip()
    return t


def _norm_openalex_url(x: Any) -> str:
    s = _safe_text(x)
    if not s:
        return ""
    if s.startswith("http://") or s.startswith("https://"):
        return s
    m = re.search(r"(A\d+)", s)
    return f"https://openalex.org/{m.group(1)}" if m else ""


def _norm_orcid_url(x: Any) -> str:
    s = _safe_text(x)
    if not s:
        return ""
    if s.startswith("http://") or s.startswith("https://"):
        return s
    m = re.search(r"(0000-[0-9X]{4}-[0-9X]{4}-[0-9X]{4})", s)
    return f"https://orcid.org/{m.group(1)}" if m else ""


def _openalex_id_from_url(url: str) -> str:
    s = _safe_text(url)
    m = re.search(r"\b(A\d+)\b", s)
    return m.group(1) if m else ""


def _orcid_id_from_url(url: str) -> str:
    s = _safe_text(url)
    m = re.search(r"\b(0000-[0-9X]{4}-[0-9X]{4}-[0-9X]{4})\b", s)
    return m.group(1) if m else ""


def _keyify(label: str) -> str:
    return "k_" + "".join(ch if ch.isalnum() else "_" for ch in label).strip("_")


def _show_col_key(col: str) -> str:
    return _keyify("show_" + col)


# ==========================
# ФИЛЬТРЫ (UI)
# ==========================

SCIENCE_LABELS = [
    "Архитектура",
    "Биологические",
    "Ветеринарные",
    "Географические",
    "Геолого-минералогические",
    "Искусствоведение",
    "Исторические",
    "Культурология",
    "Медицинские",
    "Педагогические",
    "Политические",
    "Сельскохозяйственные",
    "Технические",
    "Фармацевтические",
    "Физико-математические",
    "Филологические",
    "Философские",
    "Химические",
    "Экономические",
    "Юридические науки",
]
SCIENCE_LABELS = sorted(list(dict.fromkeys(SCIENCE_LABELS)), key=lambda s: s.casefold())
DEFAULT_SCIENCES = {"Технические", "Физико-математические", "Химические", "Биологические"}

SCIENCE_PATTERNS = {
    "Архитектура": ["архитектур"],
    "Биологические": ["биолог"],
    "Ветеринарные": ["ветеринар"],
    "Географические": ["географ"],
    "Геолого-минералогические": ["геолого-минералог", "геол.-минералог", "геолого минералог"],
    "Искусствоведение": ["искусствовед"],
    "Исторические": ["историч"],
    "Культурология": ["культуролог"],
    "Медицинские": ["медицин"],
    "Педагогические": ["педагог"],
    "Политические": ["политич"],
    "Сельскохозяйственные": ["сельскохозяй"],
    "Технические": ["технич"],
    "Фармацевтические": ["фармацевт"],
    "Физико-математические": ["физико-математ", "физ-мат", "физмат"],
    "Филологические": ["филолог"],
    "Философские": ["философ"],
    "Химические": ["химич"],
    "Экономические": ["экономич"],
    "Юридические науки": ["юридич"],
}


# ==========================
# ЗАПИСЬ ЗАЯВОК В HF REPO
# ==========================

def save_request_to_hub(payload: dict) -> str:
    token = HF_WRITE_TOKEN or HF_TOKEN
    if not token:
        raise RuntimeError("Не задан HF_WRITE_TOKEN (и нет HF_TOKEN).")

    api = HfApi(token=token)

    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    rid = uuid.uuid4().hex[:10]
    path_in_repo = f"requests/{ts}_{rid}.json"

    data = json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8")

    api.upload_file(
        path_or_fileobj=io.BytesIO(data),
        path_in_repo=path_in_repo,
        repo_id=HF_REQUESTS_REPO,
        repo_type=HF_REQUESTS_REPO_TYPE,
        commit_message=f"New PostDoc request {ts}",
    )
    return path_in_repo


# ==========================
# OA ENRICHMENT (КОМПАКТНЫЙ)
# ==========================

@st.cache_data(show_spinner="Загрузка OpenAlex/ORCID обогащения...")
def load_oa_enrichment() -> pd.DataFrame:
    ds = load_dataset(OA_ENRICH_REPO, split="train")
    df = ds.to_pandas()

    if "registration_number" not in df.columns:
        return pd.DataFrame().set_index(pd.Index([], name="reg_norm"))

    df["reg_norm"] = df["registration_number"].apply(_norm_regnum)

    if "openalex_url" not in df.columns and "openalex" in df.columns:
        df["openalex_url"] = df["openalex"]
    if "orcid_url" not in df.columns and "orcid" in df.columns:
        df["orcid_url"] = df["orcid"]

    if "openalex_url" in df.columns:
        df["openalex_url"] = df["openalex_url"].map(_norm_openalex_url)
    if "orcid_url" in df.columns:
        df["orcid_url"] = df["orcid_url"].map(_norm_orcid_url)

    for c in ["h_index", "i10_index", "works_count", "cited_by_count"]:
        if c in df.columns:
            df[c] = pd.to_numeric(df[c], errors="coerce")

    keep = ["reg_norm", "openalex_url", "orcid_url", "h_index", "i10_index", "works_count", "cited_by_count"]
    keep = [c for c in keep if c in df.columns]
    df = df[keep].copy()

    df = df.drop_duplicates(subset=["reg_norm"]).set_index("reg_norm", drop=True)
    return df


oa_enrich = load_oa_enrichment()


# ==========================
# ДАННЫЕ И МОДЕЛЬ
# ==========================

@st.cache_data(show_spinner="Загрузка данных...")
def load_data():
    ds_meta = load_dataset(HF_MERGED_REPO, split="train")
    df_meta = ds_meta.to_pandas()
    df_meta["registration_number"] = df_meta["registration_number"].astype(str).map(_norm_regnum)
    df_meta = df_meta.set_index("registration_number", drop=False)

    ds_emb = load_dataset(HF_EMB_REPO, split="train")
    df_emb = ds_emb.to_pandas()
    df_emb["registration_number"] = df_emb["registration_number"].astype(str).map(_norm_regnum)

    reg_nums = df_emb["registration_number"].tolist()

    emb_matrix = np.vstack(df_emb["embedding"].to_list()).astype("float32")
    norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
    emb_matrix = emb_matrix / np.maximum(norms, 1e-8)

    meta_aligned = df_meta.reindex(reg_nums)

    type_s = meta_aligned.get("dissertation_type", pd.Series([""] * len(reg_nums))).fillna("").astype(str)
    is_candidate = type_s.str.contains("кандид", case=False, na=False).to_numpy()
    is_doctor = type_s.str.contains("доктор", case=False, na=False).to_numpy()

    deg_s = meta_aligned.get("degree_pursued", pd.Series([""] * len(reg_nums))).fillna("").astype(str)
    degree_lower = np.char.lower(deg_s.to_numpy().astype(str))

    if "protection_date" in meta_aligned.columns:
        dt = pd.to_datetime(meta_aligned["protection_date"], errors="coerce")
        year_arr = dt.dt.year.astype("float").to_numpy()
    elif "protection_year" in meta_aligned.columns:
        year_arr = pd.to_numeric(meta_aligned["protection_year"], errors="coerce").astype("float").to_numpy()
    else:
        year_arr = np.full(len(reg_nums), np.nan, dtype="float")

    return df_meta, reg_nums, emb_matrix, is_candidate, is_doctor, degree_lower, year_arr


@st.cache_resource(show_spinner="Загрузка модели...")
def load_model():
    return SentenceTransformer(MODEL_NAME)


try:
    df_all, reg_nums, emb_matrix, is_candidate, is_doctor, degree_lower, year_arr = load_data()
    model = load_model()
except Exception as e:
    st.error(f"Ошибка при загрузке данных или модели: {e}")
    st.stop()


# ==========================
# ПОИСК
# ==========================

def _contains_any(deg_lower_arr: np.ndarray, patterns: List[str]) -> np.ndarray:
    m = np.zeros(len(deg_lower_arr), dtype=bool)
    for p in patterns:
        p = (p or "").strip().lower()
        if not p:
            continue
        m |= (np.char.find(deg_lower_arr, p) >= 0)
    return m


def build_filter_mask(
    candidate_selected: bool,
    doctor_selected: bool,
    science_selected: List[str],
    year_range: Optional[Tuple[int, int]],
) -> np.ndarray:
    mask = np.ones(len(reg_nums), dtype=bool)

    type_mask = np.zeros(len(reg_nums), dtype=bool)
    if candidate_selected:
        type_mask |= is_candidate
    if doctor_selected:
        type_mask |= is_doctor
    mask &= type_mask

    if science_selected:
        sci_mask = np.zeros(len(reg_nums), dtype=bool)
        for label in science_selected:
            patterns = SCIENCE_PATTERNS.get(label, [label])
            sci_mask |= _contains_any(degree_lower, patterns)
        mask &= sci_mask

    if year_range is not None:
        y0, y1 = int(year_range[0]), int(year_range[1])
        yr = year_arr
        mask &= (np.isnan(yr) | ((yr >= y0) & (yr <= y1)))

    return mask


def search_core(query: str, top_k: int = 10, mask=None):
    query = query.strip()
    if not query:
        return []

    q_emb = model.encode("query: " + query, normalize_embeddings=True)

    idx_pool = np.arange(len(reg_nums)) if mask is None else np.flatnonzero(mask)
    if idx_pool.size == 0:
        return []

    scores_pool = emb_matrix[idx_pool] @ q_emb
    top_k = min(int(top_k), len(scores_pool))
    top_local = np.argsort(-scores_pool)[:top_k]

    top_idx = idx_pool[top_local]
    top_scores = scores_pool[top_local]

    return [
        {"rank": i + 1, "registration_number": reg_nums[idx], "score": float(sc)}
        for i, (idx, sc) in enumerate(zip(top_idx, top_scores))
    ]


def extract_year_int(value) -> Optional[int]:
    if value is None:
        return None
    try:
        if not isinstance(value, str):
            if pd.isna(value):
                return None
            dt = pd.to_datetime(value, errors="coerce")
            if pd.isna(dt):
                return None
            return int(dt.year)
    except Exception:
        pass

    s = str(value).strip()
    if len(s) >= 4 and s[:4].isdigit():
        return int(s[:4])
    return None


def build_result_df(results):
    rows = []
    for r in results:
        reg = _norm_regnum(r["registration_number"])
        score = r["score"]

        if reg in df_all.index:
            meta = df_all.loc[reg]
            if isinstance(meta, pd.DataFrame):
                meta = meta.iloc[0]
        else:
            meta = pd.Series({}, index=df_all.columns)

        protection_year = extract_year_int(meta.get("protection_date", None))

        org_short = meta.get("author_org_short", None)
        if (
            org_short is None
            or (isinstance(org_short, float) and pd.isna(org_short))
            or str(org_short).lower() in {"none", "nan"}
        ):
            org_short = meta.get("author_org_name", None)

        rows.append(
            {
                "№": r["rank"],
                "score": float(round(score, 4)),
                "fio": meta.get("fio", None),
                "title": meta.get("title", None),
                "author_org_short": org_short,
                "dissertation_type": meta.get("dissertation_type", None),
                "protection_year": protection_year,
                "registration_number": meta.get("registration_number", reg),
                "vak_link": meta.get("vak_link", ""),
            }
        )

    if not rows:
        return pd.DataFrame(columns=DISPLAY_COLUMNS_ALL + ["vak_link"])

    return pd.DataFrame(rows).reset_index(drop=True)[DISPLAY_COLUMNS_ALL + ["vak_link"]]


def run_search(
    query: str,
    top_k: int,
    candidate_selected: bool,
    doctor_selected: bool,
    science_selected: List[str],
    year_range: Optional[Tuple[int, int]],
    only_openalex: bool,
    only_orcid: bool,
):
    mask = build_filter_mask(candidate_selected, doctor_selected, science_selected, year_range)

    prefetch_k = min(max(int(top_k) * 5, int(top_k)), 500)
    results = search_core(query, prefetch_k, mask=mask)

    df_raw = build_result_df(results)

    if df_raw.empty:
        empty_ui = pd.DataFrame(columns=["Выбрать"] + UI_TABLE_COLUMNS)
        empty_ui.index.name = "reg_norm"
        out = io.BytesIO()
        with pd.ExcelWriter(out, engine="xlsxwriter") as writer:
            pd.DataFrame().to_excel(writer, index=False)
        out.seek(0)
        return empty_ui, out, df_raw

    df_raw["reg_norm"] = df_raw["registration_number"].map(_norm_regnum)

    if not oa_enrich.empty:
        en = oa_enrich.reindex(df_raw["reg_norm"]).reset_index(drop=True)

        def _get(col: str, default):
            if col in en.columns:
                return en[col]
            return pd.Series([default] * len(df_raw))

        df_raw["openalex_url"] = _get("openalex_url", "").map(_norm_openalex_url)
        df_raw["orcid_url"] = _get("orcid_url", "").map(_norm_orcid_url)

        df_raw["h_index"] = pd.to_numeric(_get("h_index", np.nan), errors="coerce").astype("float64")
        df_raw["i10_index"] = pd.to_numeric(_get("i10_index", np.nan), errors="coerce").astype("float64")
        df_raw["works_count"] = pd.to_numeric(_get("works_count", np.nan), errors="coerce").astype("float64")
        df_raw["cited_by_count"] = pd.to_numeric(_get("cited_by_count", np.nan), errors="coerce").astype("float64")
    else:
        df_raw["openalex_url"] = ""
        df_raw["orcid_url"] = ""
        df_raw["h_index"] = np.nan
        df_raw["i10_index"] = np.nan
        df_raw["works_count"] = np.nan
        df_raw["cited_by_count"] = np.nan

    if only_openalex:
        df_raw = df_raw[df_raw["openalex_url"].map(_safe_text) != ""]
    if only_orcid:
        df_raw = df_raw[df_raw["orcid_url"].map(_safe_text) != ""]

    df_raw = df_raw.reset_index(drop=True)
    if len(df_raw) > int(top_k):
        df_raw = df_raw.iloc[: int(top_k)].copy()

    df_raw["reg_norm"] = df_raw["registration_number"].map(_norm_regnum)
    df_raw = df_raw.drop_duplicates(subset=["reg_norm"]).set_index("reg_norm", drop=True)

    if "№" in df_raw.columns:
        df_raw["№"] = np.arange(1, len(df_raw) + 1)

    fio_txt = df_raw["fio"].map(_safe_text)

    title_txt = df_raw["title"].map(_safe_text)
    title_frag = title_txt.map(_safe_fragment)
    vak_url = df_raw["vak_link"].map(_safe_text)
    title_cell = np.where(vak_url != "", vak_url + "#" + title_frag, title_txt)

    oa_url = df_raw["openalex_url"].map(_safe_text)
    oa_id = oa_url.map(_openalex_id_from_url)
    openalex_cell = np.where((oa_url != "") & (oa_id != ""), oa_url + "#" + oa_id, "")

    orcid_url = df_raw["orcid_url"].map(_safe_text)
    orcid_id = orcid_url.map(_orcid_id_from_url)
    orcid_cell = np.where((orcid_url != "") & (orcid_id != ""), orcid_url + "#" + orcid_id, "")

    df_ui = pd.DataFrame(
        {
            "Сходство": pd.to_numeric(df_raw["score"], errors="coerce").astype("float64"),
            "ФИО": fio_txt,
            "Название диссертации": title_cell,
            "Организация": df_raw["author_org_short"].map(_safe_text),
            "Тип": df_raw["dissertation_type"].map(_safe_text),
            "Год": pd.to_numeric(df_raw["protection_year"], errors="coerce").astype("float64"),
            "OpenAlex": openalex_cell,
            "ORCID": orcid_cell,
            "Регистрационный номер": df_raw["registration_number"].map(_safe_text),
            "h-index": df_raw["h_index"],
            "i10-index": df_raw["i10_index"],
            "Работ": df_raw["works_count"],
            "Цитат": df_raw["cited_by_count"],
        },
        index=df_raw.index,
    )
    df_ui.index.name = "reg_norm"

    df_excel_ru = df_raw.reset_index(drop=True).rename(columns=COLUMN_LABELS_RU_EXCEL)
    output = io.BytesIO()
    with pd.ExcelWriter(output, engine="xlsxwriter") as writer:
        df_excel_ru.to_excel(writer, index=False)
    output.seek(0)

    return df_ui, output, df_raw


def format_selected_list_from_raw(df_raw: pd.DataFrame, selected_regnorms: List[str]) -> str:
    lines = []
    for reg_norm in selected_regnorms:
        if df_raw is None or reg_norm not in df_raw.index:
            continue

        raw = df_raw.loc[reg_norm]
        fio = _safe_text(raw.get("fio"))
        title = _safe_text(raw.get("title"))
        year = _safe_text(raw.get("protection_year"))

        vak = _safe_text(raw.get("vak_link"))
        orcid = _safe_text(raw.get("orcid_url"))
        oa = _safe_text(raw.get("openalex_url"))

        link_parts = []
        if vak:
            link_parts.append(f"[ВАК]({vak})")
        if orcid:
            link_parts.append(f"[ORCID]({orcid})")
        if oa:
            link_parts.append(f"[OpenAlex]({oa})")

        links_line = ("  \n  " + " ".join(link_parts)) if link_parts else ""
        lines.append(f"- **{fio}** — {title} ({year}){links_line}")
    return "\n".join(lines)


# ==========================
# UI
# ==========================

st.markdown(
    "<h1 style='text-align: center; margin-bottom: 0.5rem;'>Поиск постдока🎓</h1>",
    unsafe_allow_html=True,
)

if "last_df_ui" not in st.session_state:
    st.session_state.last_df_ui = None
if "last_df_raw" not in st.session_state:
    st.session_state.last_df_raw = None
if "last_excel" not in st.session_state:
    st.session_state.last_excel = None

if "selected_regnorms" not in st.session_state:
    st.session_state.selected_regnorms = set()

if "search_id" not in st.session_state:
    st.session_state.search_id = 0

data_has_years = np.isfinite(year_arr).any()
year_max = int(np.nanmax(year_arr)) if data_has_years else None

with st.form("search_form"):
    top_k = st.slider("Сколько результатов показать", 1, 100, 20, 1)

    query = st.text_area(
        "Введите запрос",
        height=120,
        placeholder="Например: пластификаторы для самоуплотняющихся бетонов",
        key="query",
    )

    with st.expander("Расширенные настройки", expanded=False):
        st.markdown("**Диссертации:**")
        c1, c2 = st.columns(2)
        with c1:
            candidate_selected = st.checkbox("Кандидатские", value=True, key="dtype_candidate")
        with c2:
            doctor_selected = st.checkbox("Докторские", value=False, key="dtype_doctor")

        st.markdown("**Науки:**")
        cols = st.columns(3)
        science_selected = []
        for i, label in enumerate(SCIENCE_LABELS):
            default_val = label in DEFAULT_SCIENCES
            with cols[i % 3]:
                if st.checkbox(label, value=default_val, key=_keyify("sci_" + label)):
                    science_selected.append(label)

        st.markdown("**Годы защиты:**")
        if year_max is None:
            st.info("Годы защиты не найдены в данных — фильтр по годам недоступен.")
            year_range = None
        elif year_max < SLIDER_MIN_YEAR:
            st.info("В данных нет защит с 2005 года и позже — фильтр по годам недоступен.")
            year_range = None
        else:
            year_range = st.slider(
                "Диапазон лет",
                min_value=SLIDER_MIN_YEAR,
                max_value=year_max,
                value=(SLIDER_MIN_YEAR, year_max),
                step=1,
            )

        st.markdown("**Фильтрация по профилям:**")
        only_openalex = st.checkbox("Отображать только с OpenAlex", value=False, key="only_openalex")
        only_orcid = st.checkbox("Отображать только с ORCID", value=False, key="only_orcid")

        st.markdown("**Настройки отображения:**")
        disp_cols = st.columns(3)
        for i, col in enumerate(UI_TABLE_COLUMNS):
            with disp_cols[i % 3]:
                st.checkbox(
                    col,
                    value=DEFAULT_VISIBLE_UI.get(col, True),
                    key=_show_col_key(col),
                )

    c1, c2, c3 = st.columns([1, 1, 1])
    with c2:
        do_search = st.form_submit_button("🔍 Поиск", type="primary", use_container_width=True)

visible_ui_cols = [c for c in UI_TABLE_COLUMNS if st.session_state.get(_show_col_key(c), True)]
if not visible_ui_cols:
    visible_ui_cols = [c for c in UI_TABLE_COLUMNS if c not in {"Организация", "Тип", "Регистрационный номер"}]

if do_search:
    if not candidate_selected and not doctor_selected:
        st.warning("Выключены оба типа диссертаций. Включите «Кандидатские» и/или «Докторские».")
    else:
        with st.spinner("Идёт поиск по базе диссертаций..."):
            df_ui, excel_bytes, df_raw = run_search(
                query=query,
                top_k=top_k,
                candidate_selected=candidate_selected,
                doctor_selected=doctor_selected,
                science_selected=science_selected,
                year_range=year_range,
                only_openalex=only_openalex,
                only_orcid=only_orcid,
            )

        st.session_state.last_df_ui = df_ui
        st.session_state.last_df_raw = df_raw
        st.session_state.last_excel = excel_bytes

        if isinstance(df_ui, pd.DataFrame) and not df_ui.empty:
            st.session_state.selected_regnorms = set(st.session_state.selected_regnorms) & set(df_ui.index)
        else:
            st.session_state.selected_regnorms = set()

        st.session_state.search_id += 1

df_ui_saved = st.session_state.last_df_ui
df_raw_saved = st.session_state.last_df_raw
excel_saved = st.session_state.last_excel

if isinstance(df_ui_saved, pd.DataFrame) and not df_ui_saved.empty:
    st.success(f"Найдено записей: {len(df_ui_saved)}")

    selected_set = set(st.session_state.selected_regnorms) & set(df_ui_saved.index)
    st.session_state.selected_regnorms = selected_set

    df_display = df_ui_saved.copy()
    df_display.insert(0, "Выбрать", df_display.index.map(lambda x: x in selected_set))

    show_cols = [c for c in visible_ui_cols if c in df_display.columns]
    df_edit = df_display[["Выбрать"] + show_cols].copy()

    full_column_config = {
        "Выбрать": st.column_config.CheckboxColumn("Выбрать", width="small"),
        "Сходство": st.column_config.NumberColumn("Сходство", format="%.4f", width="small"),
        "ФИО": st.column_config.TextColumn("ФИО", width="medium"),
        "Название диссертации": st.column_config.LinkColumn(
            "Название диссертации",
            display_text=r"(?:.*#)?(.*)$",
            width="large",
            help="Название ведёт на ВАК (если ссылка есть).",
            validate=r"^https?://.+#.+$|^.+$",
        ),
        "Организация": st.column_config.TextColumn("Организация", width="large"),
        "Тип": st.column_config.TextColumn("Тип", width="small"),
        "Год": st.column_config.NumberColumn("Год", format="%.0f", width="small"),
        "OpenAlex": st.column_config.LinkColumn(
            "OpenAlex",
            display_text=r"(?:.*#)?(.*)$",
            width="small",
            help="ID автора в OpenAlex (если найден).",
            validate=r"^https?://openalex\.org/A\d+#A\d+$|^$",
        ),
        "ORCID": st.column_config.LinkColumn(
            "ORCID",
            display_text=r"(?:.*#)?(.*)$",
            width="small",
            help="ORCID автора (если найден).",
            validate=r"^https?://orcid\.org/0000-[0-9X]{4}-[0-9X]{4}-[0-9X]{4}#0000-[0-9X]{4}-[0-9X]{4}-[0-9X]{4}$|^$",
        ),
        "Регистрационный номер": st.column_config.TextColumn("Регистрационный номер", width="medium"),
        "h-index": st.column_config.NumberColumn("h-index", format="%.0f", width="small"),
        "i10-index": st.column_config.NumberColumn("i10-index", format="%.0f", width="small"),
        "Работ": st.column_config.NumberColumn("Работ", format="%.0f", width="small"),
        "Цитат": st.column_config.NumberColumn("Цитат", format="%.0f", width="small"),
    }
    column_config_filtered = {k: v for k, v in full_column_config.items() if k in df_edit.columns}

    edited = st.data_editor(
        df_edit,
        use_container_width=True,
        hide_index=True,
        num_rows="fixed",
        column_config=column_config_filtered,
        disabled=[c for c in df_edit.columns if c != "Выбрать"],
        key=f"editor_{st.session_state.search_id}",
    )

    if isinstance(edited, pd.DataFrame) and "Выбрать" in edited.columns:
        st.session_state.selected_regnorms = set(edited.index[edited["Выбрать"] == True].tolist())

    selected_regnorms = sorted(list(st.session_state.selected_regnorms))

    if selected_regnorms:
        with st.expander("Полный текст и ссылки (для выбранных строк)", expanded=False):
            for reg_norm in selected_regnorms[:50]:
                if df_raw_saved is None or reg_norm not in df_raw_saved.index:
                    continue
                raw = df_raw_saved.loc[reg_norm]

                fio = _safe_text(raw.get("fio"))
                title = _safe_text(raw.get("title"))
                org = _safe_text(raw.get("author_org_short"))
                year = _safe_text(raw.get("protection_year"))

                vak = _safe_text(raw.get("vak_link"))
                orcid = _safe_text(raw.get("orcid_url"))
                oa = _safe_text(raw.get("openalex_url"))

                links = []
                if vak:
                    links.append(f"[ВАК]({vak})")
                if orcid:
                    links.append(f"[ORCID]({orcid})")
                if oa:
                    links.append(f"[OpenAlex]({oa})")

                links_md = ("  \n  " + " ".join(links)) if links else ""

                st.markdown(
                    f"- **{fio}** — {title}\n"
                    f"  \n  {org} ({year}){links_md}"
                )

    if excel_saved is not None:
        st.download_button(
            label="💾 Скачать результаты в Excel",
            data=excel_saved,
            file_name="search_results.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        )

    st.markdown("---")
    st.subheader("Запрос расширенной информации")

    with st.form("request_form"):
        requester_fio = st.text_input("Ваше ФИО", placeholder="Иванов Иван Иванович")
        requester_email = st.text_input("Email", placeholder="name@example.com")

        st.markdown("**Перечень диссертаций:**")
        if selected_regnorms:
            st.markdown(format_selected_list_from_raw(df_raw_saved, selected_regnorms))
        else:
            st.info("Отметьте диссертации чекбоксом «Выбрать» — здесь появится перечень.")

        comment = st.text_area(
            "Комментарий",
            height=160,
            placeholder=(
                "Какую дополнительную информацию по авторам диссертаций вы хотите получить?\n"
                "Какие замечания/пожелания по функционалу системы?"
            ),
        )

        send_request = st.form_submit_button("📨 Отправить запрос", type="primary", use_container_width=True)

    if send_request:
        if not requester_fio.strip() or not requester_email.strip():
            st.warning("Поля «Ваше ФИО» и «Email» обязательны. Заполните их, чтобы отправить запрос.")
        elif len(selected_regnorms) == 0:
            st.warning("Выберите хотя бы одну диссертацию (чекбокс «Выбрать»).")
        else:
            items = []
            for reg_norm in selected_regnorms:
                if df_raw_saved is None or reg_norm not in df_raw_saved.index:
                    continue
                raw = df_raw_saved.loc[reg_norm].to_dict()
                items.append(
                    {
                        "author_fio": raw.get("fio"),
                        "title": raw.get("title"),
                        "org": raw.get("author_org_short"),
                        "year": raw.get("protection_year"),
                        "vak_link": raw.get("vak_link"),
                        "registration_number": raw.get("registration_number"),
                        "score": raw.get("score"),
                        "openalex_url": raw.get("openalex_url", ""),
                        "orcid_url": raw.get("orcid_url", ""),
                        "h_index": raw.get("h_index"),
                        "i10_index": raw.get("i10_index"),
                        "works_count": raw.get("works_count"),
                        "cited_by_count": raw.get("cited_by_count"),
                    }
                )

            payload = {
                "created_at_utc": datetime.now(timezone.utc).isoformat(),
                "requester": {
                    "fio": requester_fio.strip(),
                    "email": requester_email.strip(),
                    "comment": (comment or "").strip(),
                },
                "items_count": len(items),
                "items": items,
            }

            try:
                path = save_request_to_hub(payload)
                st.success(f"Запрос {path} сохранен")
            except Exception as e:
                st.error(
                    "Не удалось сохранить запрос в репозиторий.\n\n"
                    f"Ошибка: {e}\n\n"
                    "Проверьте HF_WRITE_TOKEN (write) и repo_type."
                )
else:
    st.info("Введите запрос и нажмите «Поиск». После этого можно выбрать диссертации и отправить запрос.")

st.markdown(
    "<p style='font-size: 0.8rem; text-align: right; color: gray;'>(с) Антон Лощилов, 2025</p>",
    unsafe_allow_html=True,
)