"""Role suggestion API routes for retrieval planning."""
from __future__ import annotations

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

from fastapi import APIRouter, Depends
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from normbase.db.connection import get_db
from normbase.models.norm import Norm
from normbase.models.schemas import RoleCandidate, RoleSuggestRequest, RoleSuggestResponse

router = APIRouter(prefix="/roles", tags=["roles"])


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}


@router.post("/suggest", response_model=RoleSuggestResponse)
async def suggest_roles(
    request: RoleSuggestRequest,
    db: AsyncSession = Depends(get_db),
):
    """
    Suggest candidate agent roles based on active norms and context.

    Notes:
    - Uses active norms only.
    - Optionally narrows by field_ids.
    - Treats '*' as fallback role with lower priority.
    """
    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(
        lambda: {"norm_count": 0, "matched_contexts": set(), "score": 0.0}
    )

    for norm in norms:
        roles = norm.scope_roles or []
        if not roles:
            continue
        norm_contexts = {str(x).strip().lower() for x in (norm.scope_contexts or []) if str(x).strip()}
        matched = sorted(norm_contexts.intersection(context_terms))
        for role in roles:
            role_name = str(role).strip()
            if not role_name:
                continue
            item = role_stats[role_name]
            item["norm_count"] += 1
            item["matched_contexts"].update(matched)
            base = 0.3 if role_name == "*" else 1.0
            item["score"] += base + (0.2 * len(matched))

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

    return RoleSuggestResponse(
        candidates=candidates,
        recommended_role=candidates[0].role if candidates else None,
    )

