import json
import zipfile
import pandas as pd
from huggingface_hub import hf_hub_download, list_repo_files
from llama_index.core import Document
from llama_index.core.text_splitter import SentenceSplitter
from my_logging import log_message

# Configuration
CHUNK_SIZE = 512
CHUNK_OVERLAP = 50

def chunk_text_documents(documents):
    """Simple text chunking with sentence awareness"""
    text_splitter = SentenceSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP
    )
    
    chunked = []
    for doc in documents:
        chunks = text_splitter.get_nodes_from_documents([doc])
        for i, chunk in enumerate(chunks):
            chunk.metadata.update({
                'chunk_id': i,
                'total_chunks': len(chunks)
            })
            chunked.append(chunk)
    
    log_message(f"✓ Text: {len(documents)} docs → {len(chunked)} chunks")
    return chunked


def chunk_table_by_rows(table_data, doc_id, max_rows=30):
    """Split large tables into row blocks - OPTIMIZED"""
    headers = table_data.get('headers', [])
    rows = table_data.get('data', [])
    table_num = table_data.get('table_number', 'unknown')
    table_title = table_data.get('table_title', '')
    section = table_data.get('section', '')
    
    # Normalize table number
    table_num_clean = str(table_num).replace('№', '').strip()
    
    if not rows:
        return []
    
    # Ensure document_id is set
    if 'document_id' not in table_data:
        table_data['document_id'] = doc_id
    
    # Small table: keep whole
    if len(rows) <= max_rows:
        content = format_table_content(table_data, headers, rows)
        return [Document(
            text=content,
            metadata={
                'type': 'table',
                'document_id': doc_id,
                'table_number': table_num_clean,
                'table_number_original': table_num,
                'table_title': table_title,
                'section': section,
                'total_rows': len(rows),
                'is_complete_table': True
            }
        )]
    
    # Large table: split by row blocks
    chunks = []
    for i in range(0, len(rows), max_rows):
        chunk_rows = rows[i:i+max_rows]
        content = format_table_content(
            table_data, 
            headers, 
            chunk_rows, 
            chunk_info=f"Часть таблицы: строки {i+1}-{i+len(chunk_rows)} из {len(rows)} всего"
        )
        
        chunks.append(Document(
            text=content,
            metadata={
                'type': 'table',
                'document_id': doc_id,
                'table_number': table_num_clean,
                'table_number_original': table_num,
                'table_title': table_title,
                'section': section,
                'chunk_id': i // max_rows,
                'row_start': i,
                'row_end': i + len(chunk_rows),
                'total_rows': len(rows),
                'total_chunks': (len(rows) + max_rows - 1) // max_rows,
                'is_complete_table': False
            }
        ))
    
    log_message(f"  📊 Table {table_num_clean} ({doc_id}): {len(rows)} rows → {len(chunks)} chunks")
    return chunks


def format_table_content(table_data, headers, rows, chunk_info=""):
    """Format table for semantic search - OPTIMIZED"""
    doc_id = table_data.get('document_id', table_data.get('document', 'unknown'))
    table_num = table_data.get('table_number', 'unknown')
    table_title = table_data.get('table_title', '')
    section = table_data.get('section', '')
    
    # Normalize table number - remove № prefix for consistent searching
    table_num_clean = str(table_num).replace('№', '').strip()
    
    # Create highly searchable header
    content = f"ТАБЛИЦА {table_num_clean}\n"
    content += f"Документ: {doc_id}\n"
    if table_title:
        content += f"Название таблицы: {table_title}\n"
    if section:
        content += f"Раздел: {section}\n"
    if chunk_info:
        content += f"{chunk_info}\n"
    content += f"{'='*60}\n\n"
    
    # Add multiple search-friendly descriptions
    content += f"Это таблица {table_num_clean} из документа {doc_id}. "
    content += f"Таблица номер {table_num_clean}. "
    if table_title:
        content += f"Таблица называется: {table_title}. "
        content += f"Содержание таблицы: {table_title}. "
    if section:
        content += f"Таблица находится в разделе {section}. "
    
    # Add searchable patterns
    content += f"Если вы ищете таблицу {table_num_clean} в {doc_id}, это она. "
    content += f"\n\nСОДЕРЖАНИЕ ТАБЛИЦЫ:\n"
    
    # Headers with emphasis
    if headers:
        header_str = ' | '.join(str(h) for h in headers)
        content += f"\nКОЛОНКИ: {header_str}\n"
        content += f"Заголовки столбцов: {header_str}\n\n"
    
    # Row data with better formatting
    content += "ДАННЫЕ:\n"
    for idx, row in enumerate(rows, 1):
        if isinstance(row, dict):
            parts = [f"{k}: {v}" for k, v in row.items() 
                    if v and str(v).strip() and str(v).lower() != 'nan']
            if parts:
                content += f"Строка {idx}: {' | '.join(parts)}\n"
        elif isinstance(row, list):
            parts = [str(v) for v in row if v and str(v).strip() and str(v).lower() != 'nan']
            if parts:
                content += f"Строка {idx}: {' | '.join(parts)}\n"
    
    return content

def load_json_documents(repo_id, hf_token, json_dir):
    """Load text sections from JSON (including ZIPs)"""
    import zipfile
    import tempfile
    import os
    
    log_message("Loading JSON documents...")
    
    files = list_repo_files(repo_id=repo_id, repo_type="dataset", token=hf_token)
    json_files = [f for f in files if f.startswith(json_dir) and f.endswith('.json')]
    zip_files = [f for f in files if f.startswith(json_dir) and f.endswith('.zip')]
    
    log_message(f"Found {len(json_files)} JSON files and {len(zip_files)} ZIP files")
    
    documents = []
    
    # Load direct JSON files
    for file_path in json_files:
        try:
            log_message(f"  Loading: {file_path}")
            local_path = hf_hub_download(
                repo_id=repo_id,
                filename=file_path,
                repo_type="dataset",
                token=hf_token
            )
            
            docs = extract_sections_from_json(local_path)
            documents.extend(docs)
            log_message(f"    ✓ Extracted {len(docs)} sections")
            
        except Exception as e:
            log_message(f"    ✗ Error loading {file_path}: {e}")
    
    # Extract and load ZIP files
    for zip_path in zip_files:
        try:
            log_message(f"  Processing ZIP: {zip_path}")
            local_zip = hf_hub_download(
                repo_id=repo_id,
                filename=zip_path,
                repo_type="dataset",
                token=hf_token
            )
            
            with zipfile.ZipFile(local_zip, 'r') as zf:
                json_files_in_zip = [f for f in zf.namelist() 
                                    if f.endswith('.json') 
                                    and not f.startswith('__MACOSX')
                                    and not f.startswith('.')]
                
                log_message(f"    Found {len(json_files_in_zip)} JSON files in ZIP")
                
                for json_file in json_files_in_zip:
                    try:
                        log_message(f"      Processing: {json_file}")
                        
                        # Read file content
                        file_content = zf.read(json_file)
                        
                        # Try to detect encoding
                        try:
                            # First try UTF-8
                            text_content = file_content.decode('utf-8')
                        except UnicodeDecodeError:
                            try:
                                # Try UTF-8 with BOM
                                text_content = file_content.decode('utf-8-sig')
                            except UnicodeDecodeError:
                                try:
                                    # Try Windows-1251 (common for Cyrillic)
                                    text_content = file_content.decode('windows-1251')
                                except UnicodeDecodeError:
                                    log_message(f"        ✗ Skipping: Cannot decode {json_file}")
                                    continue
                        
                        # Write to temp file
                        with tempfile.NamedTemporaryFile(mode='w', delete=False, 
                                                        suffix='.json', encoding='utf-8') as tmp:
                            tmp.write(text_content)
                            tmp_path = tmp.name
                        
                        # Extract sections
                        docs = extract_sections_from_json(tmp_path)
                        documents.extend(docs)
                        log_message(f"        ✓ Extracted {len(docs)} sections")
                        
                        # Clean up temp file
                        os.unlink(tmp_path)
                        
                    except json.JSONDecodeError as e:
                        log_message(f"        ✗ Invalid JSON in {json_file}: {e}")
                    except Exception as e:
                        log_message(f"        ✗ Error processing {json_file}: {e}")
                        
        except Exception as e:
            log_message(f"    ✗ Error with ZIP {zip_path}: {e}")
    
    log_message(f"✓ Loaded {len(documents)} text sections total")
    return documents

def extract_sections_from_json(json_path):
    """Extract sections from a single JSON file"""
    documents = []
    
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        doc_id = data.get('document_metadata', {}).get('document_id', 'unknown')
        
        # Extract all section levels
        for section in data.get('sections', []):
            if section.get('section_text', '').strip():
                documents.append(Document(
                    text=section['section_text'],
                    metadata={
                        'type': 'text',
                        'document_id': doc_id,
                        'section_id': section.get('section_id', '')
                    }
                ))
            
            # Subsections
            for subsection in section.get('subsections', []):
                if subsection.get('subsection_text', '').strip():
                    documents.append(Document(
                        text=subsection['subsection_text'],
                        metadata={
                            'type': 'text',
                            'document_id': doc_id,
                            'section_id': subsection.get('subsection_id', '')
                        }
                    ))
                
                # Sub-subsections
                for sub_sub in subsection.get('sub_subsections', []):
                    if sub_sub.get('sub_subsection_text', '').strip():
                        documents.append(Document(
                            text=sub_sub['sub_subsection_text'],
                            metadata={
                                'type': 'text',
                                'document_id': doc_id,
                                'section_id': sub_sub.get('sub_subsection_id', '')
                            }
                        ))
    
    except Exception as e:
        log_message(f"Error extracting from {json_path}: {e}")
    
    return documents


def load_table_documents(repo_id, hf_token, table_dir):
    """Load and chunk tables"""
    log_message("Loading tables...")
    
    files = list_repo_files(repo_id=repo_id, repo_type="dataset", token=hf_token)
    table_files = [f for f in files if f.startswith(table_dir) and f.endswith('.json')]
    
    all_chunks = []
    for file_path in table_files:
        try:
            local_path = hf_hub_download(
                repo_id=repo_id,
                filename=file_path,
                repo_type="dataset",
                token=hf_token
            )
            
            with open(local_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # Extract file-level document_id
            file_doc_id = data.get('document_id', data.get('document', 'unknown'))
            
            for sheet in data.get('sheets', []):
                # Use sheet-level document_id if available, otherwise use file-level
                sheet_doc_id = sheet.get('document_id', sheet.get('document', file_doc_id))
                
                # CRITICAL: Pass document_id to chunk function
                chunks = chunk_table_by_rows(sheet, sheet_doc_id)
                all_chunks.extend(chunks)
                
        except Exception as e:
            log_message(f"Error loading {file_path}: {e}")
    
    log_message(f"✓ Loaded {len(all_chunks)} table chunks")
    return all_chunks


def load_image_documents(repo_id, hf_token, image_dir):
    """Load image descriptions"""
    log_message("Loading images...")
    
    files = list_repo_files(repo_id=repo_id, repo_type="dataset", token=hf_token)
    csv_files = [f for f in files if f.startswith(image_dir) and f.endswith('.csv')]
    
    documents = []
    for file_path in csv_files:
        try:
            local_path = hf_hub_download(
                repo_id=repo_id,
                filename=file_path,
                repo_type="dataset",
                token=hf_token
            )
            
            df = pd.read_csv(local_path)
            
            for _, row in df.iterrows():
                content = f"Документ: {row.get('Обозначение документа', 'unknown')}\n"
                content += f"Рисунок: {row.get('№ Изображения', 'unknown')}\n"
                content += f"Название: {row.get('Название изображения', '')}\n"
                content += f"Описание: {row.get('Описание изображение', '')}\n"
                content += f"Раздел: {row.get('Раздел документа', '')}\n"
                
                documents.append(Document(
                    text=content,
                    metadata={
                        'type': 'image',
                        'document_id': str(row.get('Обозначение документа', 'unknown')),
                        'image_number': str(row.get('№ Изображения', 'unknown')),
                        'section': str(row.get('Раздел документа', ''))
                    }
                ))
        except Exception as e:
            log_message(f"Error loading {file_path}: {e}")
    
    log_message(f"✓ Loaded {len(documents)} images")
    return documents


def load_all_documents(repo_id, hf_token, json_dir, table_dir, image_dir):
    """Main loader - combines all document types"""
    log_message("="*60)
    log_message("STARTING DOCUMENT LOADING")
    log_message("="*60)
    
    # Load text sections
    text_docs = load_json_documents(repo_id, hf_token, json_dir)
    text_chunks = chunk_text_documents(text_docs)
    
    # Load tables (already chunked)
    table_chunks = load_table_documents(repo_id, hf_token, table_dir)
    
    # Load images (no chunking needed)
    image_docs = load_image_documents(repo_id, hf_token, image_dir)
    
    all_docs = text_chunks + table_chunks + image_docs
    
    log_message("="*60)
    log_message(f"TOTAL DOCUMENTS: {len(all_docs)}")
    log_message(f"  Text chunks: {len(text_chunks)}")
    log_message(f"  Table chunks: {len(table_chunks)}")
    log_message(f"  Images: {len(image_docs)}")
    log_message("="*60)
    
    return all_docs