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

# New Modular Imports
from backend.agent.nodes.models_node import StylistProState
from backend.agent.nodes.classifier_node import classifier_node
from backend.agent.nodes.stylist_node import stylist_node
from backend.agent.nodes.utils import _extract_text

# StarRocks connection for Enrichment
from agent.tools.tool_module.search_engine import SELECT_COLUMNS, TABLE_NAME
from common.starrocks_connection import get_db_connection

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 StylistProGraph:
    """
    Architect Stylist Graph (Version 2.0).
    Modular Architecture: Classifier -> Stylist -> END
    """
    
    def __init__(self):
        self._compiled = None
        
    def build(self):
        if self._compiled:
            return self._compiled
            
        workflow = StateGraph(StylistProState)
        
        workflow.add_node("classifier", classifier_node)
        workflow.add_node("stylist", stylist_node)
        
        workflow.set_entry_point("classifier")
        
        workflow.add_conditional_edges(
            "classifier",
            route_after_classifier,
            {
                "stylist": "stylist",
                END: END
            }
        )
        
        workflow.add_edge("stylist", END)
        
        self._compiled = workflow.compile()
        logger.info("✅ StylistProGraph compiled (Modular)")
        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.
        """
        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": [],
        }

        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 ---
        final_products = []
        ai_product_ids = result.get("product_ids", [])
        tool_result_raw = result.get("tool_result")
        
        product_dict = {}
        if tool_result_raw:
            try:
                parsed = json.loads(tool_result_raw)
                if parsed.get("status") == "success":
                    all_products = parsed.get("products", [])
                    for p in all_products:
                        sku = str(p.get("magento_ref_code") or p.get("sku") or "").upper().strip()
                        if sku: 
                            if "sku" not in p: p["sku"] = sku
                            product_dict[sku] = p
                        
                        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

        # Enrich missing IDs 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:
                from agent.tools.tool_module.db_connector import DBConnector, strip_local_columns
                use_sqlite = config.get("configurable", {}).get("db_source") == "sqlite"
                db = DBConnector(use_sqlite=use_sqlite)
                
                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))
                safe_columns = strip_local_columns(SELECT_COLUMNS)
                sql = f"SELECT {safe_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(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"⚠️ Enrichment failed: {e}")

        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,
            "elapsed_ms": round(elapsed_ms, 2),
            "lead_stage": result.get("lead_stage"),
            "updated_insight": result.get("updated_insight"),
            "product_ids": ai_product_ids,
            "products": final_products,
            "tool_name_used": result.get("tool_name_used"),
            "tool_result": result.get("tool_result"),
            "pipeline": result.get("diagnostics", [])
        }

# Singleton
stylist_pro_graph = StylistProGraph()
