from __future__ import annotations

import logging
from pathlib import Path
from typing import Iterable, List, Sequence

from langchain_core.documents import Document

from smart_bot_factory.rag.vectorstore import VectorStore
from langchain_text_splitters import CharacterTextSplitter

logger = logging.getLogger(__name__)

BASE_PATH = Path("mdclinica")
BOT_ID = "mdclinica"

def read_text(file_path: Path) -> str:
    try:
        content = file_path.read_text(encoding="utf-8").strip()
        return content
    except FileNotFoundError:
        logger.warning("Файл не найден: %s", file_path)
        return ""


def create_documents(base_path: Path) -> tuple[List[Document], List[Document]]:
    category_docs: List[Document] = []
    service_docs: List[Document] = []

    if not base_path.exists():
        logger.error("Папка с данными не найдена: %s", base_path.resolve())
        return category_docs, service_docs

    for category_dir in base_path.iterdir():
        if not category_dir.is_dir():
            continue

        category_slug = category_dir.name
        category_name = read_text(category_dir / "name.txt") or category_slug

        overview_dir = category_dir / "overview"
        sections_dir = overview_dir / "sections"
        if sections_dir.exists():
            for section_file in sections_dir.glob("*.txt"):
                section_content = read_text(section_file)
                if not section_content:
                    continue
                section_name = section_file.stem
                category_docs.append(
                    Document(
                        page_content=section_content,
                        metadata={
                            "bot_id": BOT_ID,
                            "doc_type": "category_section",
                            "category_slug": category_slug,
                            "category_name": category_name,
                            "section": section_name,
                            "source_path": str(section_file.relative_to(base_path)),
                        },
                    )
                )

        services_root = category_dir / "services"
        if not services_root.exists():
            continue

        for service_dir in services_root.iterdir():
            if not service_dir.is_dir():
                continue

            service_slug = service_dir.name
            service_name = read_text(service_dir / "name.txt") or service_slug

            service_sections_dir = service_dir / "sections"
            if not service_sections_dir.exists():
                continue

            for section_file in service_sections_dir.glob("*.txt"):
                section_content = read_text(section_file)
                if not section_content:
                    continue
                section_name = section_file.stem
                service_docs.append(
                    Document(
                        page_content=section_content,
                        metadata={
                            "bot_id": BOT_ID,
                            "doc_type": "service_section",
                            "category_slug": category_slug,
                            "category_name": category_name,
                            "service_slug": service_slug,
                            "service_name": service_name,
                            "section": section_name,
                            "source_path": str(section_file.relative_to(base_path)),
                        },
                    )
                )

    return category_docs, service_docs


DEFAULT_SPLITTER = CharacterTextSplitter(
    chunk_size=800,
    chunk_overlap=140,
    separator="",
)

PRICE_ENTRIES_PER_CHUNK = 15


def _clone_document(doc: Document, content: str, chunk_index: int, chunk_type: str) -> Document:
    metadata = dict(doc.metadata)
    metadata.update(
        {
            "chunk_index": chunk_index,
            "chunk_type": chunk_type,
        }
    )
    return Document(page_content=content.strip(), metadata=metadata)


def _chunk_with_splitter(doc: Document, splitter: CharacterTextSplitter) -> Iterable[Document]:
    if not doc.page_content:
        return []

    # Получаем название категории/услуги (из контента или metadata)
    category_or_service_name = _get_category_or_service_name(doc)
    
    # Извлекаем заголовок секции из контента
    section_title = _extract_section_title(doc.page_content, category_or_service_name)
    
    # Удаляем название и заголовок секции из контента перед чанкингом
    content_for_chunking = doc.page_content
    lines = content_for_chunking.splitlines()
    filtered_lines = []
    skip_next_empty = False
    
    for line in lines:
        stripped = line.strip()
        
        # Пропускаем название категории/услуги
        if category_or_service_name and stripped == category_or_service_name:
            skip_next_empty = True
            continue
        
        # Пропускаем пустую строку после названия
        if not stripped and skip_next_empty:
            skip_next_empty = False
            continue
        
        # Пропускаем заголовок секции (если он есть в начале)
        if section_title and stripped == section_title:
            skip_next_empty = True
            continue
        
        # Пропускаем пустую строку после заголовка секции
        if not stripped and skip_next_empty:
            skip_next_empty = False
            continue
        
        filtered_lines.append(line)
    
    content_for_chunking = "\n".join(filtered_lines).strip()
    
    # Если после удаления заголовков ничего не осталось, используем весь контент
    if not content_for_chunking:
        content_for_chunking = doc.page_content

    chunks = splitter.split_text(content_for_chunking)
    if not chunks:
        return []

    # Добавляем название категории/услуги в начало каждого чанка
    result = []
    for index, chunk_text in enumerate(chunks):
        chunk_lines = []
        if category_or_service_name:
            chunk_lines.append(category_or_service_name)
            chunk_lines.append("")
        chunk_lines.append(chunk_text)
        content = "\n".join(chunk_lines).strip()
        result.append(_clone_document(doc, content, index, "text"))

    return result


def _extract_category_or_service_name(content: str) -> str:
    """Извлекает название категории или услуги из начала контента (первая непустая строка)."""
    lines = content.splitlines()
    
    # Ключевые слова, которые указывают на название секции, а не категории/услуги
    section_keywords = ["цены", "о процедуре", "показания", "противопоказания", 
                       "специалисты", "вопрос", "ответ", "преимущества", "услуги"]
    
    for line in lines:
        stripped = line.strip()
        if not stripped:
            continue
            
        # Проверяем, не начинается ли строка с ключевого слова секции
        line_lower = stripped.lower()
        
        # Если строка начинается с ключевого слова секции - это не название категории/услуги
        is_section_title = any(line_lower.startswith(keyword) for keyword in section_keywords)
        
        if not is_section_title:
            # Это название категории/услуги - возвращаем его
            return stripped
    
    return ""


def _extract_section_title(content: str, category_or_service_name: str = "") -> str:
    """Извлекает заголовок секции из контента (первая строка после названия категории/услуги)."""
    lines = content.splitlines()
    found_name = False
    
    for line in lines:
        stripped = line.strip()
        if not stripped:
            if found_name:
                continue
            else:
                continue
        
        # Если нашли название категории/услуги, следующая непустая строка - заголовок секции
        if category_or_service_name and stripped == category_or_service_name:
            found_name = True
            continue
        
        if found_name:
            # Ключевые слова, которые указывают на заголовок секции
            section_keywords = ["цены", "о процедуре", "показания", "противопоказания", 
                               "специалисты", "вопрос", "ответ", "преимущества", "услуги"]
            line_lower = stripped.lower()
            if any(line_lower.startswith(keyword) for keyword in section_keywords):
                return stripped
    
    return ""


def _get_category_or_service_name(doc: Document) -> str:
    """
    Получает название категории или услуги из документа.
    Сначала пытается извлечь из контента, потом использует metadata как fallback.
    """
    # Сначала пытаемся извлечь из контента
    name_from_content = _extract_category_or_service_name(doc.page_content)
    if name_from_content:
        return name_from_content
    
    # Если не получилось, используем metadata
    # Для услуг используем service_name, для категорий - category_name
    service_name = doc.metadata.get("service_name", "")
    if service_name:
        return service_name
    
    category_name = doc.metadata.get("category_name", "")
    if category_name:
        return category_name
    
    return ""


def _group_price_entries(lines: Sequence[str], category_or_service_name: str = "", section_title: str = "") -> tuple[list[str], list[list[str]]]:
    header: list[str] = []
    entries: list[list[str]] = []
    current_entry: list[str] = []
    skip_next_empty = False

    for line in lines:
        stripped = line.strip()
        
        # Пропускаем название категории/услуги
        if category_or_service_name and stripped == category_or_service_name:
            skip_next_empty = True
            continue
        
        # Пропускаем пустую строку после названия
        if not stripped and skip_next_empty:
            skip_next_empty = False
            continue
        
        # Пропускаем заголовок секции
        if section_title and stripped == section_title:
            skip_next_empty = True
            continue
        
        # Пропускаем пустую строку после заголовка секции
        if not stripped and skip_next_empty:
            skip_next_empty = False
            continue
            
        if stripped.startswith("•"):
            if current_entry:
                entries.append(current_entry)
                current_entry = []
            current_entry = [line]
            continue

        if entries or current_entry:
            if stripped or current_entry:
                current_entry.append(line)
        else:
            header.append(line)

    if current_entry:
        entries.append(current_entry)

    return header, entries


def _chunk_prices(doc: Document) -> Iterable[Document]:
    # Получаем название категории/услуги (из контента или metadata)
    category_or_service_name = _get_category_or_service_name(doc)
    
    # Извлекаем заголовок секции из контента
    section_title = _extract_section_title(doc.page_content, category_or_service_name)
    
    lines = doc.page_content.splitlines()
    header, entries = _group_price_entries(lines, category_or_service_name, section_title)

    if not entries:
        return _chunk_with_splitter(doc, DEFAULT_SPLITTER)

    chunks: list[Document] = []
    for index in range(0, len(entries), PRICE_ENTRIES_PER_CHUNK):
        subset = entries[index : index + PRICE_ENTRIES_PER_CHUNK]
        start_no = index + 1
        end_no = index + len(subset)
        chunk_lines = []
        
        # ВАЖНО: Всегда добавляем название категории/услуги в начало каждого чанка
        if category_or_service_name:
            chunk_lines.append(category_or_service_name)
            chunk_lines.append("")
        
        if header:
            chunk_lines.extend(header)
            chunk_lines.append("")
            
        for entry in subset:
            if chunk_lines and chunk_lines[-1]:
                chunk_lines.append("")
            chunk_lines.extend(entry)
        content = "\n".join(chunk_lines).strip()
        chunk = _clone_document(
            doc,
            content,
            len(chunks),
            "prices",
        )
        chunk.metadata["price_range"] = f"{start_no}-{end_no}"
        chunks.append(chunk)

    return chunks


def chunk_documents(documents: Iterable[Document]) -> Iterable[Document]:
    for doc in documents:
        content = doc.page_content
        if not content:
            continue

        section = (doc.metadata.get("section") or "").lower()
        content_lower = content.lower()

        if "цены" in section or content_lower.startswith("цены"):
            for chunk in _chunk_prices(doc):
                if chunk.page_content:
                    yield chunk
            continue

        for chunk in _chunk_with_splitter(doc, DEFAULT_SPLITTER):
            if chunk.page_content:
                yield chunk


def main() -> None:
    logging.basicConfig(level=logging.INFO)

    logger.info("Создание документов из %s", BASE_PATH.resolve())
    category_docs, service_docs = create_documents(BASE_PATH)

    logger.info("Документов по разделам: %d", len(category_docs))
    logger.info("Документов по услугам: %d", len(service_docs))

    if not category_docs and not service_docs:
        logger.warning("Нет данных для загрузки")
        return

    category_store = VectorStore(bot_id=BOT_ID, table_name="service_categories")
    service_store = VectorStore(bot_id=BOT_ID, table_name="serv")

    if category_docs:
        logger.info("Добавление документов по разделам...")
        category_store.add_documents(list(chunk_documents(category_docs)))

    if service_docs:
        logger.info("Добавление документов по услугам...")
        service_store.add_documents(list(chunk_documents(service_docs)))

    logger.info("Готово")


if __name__ == "__main__":
    main()