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

Flow: User → Classifier ⇄ Tools → Stylist → END

2 con AI:
  - Classifier (Planner + Tool Caller): Xác định lead stage, gọi tools khi cần,
    có thể trả lời trực tiếp nếu câu hỏi đơn giản.
  - Stylist (Responder): KHÔNG gọi tools. Chỉ nhận kết quả + format câu trả lời đẹp.

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

import json
import logging
import re
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
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, STYLIST_SYSTEM_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__)

# Regex to extract STAGE_JSON: {...} from classifier response
_STAGE_JSON_RE = re.compile(r'STAGE_JSON:\s*(\{.*\})', re.DOTALL)


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 "")


def _parse_stage_from_text(text: str) -> dict | None:
    """Try to parse lead stage JSON from classifier response text.
    
    Supports 2 formats:
    1. STAGE_JSON: {...}  (inline marker)
    2. Pure JSON response  (entire response is JSON)
    """
    if not text:
        return None
    
    # Format 1: STAGE_JSON: {...}
    m = _STAGE_JSON_RE.search(text)
    if m:
        try:
            result = json.loads(m.group(1))
            if "stage" in result and "stage_name" in result:
                return result
        except (json.JSONDecodeError, TypeError):
            pass
    
    # Format 2: Pure JSON
    stripped = text.strip()
    if stripped.startswith("{"):
        # Strip markdown code fences if present
        if stripped.startswith("```"):
            stripped = re.sub(r'^```\w*\n?', '', stripped)
            stripped = re.sub(r'\n?```$', '', stripped).strip()
        try:
            result = json.loads(stripped)
            if "stage" in result and "stage_name" in result:
                return result
        except (json.JSONDecodeError, TypeError):
            pass
    
    return None



# ═══════════════════════════════════════════════
# 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
    classifier_responded: bool  # True nếu Classifier trả lời trực tiếp (skip Stylist)
    diagnostics: Annotated[list[dict[str, Any]], lambda x, y: x + y]


# ═══════════════════════════════════════════════
# Graph Builder — 2-Agent Architecture
# Classifier (1 call: classify + tools) → Stylist (1 call: format)
# ═══════════════════════════════════════════════
class LeadStageGraph:
    """2-Agent LangGraph: Classifier ⇄ Tools → Stylist → END."""

    def __init__(self, model_name: str | None = None):
        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 + Tool Caller (streaming off, has tools)
        self.classifier_llm = create_llm(
            model_name=self.model_name,
            streaming=False,
        )
        self.classifier_with_tools = self.classifier_llm.bind_tools(self.tools)

        # AI 2: Stylist — response only (streaming, NO tools)
        self.stylist_llm = create_llm(
            model_name=self.model_name,
            streaming=True,
        )

        self._compiled = None
        logger.info(f"✅ LeadStageGraph initialized | model: {self.model_name}")

    # ─────────────────────────────────────
    # Node 1: CLASSIFIER — 1 LLM call: classify stage + call tools
    # ─────────────────────────────────────
    async def _classifier_node(self, state: LeadStageState, config: RunnableConfig | None = None) -> dict:
        """
        Classifier AI — 1 LLM call duy nhất:
        - Xác định lead stage (qua STAGE_JSON marker)
        - Gọi tools nếu cần (search SP, check stock...)
        - Hoặc trả lời trực tiếp nếu câu hỏi đơn giản

        Lần 2+ (sau tools return): pass-through → route thẳng tới Stylist
        """
        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ử."

        has_tool_results = any(isinstance(m, ToolMessage) for m in messages)

        # ── PASS 2+: Tool results back → skip LLM, go to Stylist ──
        if has_tool_results:
            logger.info("🔄 Classifier pass-through → Stylist")
            return {
                "lead_stage": state.get("lead_stage"),
                "stage_injection": state.get("stage_injection"),
                "classifier_responded": False,
                "diagnostics": [],
            }

        # ── PASS 1: Single LLM call — classify + tools ──
        context_block = (
            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}"
        )

        classifier_messages = [
            SystemMessage(content=LEAD_STAGE_CLASSIFIER_PROMPT),
        ]
        for msg in messages:
            if not isinstance(msg, SystemMessage):
                classifier_messages.append(msg)
        classifier_messages.append(HumanMessage(content=context_block))

        start = time.time()
        response = await self.classifier_with_tools.ainvoke(classifier_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

        # ── Parse lead stage from response text (STAGE_JSON: {...}) ──
        lead_stage = None
        stage_injection = None
        stage_result = _parse_stage_from_text(response_text)
        if stage_result:
            lead_stage = stage_result
            stage_injection = format_stage_injection(stage_result)
            logger.info(
                f"🎯 Stage: {stage_result.get('stage_name')} "
                f"(conf={stage_result.get('confidence')}) | {elapsed_ms:.0f}ms"
            )


        # ── Build diagnostic ──
        diag = {"elapsed_ms": round(elapsed_ms)}

        if has_tool_calls:
            diag["step"] = "classifier_tool_call"
            diag["label"] = "🔧 Classifier (Stage + Tool Call)"
            tool_names = ', '.join(tc['name'] for tc in response.tool_calls)
            stage_info = f"Stage {lead_stage['stage']}: {lead_stage['stage_name']}" if lead_stage else "Stage: unknown"
            diag["content"] = f"{stage_info} | Tools: {tool_names}"
            diag["tool_calls"] = [{"name": tc["name"], "args": tc["args"]} for tc in response.tool_calls]
            if lead_stage:
                diag["raw_json"] = json.dumps(lead_stage, ensure_ascii=False, indent=2)
        elif lead_stage:
            diag["step"] = "classifier_stage"
            diag["label"] = "🧠 Classifier (Stage Only)"
            diag["content"] = (
                f"Stage {lead_stage.get('stage')}: {lead_stage.get('stage_name')} "
                f"(Conf: {lead_stage.get('confidence')})\n"
                f"Lý do: {lead_stage.get('reasoning', '')}"
            )
            diag["raw_json"] = json.dumps(lead_stage, ensure_ascii=False, indent=2)
        else:
            diag["step"] = "classifier_direct"
            diag["label"] = "💬 Classifier (Direct Response)"
            diag["content"] = response_text[:500]

        # classifier_responded = True → skip Stylist (only for simple greetings)
        is_direct = not has_tool_calls and not lead_stage

        return {
            "messages": [response],
            "lead_stage": lead_stage,
            "stage_injection": stage_injection,
            "classifier_responded": is_direct,
            "diagnostics": [diag],
        }

    # ─────────────────────────────────────
    # Router: Classifier → tools | stylist | end
    # ─────────────────────────────────────
    def _after_classifier(self, state: LeadStageState) -> str:
        """
        Sau Classifier:
        - Có tool_calls → đi tools
        - Classifier trả lời trực tiếp (greeting, no stage) → end
        - Mọi trường hợp khác (có lead_stage, hoặc sau tools) → stylist
        """
        last = state["messages"][-1]
        if hasattr(last, "tool_calls") and last.tool_calls:
            return "tools"
        if state.get("classifier_responded"):
            return "end"
        return "stylist"

    # ─────────────────────────────────────
    # Node 2: STYLIST — Format response (NO tools)
    # ─────────────────────────────────────
    async def _stylist_node(self, state: LeadStageState, config: RunnableConfig | None = None) -> dict:
        """
        Stylist AI: Nhận lead stage + tool results, format câu trả lời tư vấn đẹp.
        KHÔNG gọi tools — chỉ viết response.
        """
        messages = state["messages"]
        injection = state.get("stage_injection") or ""

        # Build system prompt with stage context
        stylist_sys = STYLIST_SYSTEM_PROMPT.format(stage_injection=injection)

        # Collect all tool results to give stylist context
        tool_context_parts = []
        for msg in messages:
            if isinstance(msg, ToolMessage):
                tool_name = getattr(msg, "name", "tool")
                tool_content = _extract_text(msg.content)
                # Parse for cleaner summary
                try:
                    parsed = json.loads(tool_content)
                    if parsed.get("status") == "success":
                        products = parsed.get("products", [])
                        if products:
                            product_lines = []
                            for p in products[:8]:
                                line = f"- {p.get('name', '')} | {p.get('price', '')} | SKU: {p.get('sku_code', '')}"
                                if p.get("url"):
                                    line += f" | Link: {p['url']}"
                                product_lines.append(line)
                            tool_context_parts.append(
                                f"[{tool_name}] Tìm thấy {len(products)} sản phẩm:\n" + "\n".join(product_lines)
                            )
                        else:
                            tool_context_parts.append(f"[{tool_name}] {json.dumps(parsed, ensure_ascii=False)[:500]}")
                    else:
                        tool_context_parts.append(f"[{tool_name}] {parsed.get('message', tool_content[:300])}")
                except Exception:
                    tool_context_parts.append(f"[{tool_name}] {tool_content[:500]}")

        # Build stylist messages
        stylist_messages = [SystemMessage(content=stylist_sys)]

        # Add user's original question
        for msg in messages:
            if isinstance(msg, HumanMessage):
                stylist_messages.append(msg)
                break  # Only first human message (the actual query)

        # Add tool results as context
        if tool_context_parts:
            tool_summary = "=== KẾT QUẢ TRA CỨU TỪ TOOLS ===\n" + "\n\n".join(tool_context_parts)
            stylist_messages.append(HumanMessage(content=tool_summary))

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

        response_text = _extract_text(response.content)

        diag = {
            "step": "stylist",
            "label": "💬 Stylist AI (Response)",
            "content": response_text[:500] + ("..." if len(response_text) > 500 else ""),
            "elapsed_ms": round(elapsed_ms),
        }

        return {
            "messages": [response],
            "diagnostics": [diag],
        }

    # ─────────────────────────────────────
    # Build Graph
    # ─────────────────────────────────────
    def build(self):
        """Build 2-agent graph: Classifier ⇄ Tools → Stylist → END."""
        if self._compiled:
            return self._compiled

        workflow = StateGraph(LeadStageState)

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

        # Flow: Entry → Classifier
        workflow.set_entry_point("classifier")

        # Classifier → (tools | stylist | end)
        workflow.add_conditional_edges(
            "classifier",
            self._after_classifier,
            {"tools": "tools", "stylist": "stylist", "end": END},
        )

        # Tools → Classifier (loop back so classifier sees tool results)
        workflow.add_edge("tools", "classifier")

        # Stylist → END
        workflow.add_edge("stylist", END)

        self._compiled = workflow.compile()
        logger.info("✅ LeadStageGraph compiled: Classifier ⇄ Tools → Stylist → END")
        return self._compiled

    # ─────────────────────────────────────
    # Main Entry Point
    # ─────────────────────────────────────
    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,
            "classifier_responded": False,
            "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 = ""
        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
        pipeline = [{"step": "user", "label": "👤 User", "content": user_message}]

        if result.get("diagnostics"):
            for diag in result["diagnostics"]:
                pipeline.append(diag)

        # Parse tool results for product cards
        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,
                })

        logger.info(
            "🏷️ [LEAD GRAPH] query='%s' | stage=%s | products=%d | time=%.2fms",
            user_message[:50],
            result.get("lead_stage", {}).get("stage_name", "UNKNOWN") if result.get("lead_stage") else "DIRECT",
            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()
    state = {
        "messages": [HumanMessage(content=user_query)],
        "user_insight": user_insight,
        "chat_history_summary": chat_history_summary,
        "lead_stage": None,
        "stage_injection": None,
        "classifier_responded": False,
        "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"),
    }
