"""
Record Builder for creating structured JSON records from tables with headers.

This module maps hierarchical headers to cell values and generates JSON records
with source metadata for tracking.
"""

from typing import Any
from uuid import UUID

import structlog

from src.extraction_v2.html_table_extractor import HtmlTable
from src.extraction_v2.llm_header_detector import HeaderStructure
from src.llm.factory import get_token_counter

logger = structlog.get_logger(__name__)

# Default max tokens for embedding (8191 is the common 8K-context ceiling).
DEFAULT_MAX_TOKENS = 8000

# Minimum tokens threshold - tables smaller than this may be merged with consecutive tables
MIN_TOKENS_THRESHOLD = 100


class RecordBuilder:
    """Builds structured records from HTML tables with detected headers."""

    def __init__(self):
        """Initialize the record builder."""
        pass

    def build_records(
        self,
        table: HtmlTable,
        headers: HeaderStructure,
        document_id: UUID,
        sheet_name: str | None = None,
        color_map: dict | None = None
    ) -> list[dict[str, Any]]:
        """
        Convert table rows to JSON records.

        Args:
            table: HtmlTable with extracted table data
            headers: HeaderStructure with detected headers
            document_id: Document UUID for source tracking
            sheet_name: Optional sheet name (overrides table.sheet_name)

        Returns:
            List of record dictionaries with hierarchical header keys
        """
        try:
            logger.info(
                f"Building records from table: {table.row_count} rows, "
                f"{table.col_count} cols"
            )

            # Use provided sheet_name or fall back to table's sheet_name
            final_sheet_name = sheet_name or table.sheet_name or "Unknown"

            # Determine which rows are data (skip header rows)
            if not headers.header_rows:
                # No headers detected, treat all rows as data
                data_start_row = 0
            else:
                # Start after the last header row
                data_start_row = max(headers.header_rows) + 1

            logger.debug(
                f"Data starts at row {data_start_row} "
                f"(header rows: {headers.header_rows})"
            )

            # Get flattened header strings
            flat_headers = headers.get_flat_headers()

            # Skip tables with generic column-based headers (table fragments without real headers)
            if self._has_only_column_based_headers(flat_headers):
                logger.info(
                    "Skipping table with only generic column headers "
                    "(table fragment without proper headers)"
                )
                return []

            # Build records from data rows
            records = []
            for row_idx in range(data_start_row, len(table.rows)):
                record = self._build_record(
                    row_data=table.rows[row_idx],
                    flat_headers=flat_headers,
                    document_id=document_id,
                    sheet_name=final_sheet_name,
                    row_number=row_idx,
                    table_index=table.table_index,
                    color_map=color_map
                )

                # Only include non-empty records
                if self._is_record_non_empty(record):
                    records.append(record)

            logger.info(
                f"Built {len(records)} records "
                f"(skipped {table.row_count - data_start_row - len(records)} "
                f"empty rows)"
            )

            return records

        except Exception as e:
            logger.error(f"Error building records: {e}")
            return []

    def _build_record(
        self,
        row_data: list[str],
        flat_headers: list[str],
        document_id: UUID,
        sheet_name: str,
        row_number: int,
        table_index: int,
        color_map: dict | None = None
    ) -> dict[str, Any]:
        """
        Build a single record from a row of data.

        Normalized merged cell handling:
        - Only stores non-empty values from the first column of merged cells
        - Empty strings from colspan/rowspan expansion are skipped
        - This creates cleaner records without redundant empty values

        Args:
            row_data: List of cell values for this row (normalized from merged cells)
            flat_headers: List of flattened header strings
            document_id: Document UUID
            sheet_name: Sheet name
            row_number: Row number (0-indexed)
            table_index: Table index in document

        Returns:
            Record dictionary with content and metadata
        """
        # Create content dictionary mapping headers to values
        content = {}
        for col_idx, cell_value in enumerate(row_data):
            if col_idx < len(flat_headers):
                header_key = flat_headers[col_idx]

                # NORMALIZATION: Only store non-empty values
                # This skips empty strings from merged cell colspan/rowspan expansion
                # Benefits:
                # - Cleaner records without redundant empty columns
                # - Reduced storage size
                # - Easier to identify actual data vs merged cell artifacts
                if cell_value:  # Skip empty strings and None
                    content[header_key] = cell_value
            else:
                # More columns than headers (shouldn't happen, but handle it)
                if cell_value:  # Also skip empties for overflow columns
                    content[f"Column {col_idx}"] = cell_value

        # Create record with content and source metadata
        record = {
            "content": content,
            "_source": {
                "document_id": str(document_id),
                "sheet": sheet_name,
                "row": row_number,
                "table_id": table_index
            }
        }

        # Add cell colors if available (memory optimized - only this row's colors)
        if color_map:
            cell_colors = {}
            sheet_colors = color_map.get(sheet_name, {})

            for col_idx, cell_value in enumerate(row_data):
                if col_idx < len(flat_headers):
                    header_key = flat_headers[col_idx]
                    color_key = f"{row_number},{col_idx}"

                    if color_key in sheet_colors:
                        # Copy color dict (not reference) to avoid shared state
                        cell_colors[header_key] = dict(sheet_colors[color_key])

            # Only add if non-empty
            if cell_colors:
                record["_source"]["cell_colors"] = cell_colors

        return record

    def build_raw_text_records(
        self,
        table: HtmlTable,
        document_id: UUID,
        sheet_name: str | None = None
    ) -> list[dict[str, Any]]:
        """
        Build records from single-column raw text tables.

        Concatenates all row text into a single record with 'text' field.
        Sets row_number = -1 to distinguish from tabular data records.
        No LLM header detection is needed for raw text.

        Args:
            table: HtmlTable with single-column raw text content
            document_id: Document UUID for source tracking
            sheet_name: Optional sheet name (overrides table.sheet_name)

        Returns:
            List containing single record with concatenated text, or empty list if no content
        """
        try:
            logger.info(
                f"Building raw text record from table {table.table_index}: "
                f"{table.row_count} rows"
            )

            # Use provided sheet_name or fall back to table's sheet_name
            final_sheet_name = sheet_name or table.sheet_name or "Unknown"

            # Concatenate all row text (filter empty rows)
            text_lines = []
            for row in table.rows:
                # Get first column text (single-column table)
                if row and row[0]:
                    text_content = row[0].strip()
                    if text_content:
                        text_lines.append(text_content)

            # Skip empty tables
            if not text_lines:
                logger.debug(
                    f"Skipping empty raw text table {table.table_index}"
                )
                return []

            # Concatenate with newlines
            concatenated_text = "\n".join(text_lines)

            # Create single record with concatenated text
            record = {
                "content": {
                    "text": concatenated_text
                },
                "_source": {
                    "document_id": str(document_id),
                    "sheet": final_sheet_name,
                    "row": -1,  # Distinguish raw text from table rows
                    "table_id": table.table_index
                }
            }

            logger.info(
                f"Built raw text record: {len(text_lines)} lines, "
                f"{len(concatenated_text)} chars"
            )

            return [record]

        except Exception as e:
            logger.error(f"Error building raw text record: {e}")
            return []

    def build_combined_raw_text_records(
        self,
        tables: list[HtmlTable],
        document_id: UUID,
        sheet_name: str,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        min_tokens: int = MIN_TOKENS_THRESHOLD,
        color_map: dict | None = None
    ) -> list[dict[str, Any]]:
        """
        Build records from raw text tables, merging small consecutive tables.

        Logic:
        1. Each raw text table becomes one record by default
        2. If a table has < min_tokens, try to merge with next consecutive table
        3. Stop merging when combined exceeds max_tokens

        Args:
            tables: List of HtmlTable objects (single-column raw text tables)
            document_id: Document UUID for source tracking
            sheet_name: Sheet name for all tables
            max_tokens: Maximum tokens per record (default: 8000 for embedding)
            min_tokens: Minimum tokens threshold for merging (default: 100)

        Returns:
            List of records, with small consecutive tables merged
        """
        try:
            if not tables:
                return []

            # Sort tables by table_index to maintain document order
            sorted_tables = sorted(tables, key=lambda t: t.table_index)

            logger.info(
                f"Building raw text records from {len(sorted_tables)} tables "
                f"on sheet '{sheet_name}' (min_tokens={min_tokens}, max_tokens={max_tokens})"
            )

            # Provider-aware tokenizer (Azure -> tiktoken cl100k_base; vLLM -> HF tokenizer)
            count_tokens = get_token_counter()

            def extract_table_text(table: HtmlTable) -> tuple[list[str], int]:
                """Extract text lines and token count from a table."""
                texts = []
                for row in table.rows:
                    if row and row[0]:
                        text_content = row[0].strip()
                        if text_content:
                            texts.append(text_content)
                combined = "\n".join(texts)
                tokens = count_tokens(combined) if texts else 0
                return texts, tokens

            def split_texts_by_tokens(texts: list[str], max_tok: int) -> list[list[str]]:
                """Split a list of text lines into chunks that fit within max_tok."""
                if not texts:
                    return []

                chunks = []
                current_chunk = []
                current_tokens = 0

                for line in texts:
                    line_tokens = count_tokens(line)

                    # If single line exceeds max_tok, split the line itself
                    if line_tokens > max_tok:
                        # Flush current chunk first
                        if current_chunk:
                            chunks.append(current_chunk)
                            current_chunk = []
                            current_tokens = 0

                        # Split long line by words
                        words = line.split()
                        line_chunk = []
                        line_chunk_tokens = 0
                        for word in words:
                            word_tokens = count_tokens(word + ' ')
                            if line_chunk_tokens + word_tokens > max_tok:
                                if line_chunk:
                                    chunks.append([' '.join(line_chunk)])
                                line_chunk = [word]
                                line_chunk_tokens = word_tokens
                            else:
                                line_chunk.append(word)
                                line_chunk_tokens += word_tokens
                        if line_chunk:
                            chunks.append([' '.join(line_chunk)])
                        continue

                    # Check if adding this line would exceed limit
                    if current_tokens + line_tokens + 1 > max_tok:  # +1 for newline
                        if current_chunk:
                            chunks.append(current_chunk)
                        current_chunk = [line]
                        current_tokens = line_tokens
                    else:
                        current_chunk.append(line)
                        current_tokens += line_tokens + 1

                # Add remaining lines
                if current_chunk:
                    chunks.append(current_chunk)

                return chunks

            def create_record(texts: list[str], table_ids: list[int]) -> dict[str, Any]:
                """Create a record from accumulated text."""
                concatenated = "\n".join(texts)
                return {
                    "content": {
                        "text": concatenated
                    },
                    "_source": {
                        "document_id": str(document_id),
                        "sheet": sheet_name,
                        "row": -1,
                        "table_ids": table_ids
                    }
                }

            records = []
            current_texts = []
            current_table_ids = []
            current_tokens = 0
            last_table_index = None

            for table in sorted_tables:
                table_texts, table_tokens = extract_table_text(table)

                if not table_texts:
                    last_table_index = table.table_index
                    continue

                # If this single table exceeds max_tokens, split it into multiple records
                if table_tokens > max_tokens:
                    # Flush any accumulated text first
                    if current_texts:
                        records.append(create_record(current_texts, current_table_ids))
                        logger.debug(
                            f"Created raw text record: {current_tokens} tokens, "
                            f"tables {current_table_ids}"
                        )
                        current_texts = []
                        current_table_ids = []
                        current_tokens = 0

                    # Split this large table into chunks
                    text_chunks = split_texts_by_tokens(table_texts, max_tokens)
                    logger.info(
                        f"Split large table {table.table_index} ({table_tokens} tokens) "
                        f"into {len(text_chunks)} chunks"
                    )
                    for chunk_idx, chunk_texts in enumerate(text_chunks):
                        chunk_record = {
                            "content": {
                                "text": "\n".join(chunk_texts)
                            },
                            "_source": {
                                "document_id": str(document_id),
                                "sheet": sheet_name,
                                "row": -1,
                                "table_id": table.table_index,
                                "chunk_index": chunk_idx,
                                "total_chunks": len(text_chunks),
                            }
                        }
                        records.append(chunk_record)
                    last_table_index = table.table_index
                    continue

                # Check if this table is consecutive with the previous
                is_consecutive = (
                    last_table_index is not None and
                    table.table_index - last_table_index == 1
                )

                # Determine if we should try to merge
                # Merge only if:
                # 1. Current accumulated text is small (< min_tokens), AND
                # 2. Tables are consecutive, AND
                # 3. Combined won't exceed max_tokens
                should_merge = (
                    current_texts and
                    is_consecutive and
                    current_tokens < min_tokens and
                    current_tokens + table_tokens + 1 <= max_tokens  # +1 for newline
                )

                if current_texts and not should_merge:
                    # Flush current record before starting new one
                    records.append(create_record(current_texts, current_table_ids))
                    logger.debug(
                        f"Created raw text record: {current_tokens} tokens, "
                        f"tables {current_table_ids}"
                    )
                    current_texts = []
                    current_table_ids = []
                    current_tokens = 0

                # Add this table's text
                current_texts.extend(table_texts)
                current_table_ids.append(table.table_index)
                current_tokens = count_tokens("\n".join(current_texts))
                last_table_index = table.table_index

            # Flush any remaining text
            if current_texts:
                records.append(create_record(current_texts, current_table_ids))
                logger.debug(
                    f"Created raw text record: {current_tokens} tokens, "
                    f"tables {current_table_ids}"
                )

            logger.info(
                f"Built {len(records)} raw text record(s) from "
                f"{len(sorted_tables)} tables on sheet '{sheet_name}'"
            )

            return records

        except Exception as e:
            logger.error(f"Error building raw text records: {e}")
            return []

    def _has_only_column_based_headers(self, flat_headers: list[str]) -> bool:
        """
        Check if all headers are generic column-based (Column 0, Column 1, etc.).

        These are table fragments without proper headers and should be skipped.

        Args:
            flat_headers: List of flattened header strings

        Returns:
            True if all headers match "Column N" pattern
        """
        if not flat_headers:
            return False

        # Check if all headers match "Column N" pattern
        column_pattern_count = sum(
            1 for header in flat_headers
            if header.startswith("Column ") and header[7:].isdigit()
        )

        # If more than 80% are "Column N" headers, skip this table
        threshold = len(flat_headers) * 0.8
        return column_pattern_count > threshold

    def _is_record_non_empty(self, record: dict[str, Any]) -> bool:
        """
        Check if a record has any non-empty content.

        Args:
            record: Record dictionary

        Returns:
            True if record has at least one non-empty cell
        """
        content = record.get("content", {})

        # Check if any value is non-empty after stripping whitespace
        for value in content.values():
            if isinstance(value, str) and value.strip():
                return True
            elif value is not None and not isinstance(value, str):
                # Non-string, non-None value (e.g., number)
                return True

        return False


# Convenience function for easy imports
def build_records_from_table(
    table: HtmlTable,
    headers: HeaderStructure,
    document_id: UUID,
    sheet_name: str | None = None
) -> list[dict[str, Any]]:
    """
    Convenience function to build records from a table.

    Args:
        table: HtmlTable with extracted data
        headers: HeaderStructure with detected headers
        document_id: Document UUID
        sheet_name: Optional sheet name

    Returns:
        List of record dictionaries
    """
    builder = RecordBuilder()
    return builder.build_records(table, headers, document_id, sheet_name)


def build_raw_text_records_from_table(
    table: HtmlTable,
    document_id: UUID,
    sheet_name: str | None = None
) -> list[dict[str, Any]]:
    """
    Convenience function to build raw text records from a single-column table.

    Args:
        table: HtmlTable with single-column raw text content
        document_id: Document UUID
        sheet_name: Optional sheet name

    Returns:
        List containing single record with concatenated text
    """
    builder = RecordBuilder()
    return builder.build_raw_text_records(table, document_id, sheet_name)
