import re
import yaml
from typing import List, Dict, Any
import logging

logger = logging.getLogger("lightrag")

def load_config():
    with open('config.yaml', 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

config = load_config()


def enrich_document(content: str, metadata: Dict[str, Any]) -> str:
    """Обогащает документ метаданными для улучшения поиска"""
    enrichment = []
    
    # Добавляем метаданные, если они есть
    if "requirement_code" in metadata:
        enrichment.append(f"Код требования: {metadata['requirement_code']}")
    
    if "title" in metadata:
        enrichment.append(f"Краткая формулировка: {metadata['title']}")
    
    if "status" in metadata:
        enrichment.append(f"Статус: {metadata['status']}")
    
    # Добавляем информацию о сущностях
    if "entities" in metadata:
        entity_info = []
        for entity_name, entity_data in metadata["entities"].items():
            entity_desc = []
            entity_desc.append(f"Сущность: {entity_name}")
            
            if entity_data.get("attributes"):
                entity_desc.append(f"Атрибуты: {', '.join(entity_data['attributes'])}")
            
            if entity_data.get("relationships"):
                relations = [rel["related_to"] for rel in entity_data["relationships"]]
                entity_desc.append(f"Связи: {', '.join(relations)}")
            
            entity_info.append(" | ".join(entity_desc))
        
        if entity_info:
            enrichment.append("Информация о сущностях:")
            enrichment.extend(entity_info)
    
    # Если есть информация для обогащения, добавляем её в начало документа
    if enrichment:
        metadata_section = "\n".join(enrichment)
        return f"МЕТАДАННЫЕ ДОКУМЕНТА:\n{metadata_section}\n\n---\n\n{content}"
    else:
        return content

def extract_metadata(content: str) -> Dict[str, Any]:
    """Извлекает метаданные из страницы Confluence"""
    metadata = {}
    
    # Пытаемся найти код требования
    req_match = re.search(r'Код требования[^\|]*\|\s*([^\|]+)', content)
    if req_match:
        metadata["requirement_code"] = req_match.group(1).strip()
    
    # Пытаемся найти статус
    status_match = re.search(r'Статус[^\|]*\|\s*([^\|]+)', content)
    if status_match:
        metadata["status"] = status_match.group(1).strip()
    
    # Пытаемся найти краткую формулировку
    title_match = re.search(r'Краткая формулировка[^\|]*\|\s*([^\|]+)', content)
    if title_match:
        metadata["title"] = title_match.group(1).strip()
    
    # Извлекаем сущности и связи
    entities = extract_entities(content)
    if entities:
        metadata["entities"] = entities
    
    return metadata

def extract_entities(content: str) -> Dict[str, Dict[str, List]]:
    """Извлекает сущности и их атрибуты из документа"""
    entities = {}
    
    # Находим упоминания сущностей
    entity_matches = re.finditer(r'сущност[ьи]\s+([a-zA-Z0-9_]+)', content, re.IGNORECASE)
    for match in entity_matches:
        entity_name = match.group(1).strip()
        if entity_name:
            entities[entity_name] = {"attributes": [], "relationships": []}
    
    # Ищем атрибуты в таблицах (упрощенно)
    for entity_name in entities:
        # Ищем таблицы атрибутов после упоминания сущности
        entity_pos = content.find(entity_name)
        if entity_pos > -1:
            # Ищем ближайшую таблицу после упоминания сущности
            table_match = re.search(r'\|\s*([^|]+)\s*\|\s*([^|]+)\s*\|', content[entity_pos:])
            if table_match:
                # Извлекаем строки таблицы
                table_rows = re.finditer(r'\|\s*([^|]+)\s*\|\s*([^|]+)\s*\|', content[entity_pos:])
                for row_match in table_rows:
                    col1 = row_match.group(1).strip()
                    if col1.lower() not in ["атрибут", "поле", "id"]:
                        entities[entity_name]["attributes"].append(col1)
    
    # Ищем отношения между сущностями
    for entity_name in entities:
        # Ищем упоминания отношений
        rel_patterns = [
            rf'{entity_name}\s+связан[а]?\s+с\s+([a-zA-Z0-9_]+)',
            rf'Тип связи.*?{entity_name}.*?([a-zA-Z0-9_]+)',
            rf'{entity_name}.*?one-to-many.*?([a-zA-Z0-9_]+)',
            rf'{entity_name}.*?many-to-one.*?([a-zA-Z0-9_]+)'
        ]
        
        for pattern in rel_patterns:
            rel_matches = re.finditer(pattern, content, re.IGNORECASE)
            for match in rel_matches:
                related_entity = match.group(1).strip()
                if related_entity in entities and related_entity != entity_name:
                    relation = {"related_to": related_entity}
                    entities[entity_name]["relationships"].append(relation)
    
    return entities
