#!/usr/bin/env python3
"""
Phase 2 — Semantic Normalization / Identity Extraction
Compliments Reference DB Pipeline

Input: Phase 1 output (phase1_output.parquet)
Output: Phase 2 output (phase2_output.parquet)

This phase performs:
  1. Brand normalization → brand_norm + product_line (uses brand_clean from Phase 1)
  2. Identity attribute extraction from ORIGINAL title (BEFORE normalization)
  3. Fat-level and fat-percentage extraction
  4. Flavour extraction
  5. Formulation extraction
  6. Variant/size attribute parsing
  7. Core-title extraction (strip brand prefix + size)

CRITICAL: Identity information is extracted from the ORIGINAL title
before any destructive normalization removes it.

Phase 1 (data quality) has already cleaned:
  - brand → brand_clean (whitespace/case normalization)
  - title → title_clean (whitespace normalization)
  - upc → upc_raw (preserved, not altered)
  - external_id → external_id_raw (preserved, not altered)

Phase 2 uses brand_clean for semantic normalization, and original title
for identity extraction.
"""

import json
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

import pandas as pd
import numpy as np

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
BASE_DIR = Path(__file__).resolve().parent.parent
INPUT_PATH = BASE_DIR.parent / "phase1" / "outputs" / "phase1_output.parquet"
OUTPUT_DIR = BASE_DIR / "outputs"
VALIDATION_DIR = BASE_DIR / "validation"
STATISTICS_DIR = BASE_DIR / "statistics"

VERSION = "2.0.0"
TIMESTAMP = datetime.now(timezone.utc).isoformat()


def log(msg: str) -> None:
    print(f"[Phase2] {msg}")


# ---------------------------------------------------------------------------
# 1. Brand Normalization
# ---------------------------------------------------------------------------
# Maps 11 raw brand variants → (brand_norm, product_line)
# Case-insensitive matching

BRAND_MAP = {
    "compliments organic":       ("Compliments", "Organic"),
    "compliments balance":       ("Compliments", "Balance"),
    "compliments naturally simple": ("Compliments", "Naturally Simple"),
    "compliments green care":    ("Compliments", "Green"),
    "compliments green":         ("Compliments", "Green"),
    "compliments little ones":   ("Compliments", "Little Ones"),
    "sensations":                ("Sensations", "Sensations"),
    "compliments":               ("Compliments", "Core"),
    "compliments ":              ("Compliments", "Core"),   # trailing space
    " compliments":              ("Compliments", "Core"),   # leading space
}


def normalize_brand(clean_brand: str) -> tuple[str, str]:
    """Map cleaned brand string to (brand_norm, product_line).
    
    Phase 1 has already cleaned the brand field (whitespace trim, case normalization).
    This function performs SEMANTIC normalization: mapping brand variants to canonical
    brand + product_line.
    """
    if pd.isna(clean_brand):
        return ("Unknown", "Core")
    key = str(clean_brand).strip().lower()
    if key in BRAND_MAP:
        return BRAND_MAP[key]
    # Default: if it contains "compliments", treat as Core Compliments
    if "compliments" in key:
        return ("Compliments", "Core")
    if "sensations" in key:
        return ("Sensations", "Sensations")
    return (str(clean_brand).strip(), "Core")


# ---------------------------------------------------------------------------
# 2. Identity Attribute Extraction
# ---------------------------------------------------------------------------
# Extracted from the ORIGINAL title BEFORE any normalization

def extract_identity_flags(title: str) -> dict:
    """Extract 9 boolean identity flags from raw product title."""
    t = title.lower()
    return {
        "is_organic": bool(re.search(r"\borganic\b", t)),
        "is_gluten_free": bool(re.search(r"\bgluten[\s-]+free\b", t)),
        "is_naturally_simple": bool(re.search(r"\bnaturally\s+simple\b", t)),
        "is_sugar_free": bool(re.search(
            r"\bsugar[\s-]+free\b|\b(?:no sugar added|unsweetened)\b|\bzero\s+sugar\b", t)),
        "is_unsalted": bool(re.search(r"\bunsalted\b|\bno salt\b", t)),
        "is_lactose_free": bool(re.search(r"\blactose[\s-]+free\b", t)),
        "is_peanut_free": bool(re.search(r"\bpeanut[\s-]+free\b", t)),
        "is_plant_based": bool(re.search(r"\bplant[\s-]*based\b", t)),
        "is_reduced_sodium": bool(re.search(
            r"\breduced\s+sodium\b|\blow\s+sodium\b|\bno\s+salt\s+added\b", t)),
    }


# ---------------------------------------------------------------------------
# 3. Fat Level and Fat Percentage Extraction
# ---------------------------------------------------------------------------

def extract_fat_info(title: str) -> dict:
    """Extract fat_level and fat_percentage from raw product title."""
    t = title.lower()

    # Fat level
    fat_level = "regular"
    if re.search(r"\b(?:light|lite|reduced fat|low fat|lean)\b", t):
        fat_level = "reduced_fat"
    elif re.search(r"\bfat[\s-]+free\b", t):
        fat_level = "fat_free"

    # Fat percentage — extract numeric % from title
    fat_percentage = None
    pct_match = re.search(r"(\d+(?:\.\d+)?)\s*%", title)
    if pct_match:
        val = float(pct_match.group(1))
        # Skip values that are clearly not fat percentages
        # (cocoa %, alcohol %, etc.)
        skip_context = any(w in t for w in ["cocoa", "alcohol", "isopropyl", "peanuts"])
        if val < 100 and not skip_context:
            fat_percentage = val
            # Override fat_level if percentage is very low
            if 0 < val <= 0.7:
                fat_level = "fat_free"

    return {
        "fat_level": fat_level,
        "fat_percentage": fat_percentage,
    }


# ---------------------------------------------------------------------------
# 4. Flavour Extraction
# ---------------------------------------------------------------------------

FLAVOUR_KEYWORDS = [
    "almond", "apple", "banana", "blueberry", "caramel", "cherry",
    "chocolate", "cinnamon", "coconut", "cranberry", "honey", "lemon",
    "lime", "mango", "maple", "mixed berry", "peach", "peanut",
    "peppermint", "pineapple", "pomegranate", "raspberry", "strawberry",
    "tropical", "vanilla", "watermelon", "white chocolate", "berry",
    "espresso", "caramel", "butterscotch", "toffee",
]


def extract_flavours(title: str) -> list[str]:
    """Extract flavour keywords from raw product title."""
    t = title.lower()
    found = []
    for kw in FLAVOUR_KEYWORDS:
        if kw in t:
            found.append(kw)
    return sorted(set(found))


# ---------------------------------------------------------------------------
# 5. Formulation Extraction
# ---------------------------------------------------------------------------

FORMULATION_KEYWORDS = [
    "smooth", "crunchy", "creamy", "chunky", "whole", "halves",
    "sliced", "ground", "chopped", "breaded", "fresh", "frozen",
    "roasted", "smoked",
]


def extract_formulation(title: str) -> list[str]:
    """Extract formulation keywords from raw product title."""
    t = title.lower()
    found = []
    for kw in FORMULATION_KEYWORDS:
        if kw in t:
            found.append(kw)
    return sorted(set(found))


# ---------------------------------------------------------------------------
# 6. Variant Attributes Parsing
# ---------------------------------------------------------------------------

def build_variant_attributes(size_str: str) -> dict:
    """Parse raw size string into structured variant attributes."""
    if not size_str or pd.isna(size_str):
        return {}

    s = str(size_str).strip()
    attrs = {}

    # Try to match: [qty x ] amount unit [x multiplier]
    m = re.match(
        r"(\d+(?:\.\d+)?)\s*x\s*"  # qty x
        r"(\d+(?:\.\d+)?)\s*"       # amount
        r"(g|kg|ml|l|oz|lb)s?\s*$",  # unit
        s, re.IGNORECASE
    )
    if m:
        attrs["qty"] = int(float(m.group(1)))
        attrs["amount"] = float(m.group(2))
        attrs["unit"] = m.group(3).lower()
        return attrs

    # Try: amount unit [x multiplier]
    m = re.match(
        r"(\d+(?:\.\d+)?)\s*"        # amount
        r"(g|kg|ml|l|oz|lb)s?\s*"     # unit
        r"(?:x\s*(\d+))?\s*$",        # optional x multiplier
        s, re.IGNORECASE
    )
    if m:
        attrs["amount"] = float(m.group(1))
        attrs["unit"] = m.group(2).lower()
        if m.group(3):
            attrs["multiplier"] = int(m.group(3))
        return attrs

    # Try: count-based (e.g. "12 per pack", "20 count")
    m = re.match(
        r"(\d+)\s*(?:per\s+pack|count|ea|pack|piece|slice|cups?|pound)s?\s*$",
        s, re.IGNORECASE
    )
    if m:
        attrs["count"] = int(m.group(1))
        return attrs

    # Fallback: store raw string
    attrs["raw"] = s
    return attrs


# ---------------------------------------------------------------------------
# 7. Core Title Extraction
# ---------------------------------------------------------------------------

# Brand/product-line prefixes to strip from title
BRAND_PREFIXES = [
    ("Compliments Naturally Simple ", "Compliments "),
    ("Compliments Balance ", "Compliments Balance "),
    ("Compliments Organic ", "Compliments Organic "),
    ("Compliments Green Care ", "Compliments Green "),
    ("Compliments Little Ones ", "Compliments Little Ones "),
    ("Compliments ", ""),
    ("Sensations ", ""),
]

# Trailing size/weight/count patterns to strip
SIZE_PATTERNS = [
    r"\s+\d[\d,.]*\s*(?:g|kg|ml|l|oz|lb|count|ea|pack|piece|slice|cups?|litre|liters?|pound)s?\s*$",
    r"\s+\d+\s*x\s+\d+\s*(?:g|kg|ml|l)\s*$",
    r"\s+\d+\s*x\s*$",
    r"\s+\d+\s+per\s+pack\s*$",
    r"\s+\d+\s+count\s*$",
]

# Bracket content to strip
BRACKET_PATTERN = r"\s*\([^)]*\)\s*"


def extract_core_title(title: str, product_line: str) -> str:
    """Extract core product title by stripping brand prefix and size info."""
    t = title.strip()

    # Step 1: Strip brand/product-line prefix
    for prefix, replacement in BRAND_PREFIXES:
        if t.lower().startswith(prefix.lower()):
            t = replacement + t[len(prefix):]
            break

    # Step 2: Strip trailing size/weight/count patterns
    for pat in SIZE_PATTERNS:
        t = re.sub(pat, "", t, flags=re.IGNORECASE)

    # Step 3: Strip bracket content
    t = re.sub(BRACKET_PATTERN, " ", t)

    # Step 4: Clean up whitespace
    t = re.sub(r"\s+", " ", t).strip()

    return t


# ---------------------------------------------------------------------------
# 8. Identity Hash (for reference — grouping happens in Phase 3)
# ---------------------------------------------------------------------------

def build_identity_hash(row: dict) -> str:
    """Build deterministic identity hash from extracted attributes."""
    parts = []
    for col in ["is_organic", "is_gluten_free", "is_naturally_simple",
                 "is_sugar_free", "is_unsalted", "is_lactose_free",
                 "is_peanut_free", "is_plant_based", "is_reduced_sodium"]:
        parts.append("1" if row.get(col) else "0")

    fp = row.get("fat_percentage")
    if pd.notna(fp) and fp is not None:
        parts.append(f"fat{fp:.2g}" if fp else "fat0")
    else:
        parts.append("fat_none")

    parts.append(str(row.get("fat_level", "regular")))
    parts.append(str(row.get("product_line", "Core")))

    fl = row.get("flavour", [])
    parts.append(",".join(sorted(fl)) if fl else "")

    fm = row.get("formulation", [])
    parts.append(",".join(sorted(fm)) if fm else "")

    return "|".join(parts)


# ---------------------------------------------------------------------------
# Main Processing
# ---------------------------------------------------------------------------

def process_phase2(df: pd.DataFrame) -> pd.DataFrame:
    """Apply all Phase 2 transformations."""
    log(f"Processing {len(df)} products ...")

    # CRITICAL: Extract identity from ORIGINAL title BEFORE any normalization
    log("Step 1: Extracting identity flags from original titles ...")
    identity_flags = df["title"].apply(extract_identity_flags)
    id_df = pd.DataFrame(identity_flags.tolist())
    for col in id_df.columns:
        df[col] = id_df[col].values

    log("Step 2: Extracting fat info ...")
    fat_info = df["title"].apply(extract_fat_info)
    fat_df = pd.DataFrame(fat_info.tolist())
    df["fat_level"] = fat_df["fat_level"].values
    df["fat_percentage"] = fat_df["fat_percentage"].values

    log("Step 3: Extracting flavours ...")
    df["flavour"] = df["title"].apply(extract_flavours)

    log("Step 4: Extracting formulations ...")
    df["formulation"] = df["title"].apply(extract_formulation)

    log("Step 5: Normalizing brands (using brand_clean from Phase 1) ...")
    brand_col = "brand_clean" if "brand_clean" in df.columns else "brand"
    brand_results = df[brand_col].apply(normalize_brand)
    df["brand_norm"] = brand_results.apply(lambda x: x[0])
    df["product_line"] = brand_results.apply(lambda x: x[1])

    log("Step 6: Parsing variant attributes ...")
    df["variant_attributes"] = df["size"].apply(
        lambda x: json.dumps(build_variant_attributes(x)) if pd.notna(x) else "{}"
    )

    log("Step 7: Extracting core titles ...")
    df["core_title"] = df.apply(
        lambda row: extract_core_title(row["title"], row["product_line"]),
        axis=1
    )

    log("Step 8: Building identity hash ...")
    df["identity_hash"] = df.apply(lambda row: build_identity_hash(row.to_dict()), axis=1)

    log(f"Processing complete. Output shape: {df.shape}")
    return df


# ---------------------------------------------------------------------------
# Validation
# ---------------------------------------------------------------------------

def validate_phase2(df_in: pd.DataFrame, df_out: pd.DataFrame) -> dict:
    """Validate Phase 2 output."""
    checks = {}

    # Row count preserved
    checks["row_count"] = {
        "input": len(df_in),
        "output": len(df_out),
        "pass": len(df_in) == len(df_out),
    }

    # All original columns preserved
    original_cols = set(df_in.columns)
    output_cols = set(df_out.columns)
    missing_original = original_cols - output_cols
    checks["original_columns_preserved"] = {
        "missing": list(missing_original),
        "pass": len(missing_original) == 0,
    }

    # New columns added
    new_cols = output_cols - original_cols
    expected_new = {
        "brand_norm", "product_line",
        "is_organic", "is_gluten_free", "is_naturally_simple",
        "is_sugar_free", "is_unsalted", "is_lactose_free",
        "is_peanut_free", "is_plant_based", "is_reduced_sodium",
        "fat_level", "fat_percentage",
        "flavour", "formulation",
        "variant_attributes",
        "core_title", "identity_hash",
    }
    missing_new = expected_new - new_cols
    checks["new_columns_added"] = {
        "expected": sorted(expected_new),
        "actual": sorted(new_cols),
        "missing": list(missing_new),
        "pass": len(missing_new) == 0,
    }

    # Brand normalization
    checks["brand_norm"] = {
        "unique_values": sorted(df_out["brand_norm"].unique().tolist()),
        "pass": True,
    }

    # Product line
    checks["product_line"] = {
        "unique_values": sorted(df_out["product_line"].unique().tolist()),
        "pass": True,
    }

    # Fat level values
    valid_fat_levels = {"fat_free", "reduced_fat", "regular"}
    actual_fat_levels = set(df_out["fat_level"].unique())
    invalid_fat = actual_fat_levels - valid_fat_levels
    checks["fat_level"] = {
        "valid_values": sorted(valid_fat_levels),
        "actual_values": sorted(actual_fat_levels),
        "invalid": sorted(invalid_fat),
        "pass": len(invalid_fat) == 0,
    }

    # Identity flags are boolean
    flag_cols = ["is_organic", "is_gluten_free", "is_naturally_simple",
                 "is_sugar_free", "is_unsalted", "is_lactose_free",
                 "is_peanut_free", "is_plant_based", "is_reduced_sodium"]
    all_bool = all(df_out[col].dtype == bool for col in flag_cols)
    checks["identity_flags_boolean"] = {"pass": all_bool}

    # core_title not empty
    empty_core = (df_out["core_title"].str.strip() == "").sum()
    checks["core_title"] = {
        "empty_count": int(empty_core),
        "pass": empty_core == 0,
    }

    # identity_hash not empty
    empty_hash = (df_out["identity_hash"].str.strip() == "").sum()
    checks["identity_hash"] = {
        "empty_count": int(empty_hash),
        "pass": empty_hash == 0,
    }

    # No duplicate external_ids
    ext_dupes = df_out["external_id"].duplicated().sum()
    checks["external_id_uniqueness"] = {
        "duplicate_count": int(ext_dupes),
        "pass": ext_dupes == 0,
    }

    # Overall
    all_pass = all(c.get("pass", True) for c in checks.values())
    checks["overall"] = {"result": "PASS" if all_pass else "FAIL"}

    return checks


# ---------------------------------------------------------------------------
# Statistics
# ---------------------------------------------------------------------------

def build_statistics(df_in: pd.DataFrame, df_out: pd.DataFrame) -> dict:
    """Build Phase 2 statistics."""
    return {
        "version": VERSION,
        "timestamp": TIMESTAMP,
        "input": {
            "source": "phase1_output.parquet",
            "row_count": len(df_in),
            "column_count": len(df_in.columns),
        },
        "output": {
            "row_count": len(df_out),
            "column_count": len(df_out.columns),
            "columns_added": sorted(set(df_out.columns) - set(df_in.columns)),
        },
        "brand_normalization": {
            "input_brand_column": "brand_clean" if "brand_clean" in df_in.columns else "brand",
            "raw_brands": {str(k): int(v) for k, v in df_in.get("brand_clean", df_in["brand"]).value_counts().items()},
            "normalized_brands": {str(k): int(v) for k, v in df_out["brand_norm"].value_counts().items()},
            "product_lines": {str(k): int(v) for k, v in df_out["product_line"].value_counts().items()},
        },
        "identity_flags": {
            col: int(df_out[col].sum()) for col in [
                "is_organic", "is_gluten_free", "is_naturally_simple",
                "is_sugar_free", "is_unsalted", "is_lactose_free",
                "is_peanut_free", "is_plant_based", "is_reduced_sodium",
            ]
        },
        "fat_info": {
            "fat_level_distribution": {str(k): int(v) for k, v in df_out["fat_level"].value_counts().items()},
            "fat_percentage_nulls": int(df_out["fat_percentage"].isna().sum()),
            "fat_percentage_non_null": int(df_out["fat_percentage"].notna().sum()),
        },
        "flavour": {
            "products_with_flavour": int((df_out["flavour"].apply(len) > 0).sum()),
            "products_without_flavour": int((df_out["flavour"].apply(len) == 0).sum()),
        },
        "formulation": {
            "products_with_formulation": int((df_out["formulation"].apply(len) > 0).sum()),
            "products_without_formulation": int((df_out["formulation"].apply(len) == 0).sum()),
        },
        "core_title": {
            "unique_count": int(df_out["core_title"].nunique()),
        },
        "identity_hash": {
            "unique_count": int(df_out["identity_hash"].nunique()),
        },
    }


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    log("Starting Phase 2")

    # Ensure output dirs exist
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    VALIDATION_DIR.mkdir(parents=True, exist_ok=True)
    STATISTICS_DIR.mkdir(parents=True, exist_ok=True)

    # Load Phase 1 output
    log(f"Loading Phase 1 output from: {INPUT_PATH}")
    df_in = pd.read_parquet(INPUT_PATH)
    log(f"Loaded: {df_in.shape[0]} rows, {df_in.shape[1]} columns")

    # Process
    df_out = process_phase2(df_in.copy())

    # Validate
    log("Validating output ...")
    validation = validate_phase2(df_in, df_out)

    # Statistics
    log("Building statistics ...")
    statistics = build_statistics(df_in, df_out)

    # Save outputs
    log("Saving outputs ...")
    df_out.to_parquet(OUTPUT_DIR / "phase2_output.parquet", index=False)
    log(f"  Saved phase2_output.parquet ({df_out.shape[0]} rows, {df_out.shape[1]} cols)")

    with open(VALIDATION_DIR / "phase2_validation.json", "w") as f:
        json.dump(validation, f, indent=2, default=str)
    log("  Saved phase2_validation.json")

    with open(STATISTICS_DIR / "phase2_statistics.json", "w") as f:
        json.dump(statistics, f, indent=2, default=str)
    log("  Saved phase2_statistics.json")

    # Summary
    log("")
    log("=== PHASE 2 COMPLETE ===")
    log(f"Input:  {df_in.shape[0]} rows, {df_in.shape[1]} columns")
    log(f"Output: {df_out.shape[0]} rows, {df_out.shape[1]} columns")
    log(f"Columns added: {sorted(set(df_out.columns) - set(df_in.columns))}")
    log(f"Validation: {validation['overall']['result']}")
    log("========================")

    return df_out, validation, statistics


if __name__ == "__main__":
    main()