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

logger = logging.getLogger(__name__)

# Constants for Product Enrichment (Hardcoded for clarity and decoupling)
TABLE_NAME = "shared_source.magento_product_dimension_with_text_embedding"
SELECT_COLUMNS = """
    magento_ref_code, internal_ref_code, product_name, 
    sale_price, original_price, product_image_url_thumbnail, 
    product_web_url, size_scale, gender_by_product, product_line_vn
"""

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 ---
        final_products = []
        ai_product_ids = result.get("product_ids", [])
        tool_result_raw = result.get("tool_result")
        
        product_dict = {}
        
        # 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:
                        sku = str(p.get("magento_ref_code") or p.get("sku") or p.get("sku_code") or "").upper().strip()
                        if sku: 
                            if "sku" not in p: p["sku"] = sku
                            product_dict[sku] = p
                        
                        # Also map outfit recommendations
                        recs = p.get("outfit_recommendations", [])
                        if isinstance(recs, list):
                            for r in recs:
                                r_sku = str(r.get("match_product_code") or r.get("sku") or "").upper().strip()
                                if r_sku:
                                    if "sku" not in r: r["sku"] = r_sku
                                    product_dict[r_sku] = r
            except Exception: pass

        # 2. Enrich missing or incomplete product data via StarRocks
        wanted_ids = [pid.upper().strip() for pid in ai_product_ids if pid]
        missing_ids = [pid for pid in wanted_ids if pid not in product_dict or not product_dict[pid].get("image")]

        if missing_ids:
            try:
                db = get_db_connection()
                skus = list(dict.fromkeys(missing_ids))
                base_skus = list(dict.fromkeys([s.split("-")[0] for s in skus]))
                keys = list(dict.fromkeys(skus + base_skus))
                
                placeholders = ",".join(["%s"] * len(keys))
                sql = f"""
                SELECT {SELECT_COLUMNS} 
                FROM {TABLE_NAME} 
                WHERE UPPER(magento_ref_code) IN ({placeholders}) 
                   OR UPPER(internal_ref_code) IN ({placeholders}) 
                LIMIT 200
                """
                rows = await db.execute_query_async(sql, params=tuple(keys + keys))

                for r in rows or []:
                    card = {
                        "sku": (r.get("magento_ref_code") or r.get("internal_ref_code") or "").strip(),
                        "name": r.get("product_name", ""),
                        "price": int(r.get("sale_price") or 0),
                        "original_price": int(r.get("original_price") or 0),
                        "image": r.get("product_image_url_thumbnail", ""),
                        "url": r.get("product_web_url", ""),
                        "sizes": r.get("size_scale", ""),
                        "gender": r.get("gender_by_product", ""),
                        "product_line": r.get("product_line_vn", ""),
                    }
                    product_dict[card["sku"].upper()] = card
            except Exception as e:
                logger.warning(f"⚠️ StarRocks Enrichment failed: {e}")

        # 3. Assemble final products in the order requested by AI
        seen = set()
        for pid in wanted_ids:
            if pid in product_dict and pid not in seen:
                final_products.append(product_dict[pid])
                seen.add(pid)

        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]
