"""
LLM-based Header Detection using Azure OpenAI GPT-4o.

This module uses GPT-4o to intelligently detect table headers from the first 10 rows,
handling multi-level hierarchical headers and returning structured header definitions.
"""

import json
from typing import List, Dict, Any, Optional
from dataclasses import dataclass, asdict

import structlog
from openai import AzureOpenAI

from src.core.config import settings
from src.llm import get_chat_client, get_chat_model_name, strip_thinking
from src.llm.factory import create_chat_completion

logger = structlog.get_logger(__name__)


# Prompt template for header detection
HEADER_DETECTION_PROMPT = """Analyze the first 10 rows of this table and identify the header row(s).

Table data (rows are 0-indexed):
{table_rows}

Tasks:
1. Identify which row(s) contain headers (0-indexed)
2. Detect if headers are multi-level (spanning multiple rows)
3. Return hierarchical header structure using "->" separator for levels

Rules:
- Use "->" to separate hierarchy levels (e.g., "Sales -> Q1" not "Sales > Q1")
- Return empty header_rows array [] if no clear headers detected
- Include ALL columns, even single-level ones
- For multi-level headers, combine parent and child with "->"
- If a cell is empty in a header row, inherit from the cell above it or mark as empty

Output JSON format:
{{
  "header_rows": [0, 1],
  "headers": [
    {{"column": 0, "hierarchy": ["Product"]}},
    {{"column": 1, "hierarchy": ["Sales", "Q1"]}},
    {{"column": 2, "hierarchy": ["Sales", "Q2"]}},
    {{"column": 3, "hierarchy": ["Inventory", "Current"]}}
  ]
}}

Example 1 - Single-level headers:
Row 0: ["Name", "Age", "City"]
Row 1: ["Alice", "25", "NYC"]

Output:
{{
  "header_rows": [0],
  "headers": [
    {{"column": 0, "hierarchy": ["Name"]}},
    {{"column": 1, "hierarchy": ["Age"]}},
    {{"column": 2, "hierarchy": ["City"]}}
  ]
}}

Example 2 - Multi-level headers:
Row 0: ["Product", "Sales", "", "Inventory", ""]
Row 1: ["Name", "Q1", "Q2", "Current", "Reserved"]
Row 2: ["Widget A", "100", "150", "50", "20"]

Output:
{{
  "header_rows": [0, 1],
  "headers": [
    {{"column": 0, "hierarchy": ["Product", "Name"]}},
    {{"column": 1, "hierarchy": ["Sales", "Q1"]}},
    {{"column": 2, "hierarchy": ["Sales", "Q2"]}},
    {{"column": 3, "hierarchy": ["Inventory", "Current"]}},
    {{"column": 4, "hierarchy": ["Inventory", "Reserved"]}}
  ]
}}

Now analyze the provided table and return ONLY valid JSON.
"""


@dataclass
class HeaderStructure:
    """Detected header structure from LLM analysis."""
    header_rows: List[int]  # Which rows are headers (0-indexed)
    headers: List[Dict[str, Any]]  # Column definitions with hierarchy

    # Example headers:
    # [
    #   {"column": 0, "hierarchy": ["Product"]},
    #   {"column": 1, "hierarchy": ["Sales", "Q1"]},
    #   {"column": 2, "hierarchy": ["Sales", "Q2"]}
    # ]

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return asdict(self)

    def get_flat_headers(self) -> List[str]:
        """
        Get flattened header strings using -> separator.

        Returns:
            List of header strings like ["Product", "Sales -> Q1", "Sales -> Q2"]
        """
        flat = []
        for header in self.headers:
            hierarchy = header.get('hierarchy', [])
            flat_name = " -> ".join(hierarchy) if hierarchy else ""
            flat.append(flat_name)
        return flat


class LlmHeaderDetector:
    """Uses Azure OpenAI GPT-4o to detect table headers from first 10 rows."""

    def __init__(self, client: Optional[AzureOpenAI] = None):
        """
        Initialize the LLM header detector.

        Args:
            client: Optional AzureOpenAI client. If None, creates one lazily on first use.
        """
        self._client = client
        self._client_initialized = client is not None

    def _get_client(self):
        """Get or create LLM client lazily (fork-safe)."""
        if not self._client_initialized:
            self._client = get_chat_client()
            self._client_initialized = True
        return self._client

    def detect_headers(
        self,
        table_html: str
    ) -> HeaderStructure:
        """
        Analyze HTML table and detect hierarchical headers using GPT-4o.

        Args:
            table_html: Raw HTML of the table element (first 10 rows)

        Returns:
            HeaderStructure with detected header rows and column definitions
        """
        try:
            logger.info(f"Detecting headers from HTML table using LLM...")

            # Extract first 10 rows from HTML
            from bs4 import BeautifulSoup
            soup = BeautifulSoup(table_html, 'lxml')
            table_elem = soup.find('table')

            if not table_elem:
                logger.error("No <table> element found in HTML")
                return self._create_fallback_headers_from_html(table_html)

            # Get first 10 <tr> rows
            rows = table_elem.find_all('tr')[:10]

            # PREPROCESSING: Fix misclassified <th> tags that contain data values
            # This happens when Docling marks formatted cells as headers
            rows = self._fix_misclassified_header_tags(rows)

            first_10_html = f"<table>{''.join(str(row) for row in rows)}</table>"

            # Build prompt with raw HTML
            prompt = f"""Analyze this HTML table and identify the header row(s).

HTML Table (first 10 rows):
{first_10_html}

Tasks:
1. Identify which row(s) contain headers (0-indexed)
2. Detect if headers are multi-level (spanning multiple rows using rowspan/colspan)
3. Return hierarchical header structure using " -> " separator for levels

Rules:
- Parse the HTML <th> and <td> elements
- Use rowspan and colspan attributes to understand merged cells
- Use " -> " to separate hierarchy levels (e.g., "Sales -> Q1" not "Sales > Q1")
- Return empty header_rows array [] if no clear headers detected
- Include ALL columns, even single-level ones
- For multi-level headers, combine parent and child with " -> "
- If a cell is empty or spans multiple columns, use the cell content only once

Output JSON format:
{{
  "header_rows": [0, 1],
  "headers": [
    {{"column": 0, "hierarchy": ["Product"]}},
    {{"column": 1, "hierarchy": ["Sales", "Q1"]}},
    {{"column": 2, "hierarchy": ["Sales", "Q2"]}},
    {{"column": 3, "hierarchy": ["Inventory", "Current"]}}
  ]
}}

Example - Multi-level headers with merged cells:
<table>
<tr><th colspan="1">Product</th><th colspan="2">Sales</th><th colspan="2">Inventory</th></tr>
<tr><th>Name</th><th>Q1</th><th>Q2</th><th>Current</th><th>Reserved</th></tr>
<tr><td>Widget A</td><td>100</td><td>150</td><td>50</td><td>20</td></tr>
</table>

Output:
{{
  "header_rows": [0, 1],
  "headers": [
    {{"column": 0, "hierarchy": ["Product", "Name"]}},
    {{"column": 1, "hierarchy": ["Sales", "Q1"]}},
    {{"column": 2, "hierarchy": ["Sales", "Q2"]}},
    {{"column": 3, "hierarchy": ["Inventory", "Current"]}},
    {{"column": 4, "hierarchy": ["Inventory", "Reserved"]}}
  ]
}}

Now analyze the provided table and return ONLY valid JSON.
"""

            # Debug: Log the actual HTML being sent to LLM
            logger.debug(f"First 10 rows HTML length: {len(first_10_html)} chars")
            logger.debug(f"First 500 chars of HTML: {first_10_html[:500]}")

            # Call LLM (get client lazily to avoid fork issues)
            logger.debug("Calling Azure OpenAI GPT-4o for header detection...")
            client = self._get_client()
            response = create_chat_completion(
                client,
                model=get_chat_model_name(),
                messages=[
                    {
                        "role": "system",
                        "content": "You are an Excel table header analyzer. Parse HTML tables and extract header structure. Respond only with valid JSON.",
                    },
                    {"role": "user", "content": prompt},
                ],
                temperature=0.1,  # Low temperature for consistent output
                max_tokens=2000,  # Increased for larger tables
                response_format={"type": "json_object"},
            )

            # Parse response
            content = strip_thinking(response.choices[0].message.content)
            if not content:
                logger.warning("llm_empty_after_think_strip")
                return self._create_fallback_headers_from_html(table_html)
            logger.debug(f"LLM response received: {len(content)} chars")
            logger.debug(f"LLM raw response: {content[:500]}...")  # First 500 chars

            result = json.loads(content)

            # Validate and create HeaderStructure
            header_rows = result.get('header_rows', [])
            headers = result.get('headers', [])

            # If LLM returns empty headers (no clear headers detected), use fallback
            if not headers:
                logger.warning(
                    "LLM returned empty headers (no clear header row detected). "
                    "Using fallback: generating column-based headers."
                )
                return self._create_column_based_headers(table_html)

            logger.info(
                f"Header detection complete: {len(header_rows)} header row(s), "
                f"{len(headers)} column(s)"
            )

            return HeaderStructure(
                header_rows=header_rows,
                headers=headers
            )

        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse LLM response as JSON: {e}")
            # Fallback: parse HTML to extract first row as headers
            return self._create_fallback_headers_from_html(table_html)

        except Exception as e:
            logger.error(f"Error during header detection: {e}")
            # Fallback: parse HTML to extract first row as headers
            return self._create_fallback_headers_from_html(table_html)

    def _fix_misclassified_header_tags(self, rows: list) -> list:
        """
        Fix misclassified <th> tags that contain data values instead of headers.

        Docling sometimes marks formatted cells (bold, merged) as <th> even when
        they contain data values. This method detects and converts them to <td>.

        Heuristics for detecting data values in <th> tags:
        1. Common data patterns: "－", "'否'", "'可'", "None", "null", numbers
        2. Single-character values (likely checkmarks or symbols)
        3. Rows where ALL <th> tags look like data

        Args:
            rows: List of BeautifulSoup <tr> elements

        Returns:
            List of <tr> elements with corrected <th>/<td> tags
        """
        if not rows:
            return rows

        # Data value patterns that shouldn't be in headers
        DATA_VALUE_PATTERNS = {
            '－',  # Dash
            '−',  # Minus
            "'否'",  # Japanese "No" with quotes
            "'可'",  # Japanese "Yes" with quotes
            "None",
            "null",
            "NULL",
            "",  # Empty
        }

        def looks_like_data(text: str) -> bool:
            """Check if text looks like a data value rather than a header."""
            text = text.strip()

            # Empty or whitespace-only
            if not text:
                return True

            # Known data value patterns
            if text in DATA_VALUE_PATTERNS:
                return True

            # Quoted values (often data)
            if (text.startswith("'") and text.endswith("'")) or \
               (text.startswith('"') and text.endswith('"')):
                return True

            # Single character (likely a symbol or checkmark)
            if len(text) == 1:
                return True

            # Pure numbers
            if text.replace('.', '').replace(',', '').replace('-', '').isdigit():
                return True

            return False

        def should_convert_row_to_td(row) -> bool:
            """Check if a row's <th> tags should all be converted to <td>."""
            th_tags = row.find_all('th')
            if not th_tags:
                return False

            # Count how many <th> tags look like data
            data_like_count = sum(
                1 for th in th_tags
                if looks_like_data(th.get_text(strip=True))
            )

            # If more than 50% of <th> tags look like data, convert entire row
            threshold = len(th_tags) * 0.5
            return data_like_count > threshold

        # Process each row
        fixed_rows = []
        for row_idx, row in enumerate(rows):
            if should_convert_row_to_td(row):
                logger.debug(
                    f"Row {row_idx}: Converting <th> tags to <td> (detected data values)"
                )

                # Convert all <th> to <td> in this row
                for th in row.find_all('th'):
                    th.name = 'td'

            fixed_rows.append(row)

        return fixed_rows

    def _create_column_based_headers(self, table_html: str) -> HeaderStructure:
        """
        Create column-based headers (Column 0, Column 1, etc.) when no headers detected.

        This is used when the table has no clear header row (all rows are data).

        Args:
            table_html: HTML table string

        Returns:
            HeaderStructure with column-index based headers
        """
        logger.info("Generating column-based headers (Column 0, Column 1, ...)")

        try:
            from bs4 import BeautifulSoup
            soup = BeautifulSoup(table_html, 'lxml')
            table_elem = soup.find('table')

            if not table_elem:
                return HeaderStructure(header_rows=[], headers=[])

            # Get first row to determine column count
            first_row = table_elem.find('tr')
            if not first_row:
                return HeaderStructure(header_rows=[], headers=[])

            # Count columns (accounting for colspan)
            cells = first_row.find_all(['th', 'td'])
            col_count = sum(int(cell.get('colspan', 1)) for cell in cells)

            # Generate column-based headers
            headers = []
            for col_idx in range(col_count):
                headers.append({
                    "column": col_idx,
                    "hierarchy": [f"Column {col_idx}"]
                })

            logger.info(f"Generated {len(headers)} column-based headers")

            return HeaderStructure(
                header_rows=[],  # No header rows
                headers=headers
            )
        except Exception as e:
            logger.error(f"Column-based header generation failed: {e}")
            return HeaderStructure(header_rows=[], headers=[])

    def _create_fallback_headers_from_html(
        self,
        table_html: str
    ) -> HeaderStructure:
        """
        Create simple fallback headers from HTML by parsing first row.

        Args:
            table_html: HTML table string

        Returns:
            HeaderStructure with simple single-level headers
        """
        logger.warning("Using fallback header detection (first row as headers)")

        try:
            from bs4 import BeautifulSoup
            soup = BeautifulSoup(table_html, 'lxml')
            table_elem = soup.find('table')

            if not table_elem:
                return HeaderStructure(header_rows=[], headers=[])

            # Get first row
            first_row = table_elem.find('tr')
            if not first_row:
                return HeaderStructure(header_rows=[], headers=[])

            # Extract cell text
            cells = first_row.find_all(['th', 'td'])
            headers = []

            for col_idx, cell in enumerate(cells):
                cell_text = cell.get_text(strip=True)
                header_name = cell_text if cell_text else f"Column {col_idx}"
                headers.append({
                    "column": col_idx,
                    "hierarchy": [header_name]
                })

            return HeaderStructure(
                header_rows=[0],
                headers=headers
            )
        except Exception as e:
            logger.error(f"Fallback header detection failed: {e}")
            return HeaderStructure(header_rows=[], headers=[])


# Convenience function for easy imports
def detect_headers_from_html(
    table_html: str,
    client: Optional[AzureOpenAI] = None
) -> HeaderStructure:
    """
    Convenience function to detect headers from HTML table.

    Args:
        table_html: Raw HTML of the table element
        client: Optional AzureOpenAI client

    Returns:
        HeaderStructure with detected headers
    """
    detector = LlmHeaderDetector(client=client)
    return detector.detect_headers(table_html)
