"""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, NormTag
from normbase.models.schemas import (
    RetrievalPlanResponse,
    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}


def _extract_context_tag_filters(context: Dict[str, Any]) -> Dict[str, List[str]]:
    """Extract retrieval tag filters directly from context payload."""
    filters: Dict[str, List[str]] = {}
    for category in ("scenario", "domain", "tool"):
        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": {
            "scenario": set(),
            "domain": set(),
            "tool": set(),
        },
    }


async def _build_role_plan(request: RoleSuggestRequest, db: AsyncSession) -> 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_(["role", "scenario", "domain", "tool"]),
                )
            )
        ).scalars().all()
        role_tags = [
            str(tag.tag).strip()
            for tag in tag_rows
            if tag.tag_category == "role" 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 {"scenario", "domain", "tool"}
                and str(tag.tag).strip().lower() in context_terms
            }
        )

        matched_tags_by_category = {"scenario": set(), "domain": set(), "tool": set()}
        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 ("scenario", "domain", "tool"):
                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"] = [recommended_role]
    if role_is_global:
        suggested_tag_filters["role"] = ["*"]
    if recommended_role and recommended_role in role_stats:
        matched_tags = role_stats[recommended_role]["matched_tags"]
        for category in ("scenario", "domain", "tool"):
            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,
    )


@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.
    """
    plan = await _build_role_plan(request, db)
    return RoleSuggestResponse(
        candidates=plan.candidates,
        recommended_role=plan.recommended_role,
    )


@router.post("/plan", response_model=RetrievalPlanResponse)
async def plan_retrieval(
    request: RoleSuggestRequest,
    db: AsyncSession = Depends(get_db),
):
    """
    Suggest retrieval strategy (role + tag_filters) from active norms and context.
    """
    return await _build_role_plan(request, db)

