"""
Controller Helper Functions
Extracted from controller.py for better code organization.
"""

import json
import logging
import re
import time

from common.cache import redis_cache

logger = logging.getLogger(__name__)


async def save_user_insight_to_redis(identity_key: str, insight: str) -> None:
    """
    Background task to save user insight to Redis (non-blocking).
    TTL is set to expire at 00:00 Asia/Ho_Chi_Minh to reset daily.

    Args:
        identity_key: User/device identifier
        insight: JSON string of user insight data
    """
    try:
        from datetime import datetime, timedelta
        from zoneinfo import ZoneInfo
        
        client = redis_cache.get_client()
        if client:
            insight_key = f"identity_key_insight:{identity_key}"
            
            # Calculate TTL to expire at 00:00 Vietnam time
            tz = ZoneInfo("Asia/Ho_Chi_Minh")
            now = datetime.now(tz)
            reset_time = now.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1)
            ttl_seconds = int((reset_time - now).total_seconds())
            
            # Set with TTL to auto-reset at midnight
            await client.set(insight_key, insight, ex=ttl_seconds)
            logger.info(f"💾 Updated User Insight for {identity_key} (TTL: {ttl_seconds}s): {insight[:50]}...")
    except Exception as e:
        logger.error(f"❌ Failed to save user insight: {e}")


async def extract_and_save_user_insight(json_content: str, identity_key: str) -> dict | None:
    """
    Background task: Extract user_insight từ partial/full JSON và save to Redis.

    Args:
        json_content: AI response JSON string
        identity_key: User/device identifier

    Returns:
        dict | None: user_insight dict nếu extract thành công, None nếu không tìm thấy
    """
    start_time = time.time()
    logger.info(f"🔄 [Background] Starting user_insight extraction for {identity_key}")

    try:
        # Normalize double braces → single (LLM sometimes outputs {{ }} per prompt instruction)
        normalized = json_content.replace("{{", "{").replace("}}", "}")

        # --- Robust balanced-brace extraction ---
        # Find "user_insight" key position
        key_index = normalized.find('"user_insight"')
        if key_index == -1:
            logger.warning(f"⚠️ [Background] No user_insight key found for {identity_key}")
            return None

        # Find opening { after the key
        brace_start = normalized.find('{', key_index + len('"user_insight"'))
        if brace_start == -1:
            logger.warning(f"⚠️ [Background] No opening brace for user_insight for {identity_key}")
            return None

        # Count balanced braces to find matching }
        depth = 0
        in_string = False
        escape_next = False
        brace_end = -1

        for i in range(brace_start, len(normalized)):
            ch = normalized[i]
            if escape_next:
                escape_next = False
                continue
            if ch == '\\':
                escape_next = True
                continue
            if ch == '"' and not escape_next:
                in_string = not in_string
                continue
            if in_string:
                continue
            if ch == '{':
                depth += 1
            elif ch == '}':
                depth -= 1
                if depth == 0:
                    brace_end = i
                    break

        if brace_end == -1:
            logger.warning(f"⚠️ [Background] Unbalanced braces in user_insight for {identity_key}")
            return None

        insight_json_str = normalized[brace_start:brace_end + 1]

        try:
            insight_dict = json.loads(insight_json_str)
        except json.JSONDecodeError as je:
            logger.warning(f"⚠️ [Background] user_insight JSON parse error: {je} | raw: {insight_json_str[:100]}...")
            return None

        insight_str = json.dumps(insight_dict, ensure_ascii=False, indent=2)

        # Save to Redis
        await save_user_insight_to_redis(identity_key, insight_str)
        elapsed = time.time() - start_time
        logger.warning(f"✅ [user_insight] Extracted + saved in {elapsed:.2f}s | Key: {identity_key}")
        return insight_dict

    except Exception as e:
        logger.error(f"❌ [Background] Failed to extract user_insight for {identity_key}: {e}")
        return None


async def load_user_insight_from_redis(identity_key: str) -> str | None:
    """
    Load user insight from Redis cache.

    Args:
        identity_key: User/device identifier

    Returns:
        str | None: User insight JSON string or None if not found
    """
    try:
        client = redis_cache.get_client()
        if client:
            insight_key = f"identity_key_insight:{identity_key}"
            user_insight = await client.get(insight_key)
            if user_insight:
                logger.info(f"🧠 Loaded User Insight for {identity_key}: {user_insight[:50]}...")
                return user_insight
    except Exception as e:
        logger.error(f"❌ Error fetching user insight: {e}")

    return None


async def finalize_user_insight_after_stream(
    stream_task,
    streaming_callback,
    identity_key: str,
    trace_id: str = "",
) -> None:
    """
    Background task: Wait for stream to complete, then extract and save user_insight.
    NOTE: Runs outside Langfuse trace context — no need to log to Langfuse.
    """
    # Clear Langfuse observation context to avoid auto-tracing in background
    try:
        from langfuse._context import _observation_stack_context
        _observation_stack_context.set([])
    except Exception:
        pass  # Langfuse version may not have this; safe to ignore

    try:
        # Wait for stream to complete (may have already finished or errored)
        try:
            await stream_task
        except Exception as e:
            logger.warning(f"⚠️ [Background] stream_task had error: {e}")
            # Still try to extract insight from partial content

        # Extract user insight from accumulated content
        full_content = streaming_callback.accumulated_content
        if full_content and identity_key:
            await extract_and_save_user_insight(full_content, identity_key)
        else:
            logger.warning("⚠️ [Background] No content to extract user_insight from")

    except Exception as exc:
        logger.error(f"❌ [Background] user_insight finalize failed: {exc}")


def extract_products_from_state(final_state_snapshot: dict) -> list[dict]:
    """
    Extract full product objects from LangGraph state snapshot.

    Args:
        final_state_snapshot: Last event from graph.astream()

    Returns:
        list[dict]: List of product objects with name, price, sale_price, etc.
    """
    from .helper import extract_product_ids

    enriched_products = []

    try:
        # Get messages from any node's state update
        all_messages = []
        for node_name, node_state in final_state_snapshot.items():
            if isinstance(node_state, dict) and "messages" in node_state:
                all_messages.extend(node_state["messages"])

        if all_messages:
            enriched_products = extract_product_ids(all_messages)
            logger.info(f"✅ Extracted {len(enriched_products)} full products from tool messages")

            # Debug: Log sample product
            if enriched_products:
                sample = enriched_products[0]
                logger.debug(
                    f"🔍 Sample product: name={sample.get('name')}, price={sample.get('price')}, sku={sample.get('sku')}"
                )

    except Exception as e:
        logger.error(f"❌ Error extracting products from state: {e}")

    return enriched_products
