"""Query classification and category gating router for Joule Dynamics Real Estate Intelligence."""

from __future__ import annotations

import json
import re

from services.observability import setup_logger
from services.groq_client import (
    groq_call,
    is_structural_error,
    GROQ_ROUTE_MODEL,
    GROQ_FALLBACK_ROUTE_MODEL,
)
from services.prompts import ROUTER_PROMPT

logger = setup_logger(__name__)

_VALID_CLASSIFICATIONS = frozenset({
    "OUT_OF_SCOPE", "PATH_A", "PATH_B", "BOTH", "GREETING", "COMMERCIAL_HANDOFF",
    "ALERT_SUBSCRIPTION",
})


def extract_json_from_text(text: str) -> dict:
    """Extract a JSON object from raw model output, handling think tags and markdown."""
    if not text:
        return {}

    clean = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()

    try:
        return json.loads(clean)
    except json.JSONDecodeError:
        pass

    match = re.search(r"\{[^{}]*\}", clean, re.DOTALL)
    if match:
        try:
            return json.loads(match.group(0))
        except json.JSONDecodeError:
            pass

    return {}


def _normalize_tool_categories(categories: object) -> list[str]:
    """Normalize raw parsed tool categories into a clean list of strings."""
    if isinstance(categories, str):
        return [categories]
    if isinstance(categories, list):
        return [str(c) for c in categories]
    return []


def _build_router_messages(user_query: str, recent_history: list[dict] | None = None) -> list[dict]:
    """Constructs router message list with recent conversation history for multi-turn awareness."""
    messages = [{"role": "system", "content": ROUTER_PROMPT}]
    if recent_history:
        for msg in recent_history[-2:]:
            content = msg.get("content", "")
            if "User Query:" in content:
                content = content.split("User Query:", 1)[-1].strip()
            messages.append({
                "role": msg.get("role", "user"),
                "content": content[:400],
            })
    messages.append({"role": "user", "content": user_query})
    return messages


def _find_last_assistant_message(recent_history: list[dict]) -> str:
    """Extracts the text of the most recent assistant turn in the history."""
    for msg in reversed(recent_history):
        if msg.get("role") == "assistant":
            return msg.get("content", "").lower()
    return ""


def _rescue_classification(
    classification: str,
    user_query: str,
    recent_history: list[dict] | None,
) -> str:
    """Architectural guardrail: rescues false OUT_OF_SCOPE classifications during active conversational flows."""
    if classification != "OUT_OF_SCOPE" or not recent_history:
        return classification

    query_lower = user_query.strip().lower()

    if re.search(r"[\w\.-]+@[\w\.-]+\.\w+", query_lower):
        logger.info("[router] Rescuing OUT_OF_SCOPE -> ALERT_SUBSCRIPTION (email address in active session)")
        return "ALERT_SUBSCRIPTION"

    last_assistant_msg = _find_last_assistant_message(recent_history)
    alert_keywords = ("alert", "notification", "email address", "subscribe", "digest")
    if any(kw in last_assistant_msg for kw in alert_keywords) and len(query_lower.split()) <= 8:
        logger.info("[router] Rescuing OUT_OF_SCOPE -> ALERT_SUBSCRIPTION (continuation of alert setup)")
        return "ALERT_SUBSCRIPTION"

    if len(query_lower.split()) <= 5 and any(c in last_assistant_msg for c in ("?", "which", "would you", "what")):
        logger.info("[router] Rescuing OUT_OF_SCOPE -> PATH_A (short answer to assistant question)")
        return "PATH_A"

    return classification


async def classify_query(
    user_query: str,
    recent_history: list[dict] | None = None,
) -> tuple[str, list[str]]:
    """Classify user query and extract tool categories for category gating with multi-turn awareness."""
    messages = _build_router_messages(user_query, recent_history)

    for model in [GROQ_ROUTE_MODEL, GROQ_FALLBACK_ROUTE_MODEL]:
        try:
            res = await groq_call(
                model=model,
                messages=messages,
                max_tokens=1500,
                temperature=0.0,
            )
            raw = res.choices[0].message.content or "{}"
            data = extract_json_from_text(raw)
            classification = data.get("classification", "")
            tool_categories = _normalize_tool_categories(data.get("tool_categories", []))

            if classification not in _VALID_CLASSIFICATIONS:
                logger.warning(
                    f"[router] {model} returned invalid classification '{classification}', defaulting PATH_A"
                )
                classification = "PATH_A"

            classification = _rescue_classification(classification, user_query, recent_history)

            logger.info(
                f"[router] {model} -> {classification} (categories={tool_categories}): {data.get('reason', '')}"
            )
            return classification, tool_categories

        except Exception as exc:
            if is_structural_error(exc):
                logger.warning(f"[router] structural error on {model}: {exc}")
            else:
                logger.warning(f"[router] {model} failed, trying fallback: {exc}")

    fallback_classification = _rescue_classification("PATH_A", user_query, recent_history)
    return fallback_classification, []