"""
Lead Stage Graph — 2-Agent LangGraph Architecture cho phân tích khách hàng.

Flow: User query → Classifier AI → Stylist AI ⇄ Tools → answer

2 con AI:
  - Classifier (Planner): Phân tích ý định, đánh giá lead stage (Awareness, Interest, vv) → output JSON.
  - Stylist (Responder): Nhận stage injection từ Classifier, sinh ra câu trả lời tư vấn phù hợp.
    Có thể gọi tag_search_tool để tìm kiếm sản phẩm khi cần.

Trang thí nghiệm — KHÔNG ảnh hưởng production flow.
"""

import json
import logging
import time
from typing import Annotated, Any, TypedDict

from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.runnables import RunnableConfig
from langgraph.graph import END, StateGraph, START
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode

from common.llm_factory import create_llm
from config import DEFAULT_MODEL

from .prompts import LEAD_STAGE_CLASSIFIER_PROMPT, format_stage_injection
from .lead_search_tool import tag_search_tool

# --- Nhập thêm các tools từ agent/tools ---
from agent.tools.brand_knowledge_tool import canifa_knowledge_search
from agent.tools.promotion_canifa_tool import canifa_get_promotions
from agent.tools.store_search_tool import canifa_store_search
from agent.tools.check_is_stock import check_is_stock
from agent.tools.customer_info_tool import collect_customer_info
from agent.tools.add_to_cart_tool import add_to_cart

logger = logging.getLogger(__name__)


def _extract_text(content) -> str:
    """Extract plain text from msg.content — handles both str and Gemini list format."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                parts.append(item.get("text", ""))
            elif isinstance(item, str):
                parts.append(item)
        return "\n".join(parts) if parts else str(content)
    return str(content or "")

# ═══════════════════════════════════════════════
# State
# ═══════════════════════════════════════════════
class LeadStageState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    
    # Context data
    user_insight: str | None
    chat_history_summary: str | None
    
    # Internal routing & stage data
    lead_stage: dict | None
    stage_injection: str | None
    diagnostics: Annotated[list[dict[str, Any]], lambda x, y: x + y]


# ═══════════════════════════════════════════════
# Graph Builder — 2-Agent Architecture + Tools
# ═══════════════════════════════════════════════
class LeadStageGraph:
    """2-Agent LangGraph: Classifier → Stylist ⇄ Tools."""

    def __init__(self, model_name: str | None = None):
        # Dùng model_name nếu có, không thì mặc định theo DEFAULT_MODEL (từ config)
        self.model_name = model_name or DEFAULT_MODEL
        self.tools = [
            tag_search_tool,
            canifa_knowledge_search,
            canifa_get_promotions,
            canifa_store_search,
            check_is_stock,
            collect_customer_info,
            add_to_cart
        ]

        # AI 1: Classifier (buộc dùng JSON mode để lấy state info)
        self.classifier_llm = create_llm(
            model_name=self.model_name, 
            streaming=False, 
            json_mode=True
        )
        
        # AI 2: Stylist (streaming, có tool binding để tìm sản phẩm)
        self.stylist_llm = create_llm(
            model_name=self.model_name, 
            streaming=True
        )
        self.stylist_with_tools = self.stylist_llm.bind_tools(self.tools)

        self._compiled = None
        logger.info(f"✅ LeadStageGraph 2-Agent + Tools initialized with model: {self.model_name}")

    # ─────────────────────────────────────
    # Node 1: CLASSIFIER — Phân tích stage
    # ─────────────────────────────────────
    async def _classifier_node(self, state: LeadStageState, config: RunnableConfig | None = None) -> dict:
        """Classifier AI: phân tích câu hỏi, trả về JSON đại diện Lead Stage."""
        messages = state["messages"]
        user_query = _extract_text(messages[-1].content) if messages else ""
        user_insight = state.get("user_insight") or "Chưa có thông tin."
        history_summary = state.get("chat_history_summary") or "Chưa có lịch sử."

        classifier_input = (
            f"=== TIN NHẮN HIỆN TẠI ===\n{user_query}\n\n"
            f"=== USER INSIGHT (tích lũy từ trước) ===\n{user_insight}\n\n"
            f"=== TÓM TẮT LỊCH SỬ CHAT ===\n{history_summary}"
        )

        sys_msg = SystemMessage(content=LEAD_STAGE_CLASSIFIER_PROMPT)
        human_msg = HumanMessage(content=classifier_input)

        start = time.time()
        try:
            response = await self.classifier_llm.ainvoke([sys_msg, human_msg], config=config)
            content = _extract_text(response.content)
            
            # Parse JSON
            stage_result = json.loads(content)
            elapsed_ms = (time.time() - start) * 1000

            logger.info(
                f"🎯 Lead Stage: {stage_result.get('stage_name')} "
                f"(stage={stage_result.get('stage')}, conf={stage_result.get('confidence')}) "
                f"| {elapsed_ms:.0f}ms"
            )

            injection = format_stage_injection(stage_result)
            
            diag = {
                "step": "classifier",
                "label": "🧠 Classifier AI",
                "content": f"Stage {stage_result.get('stage')}: {stage_result.get('stage_name')} (Conf: {stage_result.get('confidence')})\nLý do: {stage_result.get('reasoning')}",
                "raw_json": json.dumps(stage_result, ensure_ascii=False, indent=2),
                "elapsed_ms": round(elapsed_ms)
            }

            return {
                "lead_stage": stage_result,
                "stage_injection": injection,
                "diagnostics": [diag]
            }

        except Exception as e:
            elapsed_ms = (time.time() - start) * 1000
            logger.error(f"❌ Lead Stage Classifier failed: {e}")
            
            fallback = {
                "stage": 2,
                "stage_name": "INTEREST",
                "confidence": 0.3,
                "reasoning": f"Classifier error: {str(e)[:100]}. Fallback to INTEREST.",
                "tone_directive": "Consultant",
                "behavioral_hints": ["Hỏi thêm thông tin để hiểu nhu cầu khách."],
            }
            injection = format_stage_injection(fallback)

            diag = {
                "step": "classifier",
                "label": "🧠 Classifier AI (Error Fallback)",
                "content": f"Lỗi parse JSON: {e}",
                "raw_json": json.dumps(fallback, ensure_ascii=False, indent=2),
                "elapsed_ms": round(elapsed_ms)
            }

            return {
                "lead_stage": fallback,
                "stage_injection": injection,
                "diagnostics": [diag]
            }

    # ─────────────────────────────────────
    # Node 2: STYLIST — Format câu trả lời (có tools)
    # ─────────────────────────────────────
    async def _stylist_node(self, state: LeadStageState, config: RunnableConfig | None = None) -> dict:
        """Stylist AI: nhận stage injection, đưa ra câu trả lời tư vấn. Có thể gọi tool."""
        messages = state["messages"]
        injection = state.get("stage_injection", "")

        # Stylist System Prompt with FULL tool awareness
        stylist_sys_prompt = f"""\
Bạn là Trợ lý phong cách (AI Stylist) độc quyền của CANIFA — thương hiệu thời trang hàng đầu Việt Nam.
Nhiệm vụ của bạn không chỉ là tìm đồ mà còn là KHƠI GỢI NHU CẦU và TƯ VẤN THỜI TRANG khéo léo.
Quy tắc Tư Vấn "Khéo Léo":
- Luôn khen ngợi hoặc đồng tình với gu thẩm mỹ/mong muốn của khách.
- Khi giới thiệu sản phẩm: Gợi ý cách phối đồ (mix & match), mô tả cảm giác khi mặc (chất vải thoáng mát, form dáng tôn dáng...).
- Không bao giờ trả lời cộc lốc một danh sách sản phẩm. Hãy nói như một người tư vấn tại shop đang trò chuyện cùng bạn của mình.

=== HƯỚNG DẪN DÙNG CÔNG CỤ (TOOLS) BẮT BUỘC ===
Bạn có các công cụ để truy xuất thông tin thực tế. Hãy chọn công cụ phù hợp với câu hỏi của khách:

1. 🛍️ TÌM KIẾM SẢN PHẨM: Khi khách hỏi quần áo, váy vóc, đồ đi làm, mặc nhà...
   → Dùng tool `tag_search_tool`. Đợi kết quả công cụ trả về mới bắt đầu viết câu tư vấn sản phẩm khéo léo.
   → KHÔNG BAO GIỜ tự bịa ra sản phẩm.

2. 📦 KIỂM TRA TỒN KHO: Khi khách hỏi size M mẫu này còn không, sản phẩm SKU 12345 còn hàng ở đâu...
   → Dùng tool `check_is_stock`. Đầu vào: mã SKU, color (mã), size.

3. 🏬 TÌM CỬA HÀNG: Khi khách hỏi Cửa hàng ở Hà Đông, Ở Đà Nẵng có shop nào...
   → Dùng tool `canifa_store_search`. Đầu vào: tỉnh/thành phố hoặc khu vực.

4. ❓ CSKH & CHÍNH SÁCH: Khi khách hỏi phí ship, đổi trả đồ, đăng ký thẻ VIP, điều kiện lên hạng...
   → Dùng tool `canifa_knowledge_search`. Câu trả lời sẽ có trong tài liệu quy định nội bộ của Canifa.

5. 🎁 KHUYẾN MÃI: Khi khách hỏi nay có sale gì, chương trình sinh nhật...
   → Dùng tool `canifa_get_promotions`. Trả về tin tức event/sale hiện tại.

6. 📝 LẤY THÔNG TIN LEAD: Khi khách muốn đăng ký tư vấn, nhận ưu đãi đặc biệt qua SDT/Email, hoặc gửi khiếu nại cần gọi lại...
   → Dùng tool `collect_customer_info`. Yêu cầu khách cung cấp tên và số điện thoại trước khi gọi if thiếu.

7. 🛒 GIỎ HÀNG (ADD TO CART): Khi khách đòi mua, chốt đơn mua một sản phẩm cụ thể (đã biết mã, màu, size)...
   → Dùng tool `add_to_cart`. Đảm bảo phải có mã sản phẩm, size, màu. Yêu cầu khách chọn size/màu nếu chưa ghi rõ.

Nếu khách khen/chê hoặc chỉ chào hỏi bình thường → Trả lời trực tiếp bằng phong thái khéo léo, bỏ qua việc gọi tool nếu không cần tra cứu.

{injection}
"""
        
        # Inject responder system prompt
        stylist_messages = [SystemMessage(content=stylist_sys_prompt)]
        for msg in messages:
            if not isinstance(msg, SystemMessage):
                stylist_messages.append(msg)

        start = time.time()
        response = await self.stylist_with_tools.ainvoke(stylist_messages, config=config)
        elapsed_ms = (time.time() - start) * 1000

        response_text = _extract_text(response.content)
        has_tool_calls = hasattr(response, "tool_calls") and response.tool_calls

        diag = {
            "step": "stylist" if not has_tool_calls else "stylist_tool_call",
            "label": "💬 Stylist AI" if not has_tool_calls else "🔧 Stylist AI → Tool Call",
            "elapsed_ms": round(elapsed_ms)
        }

        if has_tool_calls:
            diag["content"] = f"Gọi {len(response.tool_calls)} tool(s): {', '.join(tc['name'] for tc in response.tool_calls)}"
            diag["tool_calls"] = [
                {"name": tc["name"], "args": tc["args"]}
                for tc in response.tool_calls
            ]
        else:
            diag["content"] = response_text[:500] + ("..." if len(response_text) > 500 else "")

        # Trả về message để add vào list, và diagnostics để track timeline
        return {
            "messages": [response],
            "diagnostics": [diag]
        }

    # ─────────────────────────────────────
    # Router: Stylist → tool hoặc end
    # ─────────────────────────────────────
    def _after_stylist(self, state: LeadStageState) -> str:
        """Sau Stylist: nếu có tool_calls → đi tools, nếu không → end."""
        last = state["messages"][-1]
        if hasattr(last, "tool_calls") and last.tool_calls:
            return "tools"
        return "end"

    def build(self):
        """Build 2-agent + tools graph."""
        if self._compiled:
            return self._compiled

        workflow = StateGraph(LeadStageState)

        # 3 Nodes: Classifier, Stylist, Tools
        workflow.add_node("classifier", self._classifier_node)
        workflow.add_node("stylist", self._stylist_node)
        workflow.add_node("tools", ToolNode(self.tools))

        # Flow: Entry → Classifier → Stylist
        workflow.set_entry_point("classifier")
        workflow.add_edge("classifier", "stylist")

        # Stylist → (tools | end)
        workflow.add_conditional_edges(
            "stylist",
            self._after_stylist,
            {"tools": "tools", "end": END},
        )

        # Tools → Stylist (loop lại để Stylist format kết quả)
        workflow.add_edge("tools", "stylist")

        self._compiled = workflow.compile()
        logger.info("✅ LeadStageGraph 2-Agent + Tools compiled")
        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:
        """
        Main entry point — gửi message và nhận response.
        Returns pipeline trace for UI display.
        """
        start = time.time()
        graph = self.build()

        messages: list[BaseMessage] = []
        if history:
            messages.extend(history)
        messages.append(HumanMessage(content=user_message))

        initial_state = {
            "messages": messages,
            "user_insight": user_insight,
            "chat_history_summary": None,
            "lead_stage": None,
            "stage_injection": None,
            "diagnostics": []
        }

        result = await graph.ainvoke(initial_state, config=config)

        elapsed_ms = round((time.time() - start) * 1000, 2)
        all_messages = result["messages"]

        # Extract AI response (last AI message without tool_calls)
        ai_response = ""
        if all_messages:
            for msg in reversed(all_messages):
                if isinstance(msg, AIMessage) and not (hasattr(msg, "tool_calls") and msg.tool_calls):
                    ai_response = _extract_text(msg.content)
                    break

        # Build pipeline trace from diagnostics + messages
        pipeline = []
        pipeline.append({
            "step": "user",
            "label": "👤 User",
            "content": user_message,
        })
        
        # Append diagnostic tracking (classifier -> stylist -> tool -> stylist)
        if result.get("diagnostics"):
            for diag in result["diagnostics"]:
                pipeline.append(diag)

        # Parse tool results from messages for product cards + pipeline enrichment
        all_products = []
        for msg in all_messages:
            if isinstance(msg, ToolMessage):
                tool_content = _extract_text(msg.content)
                summary = tool_content[:600]
                raw_json = tool_content
                try:
                    parsed = json.loads(tool_content)
                    if parsed.get("status") == "success":
                        result_count = parsed.get("count", len(parsed.get("products", [])))
                        tier = parsed.get("tier", 0)
                        tier_info = f"[Tier {tier}]" if tier > 0 else "[SKU Direct]"
                        summary = f"✅ Tìm thấy {result_count} sản phẩm ({parsed.get('elapsed_ms', '?')}ms) {tier_info}"
                        products = parsed.get("products", [])
                        all_products.extend(products)
                        if products:
                            names = [f"• {p.get('name', '')}" for p in products[:5]]
                            summary += "\n" + "\n".join(names)
                            if len(products) > 5:
                                summary += f"\n... và {len(products) - 5} sản phẩm khác"
                    elif parsed.get("status") == "error":
                        summary = f"❌ {parsed.get('message', 'Error')}"
                    raw_json = json.dumps(parsed, ensure_ascii=False, indent=2, default=str)
                except Exception:
                    pass

                pipeline.append({
                    "step": "tool_result",
                    "label": f"📦 Tool Result ({getattr(msg, 'name', 'tool')})",
                    "content": summary,
                    "raw_json": raw_json,
                })

        # Add final responder step if there was a tool cycle
        # (diagnostics already captures the final stylist response)

        logger.info(
            "🏷️ [LEAD GRAPH] query='%s' | stage=%s | products=%d | time=%.2fms",
            user_message[:50], result.get("lead_stage", {}).get("stage_name", "UNKNOWN"),
            len(all_products), elapsed_ms
        )

        return {
            "response": ai_response,
            "elapsed_ms": elapsed_ms,
            "lead_stage": result.get("lead_stage"),
            "pipeline": pipeline,
            "products": all_products,
        }


# ═══════════════════════════════════════════════
# Singleton & Public API
# ═══════════════════════════════════════════════

_instance: LeadStageGraph | None = None

def get_lead_stage_agent(model_name: str | None = None) -> LeadStageGraph:
    """Get or create LeadStageGraph singleton."""
    global _instance
    if _instance is None:
        _instance = LeadStageGraph(model_name)
    return _instance


async def run_lead_stage_classifier(
    *,
    user_query: str,
    user_insight: str | None = None,
    chat_history_summary: str | None = None,
) -> dict:
    """
    (Backward Compatibility API)
    Chạy 1 bước classifier node để lấy stage injection, bỏ qua Stylist node.
    Được dùng bởi controller chính khi Controller đóng vai trò Stylist (chứa tools, logic cứng).
    """
    agent = get_lead_stage_agent()
    # Mock state để chạy method _classifier_node độc lập
    state = {
        "messages": [HumanMessage(content=user_query)],
        "user_insight": user_insight,
        "chat_history_summary": chat_history_summary,
        "lead_stage": None,
        "stage_injection": None,
        "diagnostics": []
    }
    res = await agent._classifier_node(state)
    return {
        "lead_stage": res.get("lead_stage"),
        "classifier_latency_ms": res.get("diagnostics", [{}])[0].get("elapsed_ms", 0),
        "stage_injection": res.get("stage_injection")
    }
