"""Business logic for role suggestion and retrieval planning."""
from __future__ import annotations

from collections import defaultdict
from typing import Any, Dict, List, Set

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from normbase.models.norm import Norm, NormTag
from normbase.schemas import RetrievalPlanResponse, RoleCandidate, RoleSuggestRequest

ROLE_CATEGORY = "role"
CONTEXT_TAG_CATEGORIES = ("scenario", "domain", "tool", "skill", "capability")
PLAN_TAG_CATEGORIES = (ROLE_CATEGORY, *CONTEXT_TAG_CATEGORIES)


def _flatten_context_terms(context: Dict[str, Any]) -> Set[str]:
    terms: Set[str] = set()

    def _walk(value: Any) -> None:
        if isinstance(value, dict):
            for k, v in value.items():
                terms.add(str(k).strip().lower())
                _walk(v)
            return
        if isinstance(value, list):
            for item in value:
                _walk(item)
            return
        if value is None:
            return
        terms.add(str(value).strip().lower())

    _walk(context or {})
    return {t for t in terms if t}


def _extract_context_tag_filters(context: Dict[str, Any]) -> Dict[str, List[str]]:
    filters: Dict[str, List[str]] = {}
    for category in CONTEXT_TAG_CATEGORIES:
        raw = (context or {}).get(category)
        values: List[str] = []
        if isinstance(raw, str):
            text = raw.strip()
            if text:
                values.append(text)
        elif isinstance(raw, list):
            for item in raw:
                text = str(item).strip()
                if text and text not in values:
                    values.append(text)
        if values:
            filters[category] = values
    return filters


def _empty_role_stats() -> Dict[str, Any]:
    return {
        "norm_count": 0,
        "matched_contexts": set(),
        "score": 0.0,
        "matched_tags": {category: set() for category in CONTEXT_TAG_CATEGORIES},
    }


class RoleService:
    """Service layer for role plan calculation."""

    async def build_role_plan(
        self,
        db: AsyncSession,
        request: RoleSuggestRequest,
    ) -> RetrievalPlanResponse:
        query = select(Norm).where(Norm.status == "active")
        if request.field_ids:
            query = query.where(Norm.field_id.in_(request.field_ids))

        norms = (await db.execute(query)).scalars().all()
        context_terms = _flatten_context_terms(request.context)
        role_stats: Dict[str, Dict[str, Any]] = defaultdict(_empty_role_stats)

        for norm in norms:
            tag_rows = (
                await db.execute(
                    select(NormTag).where(
                        NormTag.norm_id == norm.id,
                        NormTag.tag_category.in_(PLAN_TAG_CATEGORIES),
                    )
                )
            ).scalars().all()
            role_tags = [
                str(tag.tag).strip()
                for tag in tag_rows
                if tag.tag_category == ROLE_CATEGORY and str(tag.tag).strip()
            ]
            if not role_tags:
                continue

            matched_contexts = sorted(
                {
                    str(tag.tag).strip().lower()
                    for tag in tag_rows
                    if tag.tag_category in CONTEXT_TAG_CATEGORIES
                    and str(tag.tag).strip().lower() in context_terms
                }
            )

            matched_tags_by_category = {
                category: set() for category in CONTEXT_TAG_CATEGORIES
            }
            for tag in tag_rows:
                category = str(tag.tag_category).strip()
                if category not in matched_tags_by_category:
                    continue
                value = str(tag.tag).strip()
                if value and value.lower() in context_terms:
                    matched_tags_by_category[category].add(value)

            for role in role_tags:
                role_name = str(role).strip()
                if not role_name:
                    continue
                item = role_stats[role_name]
                item["norm_count"] += 1
                item["matched_contexts"].update(matched_contexts)
                for category in CONTEXT_TAG_CATEGORIES:
                    item["matched_tags"][category].update(matched_tags_by_category[category])
                base = 0.3 if role_name == "*" else 1.0
                item["score"] += base + (0.2 * len(matched_contexts))

        ranked = sorted(
            role_stats.items(),
            key=lambda pair: (pair[1]["score"], pair[1]["norm_count"]),
            reverse=True,
        )

        candidates: List[RoleCandidate] = []
        for role_name, item in ranked[: request.top_k]:
            candidates.append(
                RoleCandidate(
                    role=role_name,
                    score=round(float(item["score"]), 4),
                    norm_count=int(item["norm_count"]),
                    matched_contexts=sorted(item["matched_contexts"]),
                )
            )

        suggested_tag_filters = _extract_context_tag_filters(request.context)
        recommended_role = candidates[0].role if candidates else None
        role_is_global = recommended_role == "*"
        if recommended_role and recommended_role != "*":
            suggested_tag_filters[ROLE_CATEGORY] = [recommended_role]
        if role_is_global:
            suggested_tag_filters[ROLE_CATEGORY] = ["*"]
        if recommended_role and recommended_role in role_stats:
            matched_tags = role_stats[recommended_role]["matched_tags"]
            for category in CONTEXT_TAG_CATEGORIES:
                if category in suggested_tag_filters:
                    continue
                values = sorted(matched_tags.get(category, set()))
                if values:
                    suggested_tag_filters[category] = values

        return RetrievalPlanResponse(
            candidates=candidates,
            recommended_role=recommended_role,
            role_is_global=role_is_global,
            suggested_tag_filters=suggested_tag_filters,
        )
