import gradio as gr
import pandas as pd
import pdfplumber
import os
import tempfile
import re
from datetime import datetime
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import A4
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.lib import colors


class KDChecker:
def __init__(self):
self.excel_db = pd.DataFrame()
self.cabinet_list = []
self.known_docs = ["Э3", "В4", "ПЭ3", "ВО", "ТЭ5", "СБ", "С5", "ОЛ", "Э1", "Э4", "Э7", "Д3", "Э6"]

def load_excel_db(self, excel_file):
if excel_file is None:
return "Файл не выбран", gr.update(choices=[], value=None)

all_data = []
sheets_log = []

try:
xls = pd.read_excel(excel_file.name, sheet_name=None, header=None)

for sheet_name, df_raw in xls.items():
header_row_index = -1
cab_col_idx = -1
rem_col_idx = -1

for i in range(min(20, len(df_raw))):
row_values = [str(x).lower().strip() for x in df_raw.iloc[i].values]
c_idx = -1
r_idx = -1
for idx, val in enumerate(row_values):
if "шкаф" in val or "cabinet" in val: c_idx = idx
if "примечание" in val or "remark" in val: r_idx = idx
if c_idx != -1 and r_idx != -1:
header_row_index = i
cab_col_idx = c_idx
rem_col_idx = r_idx
break

if header_row_index != -1:
df = pd.read_excel(excel_file.name, sheet_name=sheet_name, header=header_row_index)
df_subset = df.iloc[:, [cab_col_idx, rem_col_idx]]
df_subset.columns = ["Cabinet", "Remark"]

df_subset["Cabinet"] = df_subset["Cabinet"].ffill()
df_subset = df_subset.dropna(subset=["Remark"]).astype(str)

df_subset["Cabinet_Clean"] = df_subset["Cabinet"].apply(
lambda x: x.strip().replace(" ", "").replace("\n", "").replace("\r", "")
)

all_data.append(df_subset)
sheets_log.append(f"Лист '{sheet_name}': {len(df_subset)} стр.")
else:
sheets_log.append(f"Лист '{sheet_name}': заголовки не найдены")

if not all_data:
return "❌ Ошибка: Не найдены заголовки 'Шкаф' и 'Примечание'.", gr.update(choices=[], value=None)

self.excel_db = pd.concat(all_data, ignore_index=True)
self.cabinet_list = sorted(self.excel_db["Cabinet"].unique().tolist())

msg = f"✅ База знаний загружена успешно!\nВсего записей: {len(self.excel_db)}\nОбработаны листы: {', '.join(sheets_log)}"
return msg, gr.update(choices=self.cabinet_list, value=None, interactive=True)

except Exception as e:
return f"❌ Ошибка чтения Excel: {e}", gr.update(choices=[], value=None)

def extract_text(self, pdf_path):
try:
full_text = ""
with pdfplumber.open(pdf_path) as pdf:
for page in pdf.pages:
full_text += (page.extract_text() or "") + "\n"
return full_text
except:
return ""

def find_all_decimal_numbers(self, text):
pattern = r"(РЛТ|ЛДАР|ВНАР|ШТМ)[\s\.]*\d{1}[\s\.]*\d{3}[\s\.]*[А-ЯA-Z]{1,4}[\s\.]*\d{3}(-[\d]+)?"
matches = []
for match in re.finditer(pattern, text):
clean_num = match.group(0).replace(" ", "").replace("\n", "")
if clean_num not in matches:
matches.append(clean_num)
return matches

def determine_doc_type(self, filename):
fname = filename.upper()
if "С2" in fname: return "С2"
if "ПЭ3" in fname or "ПЕРЕЧЕНЬ" in fname: return "ПЭ3"
if "Э3" in fname or "СХЕМА ЭЛЕКТРИЧЕСКАЯ" in fname: return "Э3"
if "Э4" in fname: return "Э4"
if "В4" in fname or "СПЕЦИФИКАЦИЯ" in fname: return "В4"
if "ВО" in fname or "Э7" in fname or "ГАБАРИТ" in fname: return "ВО"
if "ТЭ5" in fname or "ТАБЛИЦА" in fname: return "ТЭ5"
if "СБ" in fname: return "СБ"
if "С5" in fname: return "С5"
if "ОЛ" in fname: return "ОЛ"
if "Э1" in fname: return "Э1"
if "Э6" in fname or "ЗАЗЕМЛЕНИЯ" in fname: return "Э6"
if "Д3" in fname or "МОНТАЖ" in fname: return "Д3"
return "UNKNOWN"

def get_remarks(self, cabinet_key, is_clean_key=True):
if self.excel_db.empty: return {}

if is_clean_key:
target = cabinet_key.replace(" ", "")
mask = self.excel_db['Cabinet_Clean'].str.contains(re.escape(target), case=False, na=False)
else:
mask = self.excel_db['Cabinet'] == cabinet_key

rows = self.excel_db[mask]
if rows.empty: return {}

parsed = {}

for remark_cell in rows['Remark']:
cell_text = str(remark_cell)
cell_text = re.sub(r'(\d+)\.([А-ЯA-Z])', r'\1. \2', cell_text)

items = re.split(r'(?:^|\n)\s*(?=\d+[\.\)])', cell_text)

for item in items:
if len(item) < 3: continue
clean_item = item.strip()
clean_item_no_num = re.sub(r'^\d+[\.\)]\s*', '', clean_item)

doc_pattern = r'^(?:Документ\s+|В\s+)?([А-ЯA-Z0-9\s,\(\)\-]+?)(?:[\.\:\-]|\s+)(.*)'
match = re.match(doc_pattern, clean_item_no_num, re.IGNORECASE | re.DOTALL)

detected_docs = []
final_text = clean_item

if match:
potential_docs_str = match.group(1).upper()
cleaned_codes = potential_docs_str.replace("(", " ").replace(")", " ").replace(",", " ")
parts = cleaned_codes.split()

valid_parts = [p for p in parts if p in self.known_docs]
if valid_parts:
detected_docs = valid_parts
final_text = match.group(2).strip()

if not detected_docs:
detected_docs = ["ALL"]

for doc in detected_docs:
if doc not in parsed: parsed[doc] = []
parsed[doc].append(final_text)

return parsed

def check_files(self, files, manual_cabinet):
if not files: return "Файлы не загружены", None
if self.excel_db.empty: return "Сначала загрузите Excel базу!", None

checklist = {}
detected_cabinet = "Не определен"
found_by_method = ""
is_manual = False

# 1. Ручной выбор
if manual_cabinet and manual_cabinet.strip():
detected_cabinet = manual_cabinet
found_by_method = "manual"
is_manual = True

# 2. Автопоиск
else:
all_pdf_text = ""
for file in files:
all_pdf_text += self.extract_text(file.name) + "\n"

# А. По номеру
pdf_numbers = self.find_all_decimal_numbers(all_pdf_text)
db_clean_keys = set(self.excel_db["Cabinet_Clean"].tolist())

for cand in pdf_numbers:
if cand in db_clean_keys:
detected_cabinet = cand
found_by_method = "number"
break

# Б. По имени
if detected_cabinet == "Не определен":
unique_cabinets = self.excel_db["Cabinet"].unique()
for cab_name in unique_cabinets:
sub_names = [n.strip() for n in cab_name.split(',')]
for sub_name in sub_names:
if len(sub_name) < 5: continue
if sub_name.lower() in all_pdf_text.lower():
detected_cabinet = cab_name
found_by_method = "name"
break
if found_by_method == "name": break

if detected_cabinet == "Не определен":
examples = ", ".join(self.excel_db["Cabinet"].head(3).tolist())
return f"⚠️ Шкаф не опознан автоматически.\nСовет: Выберите похожий шкаф из выпадающего списка вручную.", None

is_clean_search = (found_by_method == "number")
remarks = self.get_remarks(detected_cabinet, is_clean_key=is_clean_search)

if not remarks:
return f"⚠️ Шкаф '{detected_cabinet}' выбран, но в базе нет замечаний для него.", None

processed_count = 0
for file in files:
fname = os.path.basename(file.name)
dtype = self.determine_doc_type(fname)

tasks = []
if dtype in remarks: tasks.extend(remarks[dtype])
if "ALL" in remarks and dtype != "С2": tasks.extend(remarks["ALL"])

if tasks:
checklist[fname] = list(dict.fromkeys(tasks))
processed_count += 1

pdf_title = detected_cabinet
if is_manual: pdf_title += " (Выбор вручную)"

pdf = self.create_pdf(pdf_title, checklist)
total = sum(len(v) for v in checklist.values())

method_str = "Ручной выбор" if is_manual else (
"По децимальному номеру" if is_clean_search else "По наименованию")

return f"✅ Готово!\n\n📂 Шкаф: {detected_cabinet}\n🔍 Метод: {method_str}\n📄 Обработано файлов: {processed_count}\n🚩 Вс