import logging
import json
import time
from typing import Dict, Any, List

from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from langchain_core.runnables import RunnableConfig

# Modular Imports from the refined nodes
from agent.models import AgentConfig, get_config
from agent.nodes.models_node import StylistProState
from agent.nodes.classifier_node import classifier_node
from agent.nodes.tools_node import tools_node
from agent.nodes.stylist_node import stylist_node
from agent.nodes.utils import _extract_text

# StarRocks connection for Enrichment
from common.starrocks_connection import get_db_connection
from config import USE_LOCAL_SQLITE

logger = logging.getLogger(__name__)

def route_after_classifier(state: StylistProState) -> str:
    """Điều hướng dựa trên early_exit."""
    if state.get("early_exit"):
        return END
    return "stylist"

class CANIFAGraph:
    """
    Fashion Q&A Agent Graph (Version 2.0).
    Architecture: Classifier (Lightweight) -> Stylist (Heavyweight) -> END
    
    All logic is modularized into nodes for maximum performance and readability.
    """
    
    def __init__(self, config: AgentConfig | None = None):
        self.config = config or get_config()
        self._compiled = None
        
    def build(self):
        if self._compiled:
            return self._compiled
            
        workflow = StateGraph(StylistProState)
        
        # Nodes
        workflow.add_node("classifier", classifier_node)
        workflow.add_node("tools", tools_node) # New dedicated tools node
        workflow.add_node("stylist", stylist_node)
        
        # Entry point
        workflow.set_entry_point("classifier")
        
        # Routing logic
        def route_after_classifier(state: StylistProState):
            if state.get("early_exit"):
                return END
            if state.get("tool_name_used"):
                return "tools"
            return "stylist"

        # Edges
        workflow.add_conditional_edges(
            "classifier",
            route_after_classifier,
            {
                "tools": "tools",
                "stylist": "stylist",
                END: END
            }
        )
        
        workflow.add_edge("tools", "stylist") # After tools, always go to stylist
        workflow.add_edge("stylist", END)
        
        self._compiled = workflow.compile()
        logger.info("✅ CANIFAGraph compiled (3-Agent Architecture: Classifier -> Tools -> Stylist)")
        return self._compiled

    async def chat(
        self,
        user_message: str,
        user_insight: str | None = None,
        history: List[BaseMessage] | None = None,
        config: RunnableConfig | None = None,
    ) -> Dict[str, Any]:
        """
        Main entry point with StarRocks Product Enrichment.
        Ensures all products mentioned by AI have full metadata (image, price, url).
        """
        start_time = time.time()
        graph = self.build()

        messages = history.copy() if history else []
        messages.append(HumanMessage(content=user_message))

        initial_state = {
            "messages": messages,
            "user_insight": user_insight,
            "tool_result": None,
            "tool_name_used": None,
            "updated_insight": None,
            "lead_stage": None,
            "stage_injection": None,
            "product_ids": [],
            "early_exit": False,
            "diagnostics": [],
        }

        # Invoke Graph
        result = await graph.ainvoke(initial_state, config=config)
        elapsed_ms = (time.time() - start_time) * 1000

        # Extract AI Response
        ai_response = ""
        for msg in reversed(result["messages"]):
            if isinstance(msg, AIMessage):
                ai_response = _extract_text(msg.content)
                break

        # --- PRODUCT ENRICHMENT LAYER (Refactored for Clean Base/Variant Mapping) ---
        final_products = []
        ai_product_ids = [pid.upper().strip() for pid in result.get("product_ids", []) if pid]
        tool_result_raw = result.get("tool_result")
        
        product_dict = {}
        base_sku_mapping = {}  # Map base_sku -> list of variant skus
        
        def _add_to_dicts(p: dict):
            sku = str(p.get("magento_ref_code") or p.get("sku") or p.get("sku_code") or "").upper().strip()
            if not sku:
                return
            
            # Ensure color fields exist
            if "color_code" not in p: p["color_code"] = p.get("product_color_code", "")
            if "color_name" not in p: p["color_name"] = p.get("product_color_name", "")
            if "sku" not in p: p["sku"] = sku
            
            product_dict[sku] = p
            
            # Also map the base SKU (before the dash)
            base_sku = sku.split("-")[0]
            if base_sku not in base_sku_mapping:
                base_sku_mapping[base_sku] = []
            if sku not in base_sku_mapping[base_sku]:
                base_sku_mapping[base_sku].append(sku)

        # 1. Map products already in tool result (if any)
        if tool_result_raw:
            try:
                parsed = json.loads(tool_result_raw)
                if parsed.get("status") == "success":
                    all_products = parsed.get("products", []) or parsed.get("results", [])
                    for p in all_products:
                        _add_to_dicts(p)
                        
                        # Also map outfit recommendations
                        recs = p.get("outfit_recommendations", [])
                        if isinstance(recs, list):
                            for r in recs:
                                _add_to_dicts(r)
            except Exception: pass

        # 2. (Removed) We no longer do explicit DB queries here. 
        # The tool_result should act as the single source of truth for the AI's context. 
        # Any SKUs that are not in product_dict are hallucinations and will be skipped.

        # 3. Assemble final products in the order requested by AI
        # If AI requested a base SKU, we return the first available variant for it
        seen = set()
        for pid in ai_product_ids:
            target_sku = None
            if pid in product_dict:
                target_sku = pid
            elif pid in base_sku_mapping and base_sku_mapping[pid]:
                target_sku = base_sku_mapping[pid][0]
                
            if target_sku and target_sku not in seen:
                final_products.append(product_dict[target_sku])
                seen.add(target_sku)

        return {
            "response": ai_response,
            "ai_response": ai_response, # compatibility
            "elapsed_ms": round(elapsed_ms, 2),
            "lead_stage": result.get("lead_stage"),
            "updated_insight": result.get("updated_insight"),
            "user_insight": result.get("updated_insight"), # compatibility
            "product_ids": ai_product_ids,
            "products": final_products,
            "pipeline": result.get("diagnostics", [])
        }

# --- Singleton & Compatibility API ---
_instance: List[CANIFAGraph | None] = [None]

def reset_chain_cache():
    """Clear the singleton instance to force a rebuild."""
    global _instance
    _instance[0] = None
    logger.info("🔄 Graph cache reset")

def reset_graph():
    """Alias for reset_chain_cache for compatibility."""
    reset_chain_cache()

def build_graph(config: AgentConfig | None = None) -> Any:
    """Compatibility function for controller."""
    global _instance
    if _instance[0] is None:
        _instance[0] = CANIFAGraph(config)
    return _instance[0].build()

def get_graph_manager(config: AgentConfig | None = None) -> CANIFAGraph:
    global _instance
    if _instance[0] is None:
        _instance[0] = CANIFAGraph(config)
    return _instance[0]
