from __future__ import annotations
from typing import Dict, List, Optional, Tuple
import re
from bs4 import BeautifulSoup
from shared.utils import normalize_text
from ..registry import extractor
from ..helpers.desc import (
    split_description_container, find_description_container,
    pair_title_content_from_magento_tabs, scan_headings_for_sections,
)

# ----- value/unit parser ------------------------------------------------------

_NUM_UNIT_RE = re.compile(
    r"""
    ^\s*
    (?P<num>[-+]?\d{1,3}(?:[.,]\d{3})*(?:[.,]\d+)?|\d+(?:[.,]\d+)?)
    \s*
    (?P<unit>[a-zA-Z%µ/]+)?
    \s*$
    """,
    re.X,
)

def _parse_value_unit(s: str) -> Tuple[Optional[str], Optional[str]]:
    if not s:
        return None, None
    s = re.sub(r"\s+", " ", s.strip())
    m = _NUM_UNIT_RE.match(s)
    if not m:
        return None, None
    num = (m.group("num") or "").replace(",", "")
    unit = m.group("unit") or None
    if unit:
        u = unit.lower()
        if u in {"kcal", "kcal.", "kcalories", "kcalorie"}:
            unit = "kcal"
        elif u in {"kj", "kj.", "kilojoule", "kilojoules"}:
            unit = "kJ"
    return (num or None, unit)

# ----- section finder ---------------------------------------------------------

def _find_nutrition_section_html(soup: BeautifulSoup) -> Optional[str]:
    """
    Return the HTML for the section whose title matches 'Nutritional Information'.
    We look in the same places your description extractor does.
    """
    # 1) Magento tabs
    for t, html in (pair_title_content_from_magento_tabs(soup) or []):
        if not t or not html:
            continue
        title = normalize_text(t).rstrip(":").lower()
        if "nutritional information" in title:
            return html

    # 2) Description container split into sections
    desc_el = find_description_container(soup)
    if desc_el:
        _open_html, sections = split_description_container(desc_el)
        for sec in sections or []:
            title = normalize_text((sec.get("title") or "")).rstrip(":").lower()
            if "nutritional information" in title:
                return sec.get("html") or ""

    # 3) Fallback: generic heading scan
    for t, html in (scan_headings_for_sections(soup) or []):
        if not t or not html:
            continue
        title = normalize_text(t).rstrip(":").lower()
        if "nutritional information" in title:
            return html

    return None

# ----- table parser -----------------------------------------------------------

def _extract_rows_from_table(root: BeautifulSoup) -> List[Dict[str, str]]:
    out: List[Dict[str, str]] = []
    table = root.select_one("table")
    if not table:
        return out

    for tr in table.select("tr"):
        th = tr.find("th")
        tds = tr.find_all("td")
        if th and tds:
            key = normalize_text(th.get_text(" ").strip())
            val_raw = normalize_text(tds[0].get_text(" ").strip())
        elif len(tds) >= 2:
            key = normalize_text(tds[0].get_text(" ").strip())
            val_raw = normalize_text(tds[1].get_text(" ").strip())
        else:
            continue

        if not key or not val_raw:
            continue

        value, unit = _parse_value_unit(val_raw)
        if value is None:  # keep raw if not parseable
            value, unit = val_raw, None

        out.append({"key": key, "value": value, "unit": unit})

    # Deduplicate while preserving order
    seen = set()
    dedup: List[Dict[str, str]] = []
    for r in out:
        t = (r["key"], r.get("value"), r.get("unit"))
        if t in seen:
            continue
        seen.add(t)
        dedup.append(r)
    return dedup

# ----- extractor --------------------------------------------------------------

@extractor
def ex_nutrition(soup: BeautifulSoup, url: str) -> Dict:
    """
    Extract nutrition ONLY from the section titled 'Nutritional Information'.
    Returns: {"nutrition": [{"key": "...", "value": "...", "unit": "..."}]}
    """
    section_html = _find_nutrition_section_html(soup)
    if not section_html:
        return {"nutrition": []}
    section_soup = BeautifulSoup(section_html, "lxml")
    rows = _extract_rows_from_table(section_soup)
    return {"nutrition": rows}
