"""گراف گفتگو (LangGraph) — پورت خط لوله v13 در سه گام:

۱) بازنویسی پرسش با تاریخچه (فقط وقتی تاریخچه هست — صرفه‌جویی در تأخیر)
۲) بازیابی ترکیبی با چند نوع‌واژه (اختیاری)
۳) تولید پاسخ فارسی با قواعد تکیه بر داده
"""

import logging
import re
import json as _json
from typing import Annotated, List, Optional

from typing_extensions import TypedDict

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages

from . import content as content_mod
from .config import settings
from .retrieval import retrieve_final

log = logging.getLogger(__name__)

FALLBACK_MESSAGE = (
    "متأسفانه اطلاعات دقیقی در این باره در سایت ثبت نشده است. "
    "برای راهنمایی بیشتر می‌توانید از صفحه «تماس با ما» با ما در ارتباط باشید."
)
LLM_UNAVAILABLE = (
    "سرویس هوش مصنوعی هنوز پیکربندی نشده است. "
    "متغیرهای LLM_BASE_URL و LLM_MODEL را در تنظیمات سرویس چت‌بات تنظیم کنید."
)

_llm = None


def get_llm() -> ChatOpenAI | None:
    """LLM تولید — اگر پیکربندی نشده باشد None برمی‌گرداند."""
    global _llm
    if _llm is None and settings.llm_base_url:
        _llm = ChatOpenAI(
            model=settings.llm_model,
            base_url=settings.llm_base_url,
            api_key=settings.llm_api_key or "not-needed",
            temperature=0,
            streaming=True,
        )
    return _llm


class GraphState(TypedDict):
    messages: Annotated[List, add_messages]
    documents: List[str]
    contextualized_query: Optional[str]


# --- گام ۱: بازنویسی پرسش با تاریخچه ---

def contextualize_query(question: str, history: list) -> str:
    """ضمیرها و اشاره‌های پرسش را با تاریخچه گفتگو مشخص می‌کند («اون دومی» → موضوع واقعی)."""
    model = get_llm()
    if model is None or not history:
        return question
    try:
        history_text = "\n".join(
            f"{'کاربر' if isinstance(m, HumanMessage) else 'دستیار'}: {str(m.content)[:300]}"
            for m in history[-6:]
        )
        prompt = (
            "با توجه به تاریخچه گفتگو، پرسش کاربر را به یک پرسش مستقل و کامل فارسی بازنویسی کن.\n"
            "- ضمیرها و اشاره‌ها (او، آن، همین، دومی، این یکی) را به موضوع واقعیِ ذکرشده در تاریخچه تبدیل کن.\n"
            "- اگر پرسش از قبل مستقل است، عیناً همان را برگردان.\n"
            "- فقط متن پرسش بازنویسی‌شده را برگردان؛ هیچ توضیحی ننویس.\n\n"
            f"تاریخچه گفتگو:\n{history_text}\n\n"
            f"پرسش کاربر: {question}\n\n"
            "پرسش مستقل:"
        )
        resp = model.invoke([HumanMessage(content=prompt)], config={"tags": ["utility"]})
        rewritten = (resp.content or "").strip() if isinstance(resp.content, str) else ""
        if not rewritten or len(rewritten) > len(question) * 4:
            return question
        if rewritten != question:
            log.info("Contextualized: '%s' → '%s'", question[:60], rewritten[:60])
        return rewritten
    except Exception as e:
        log.warning("Contextualization failed: %s", e)
        return question


def expand_query(query: str) -> list[str]:
    """گسترش پرسش به چند نوع‌واژه (اختیاری — پیش‌فرض خاموش برای پیکره کوچک)."""
    model = get_llm()
    if model is None:
        return [query]
    try:
        prompt = (
            "برای جست‌وجو در محتوای سایت یک گروه ساختمانی، ۳ نوع‌واژه جایگزین برای این پرسش بساز: "
            "بازنویسی معنایی، استخراج کلیدواژه، گسترش با مفاهیم مرتبط.\n"
            'فقط JSON برگردان: {"variants": ["۱", "۲", "۳"]}\n\n'
            f"پرسش: {query}"
        )
        resp = model.invoke([HumanMessage(content=prompt)], config={"tags": ["utility"]})
        text = resp.content if isinstance(resp.content, str) else ""
        match = re.search(r'\{.*"variants".*\}', text, re.DOTALL)
        if match:
            variants = [v for v in _json.loads(match.group()).get("variants", []) if v and v != query][:3]
            if variants:
                log.info("Query variants: %s", [query] + variants)
                return [query] + variants
    except Exception as e:
        log.warning("Query expansion failed: %s", e)
    return [query]


# --- گام ۲: بازیابی ---

def retrieve_node(state: GraphState) -> dict:
    messages = state.get("messages", [])
    if not messages:
        return {"documents": []}
    question = messages[-1].content
    history = messages[:-1]

    query = contextualize_query(question, history) if settings.contextualize_enabled else question
    variants = expand_query(query) if settings.query_expansion else [query]

    docs = retrieve_final(query, variants)
    doc_texts = [f"{d['text']}\n(منبع: {d['title']} — {d['url']})" for d in docs]
    log.info("Retrieved %d documents for '%s'", len(doc_texts), question[:60])
    return {
        "documents": doc_texts,
        "contextualized_query": query if query != question else None,
    }


# --- گام ۳: تولید پاسخ ---

def build_system_prompt(documents: list[str], contextualized_query: str | None) -> str:
    site_ctx = content_mod.fetch_site_context(settings.database_url, settings.site_context_ttl)
    docs_block = "\n\n———\n\n".join(documents) if documents else "(سند مرتبطی یافت نشد)"
    reminder = (
        f"\n--- یادآوری ---\nپرسش کاربر پس از بازنویسی این است: «{contextualized_query}»\n\n"
        if contextualized_query else ""
    )
    return (
        "تو دستیار گفتگوی هوشمند سایت «گروه ساختمانی آرامش» هستی و فقط به فارسی پاسخ می‌دهی.\n"
        "دانش تو اکیداً محدود به اطلاعات زیر است.\n\n"
        f"--- اطلاعات پایه سایت (همیشه معتبر) ---\n{site_ctx}\n\n"
        f"--- اسناد بازیابی‌شده مرتبط با پرسش ---\n{docs_block}\n"
        f"{reminder}"
        "--- قوانین پاسخ ---\n"
        "۱. فقط بر اساس اطلاعات بالا پاسخ بده؛ به‌هیچ‌وجه از دانش بیرونی استفاده نکن و چیزی نساز.\n"
        "۲. پاسخ کوتاه، دقیق و مفید باشد (معمولاً ۲ تا ۵ جمله) و لحن حرفه‌ای و محترمانه داشته باشد.\n"
        f"۳. اگر اطلاعات کافی در اسناد نیست، فقط همین جمله را بنویس: «{FALLBACK_MESSAGE}»\n"
        "۴. اگر سند مسیر صفحه مرتبط دارد، در پایان پاسخ آن را معرفی کن (مثلاً: جزئیات کامل در صفحه /blog/... ).\n"
        "۵. درباره قیمت، زمان دقیق یا هر تعهدی فقط چیزی را بگو که در داده‌ها آمده است."
    )


def generate_node(state: GraphState) -> dict:
    messages = state.get("messages", [])
    question = messages[-1].content if messages else ""

    model = get_llm()
    if model is None:
        return {"messages": [AIMessage(content=LLM_UNAVAILABLE)]}

    documents = state.get("documents") or []
    system = build_system_prompt(documents, state.get("contextualized_query"))
    log.info("Generating answer with %d docs (system prompt %d chars)", len(documents), len(system))

    history = messages[:-1]
    reply = model.invoke(
        [SystemMessage(content=system), *history, HumanMessage(content=question)],
        config={"tags": ["generate"]},
    )
    return {"messages": [reply]}


# --- گراف ---

memory = MemorySaver()
workflow = StateGraph(GraphState)
workflow.add_node("retrieve", retrieve_node)
workflow.add_node("generate", generate_node)
workflow.set_entry_point("retrieve")
workflow.add_edge("retrieve", "generate")
workflow.add_edge("generate", END)
graph = workflow.compile(checkpointer=memory)