"""Service comprehension tools: workflow_graph and service_branches.

Export-based read-only tools that compute structural insights from a
single download_service API call each.

NOTE: service_map.py imports private functions from this module:
  _STATUS_TYPE_TO_EDGE, _build_branch_map, _build_form_lookup, _unwrap_export
"""

from __future__ import annotations

import json
import logging
from collections import defaultdict
from typing import Any

from fastmcp.exceptions import ToolError

from mcp_eregistrations_bpa.bpa_client import BPAClient
from mcp_eregistrations_bpa.bpa_client.errors import (
    BPAClientError,
    BPANotFoundError,
    translate_error,
)

logger = logging.getLogger(__name__)

__all__ = [
    "workflow_graph",
    "service_branches",
    "register_service_insights_tools",
]

# ---------------------------------------------------------------------------
# RoleStatusType enum (from RoleStatusType.java)
# ---------------------------------------------------------------------------
_PENDING = 0
_VALIDATED = 1
_DECLINE = 2
_REJECT = 3
_USER_DEFINED = 4

_STATUS_TYPE_TO_EDGE = {
    _VALIDATED: "forward",
    _DECLINE: "corrections",
    _REJECT: "reject",
    _USER_DEFINED: "user_defined",
}


def _extract_typed_edges(statuses: list[dict[str, Any]]) -> dict[str, list[str]]:
    """Group role status destinations by transition type.

    Args:
        statuses: List of status dicts from export (roleStatusType + destinations).

    Returns:
        Dict with forward, corrections, reject, user_defined destination name lists.
    """
    edges: dict[str, list[str]] = {
        "forward": [],
        "corrections": [],
        "reject": [],
        "user_defined": [],
    }
    for status in statuses:
        status_type: int | None = status.get("roleStatusType")
        edge_key = _STATUS_TYPE_TO_EDGE.get(status_type)  # type: ignore[arg-type]
        if not edge_key:
            continue
        destinations = status.get("destinations")
        if destinations:
            edges[edge_key].extend(destinations)
    return edges


def _bfs_forward_chain(
    graph: dict[str, dict[str, list[str]]],
    roles_meta: dict[str, dict[str, Any]],
    start_roles: list[str],
) -> tuple[list[dict[str, Any]], set[str], bool]:
    """BFS from start roles following forward (VALIDATED) edges.

    Args:
        graph: role_name -> typed edges dict (from _extract_typed_edges).
        roles_meta: role_name -> {type, startRole} metadata.
        start_roles: Names of roles with startRole=true.

    Returns:
        (chain, visited, has_cycles) where chain is list of role dicts
        in BFS order with depth, visited is set of visited names.
    """
    chain: list[dict[str, Any]] = []
    visited: set[str] = set()
    has_cycles = False
    # BFS queue: (role_name, depth)
    queue: list[tuple[str, int]] = [(name, 0) for name in start_roles]

    while queue:
        role_name, depth = queue.pop(0)
        if role_name in visited:
            has_cycles = True
            continue
        if role_name not in graph:
            continue
        visited.add(role_name)

        edges = graph[role_name]
        meta = roles_meta.get(role_name, {})
        chain.append(
            {
                "name": role_name,
                "role_type": meta.get("type", "UserRole"),
                "start": meta.get("startRole", False),
                "depth": depth,
                "forward": edges["forward"],
                "corrections": edges["corrections"],
                "reject": edges["reject"],
                "user_defined": edges["user_defined"],
            }
        )
        # Enqueue forward destinations
        for dest in edges["forward"]:
            if dest in visited:
                has_cycles = True
            else:
                queue.append((dest, depth + 1))

    return chain, visited, has_cycles


def _build_workflow_graph(roles: list[dict[str, Any]]) -> dict[str, Any]:
    """Build workflow graph from export roles data.

    Args:
        roles: Roles array from service export.

    Returns:
        Dict with chain, terminal_roles, orphaned_roles, stats.

    Raises:
        ToolError: If no roles or no start role found.
    """
    if not roles:
        raise ToolError(
            "No roles found in service export. Use 'role_list' to check if roles exist."
        )

    # Find start roles
    start_roles = [r["name"] for r in roles if r.get("startRole")]
    if not start_roles:
        raise ToolError(
            "No start role found. Check role configuration — "
            "exactly one role should have startRole=true."
        )

    # Build graph: name -> typed edges
    graph: dict[str, dict[str, list[str]]] = {}
    roles_meta: dict[str, dict[str, Any]] = {}
    for role in roles:
        name = role["name"]
        roles_meta[name] = {
            "type": role.get("type", "UserRole"),
            "startRole": role.get("startRole", False),
        }
        graph[name] = _extract_typed_edges(role.get("statuses") or [])

    # BFS forward chain
    chain, visited, has_cycles = _bfs_forward_chain(graph, roles_meta, start_roles)

    # Compute stats
    all_names = {r["name"] for r in roles}
    orphaned = sorted(all_names - visited)
    terminal = [r["name"] for r in chain if not graph.get(r["name"], {}).get("forward")]

    max_depth = max((r["depth"] for r in chain), default=0)
    max_fan_out = max(
        (len(graph.get(r["name"], {}).get("forward", [])) for r in chain),
        default=0,
    )

    return {
        "chain": chain,
        "terminal_roles": sorted(terminal),
        "orphaned_roles": orphaned,
        "stats": {
            "total_roles": len(roles),
            "chain_roles": len(chain),
            "orphaned_count": len(orphaned),
            "terminal_count": len(terminal),
            "max_depth": max_depth,
            "max_fan_out": max_fan_out,
            "has_cycles": has_cycles,
        },
    }


# ---------------------------------------------------------------------------
# Export option templates
# ---------------------------------------------------------------------------
_WORKFLOW_EXPORT_OPTIONS = {
    "serviceSelected": True,
    "rolesSelected": True,
    "determinantsSelected": False,
    "guideFormSelected": False,
    "applicantFormSelected": False,
    "sendFileFormSelected": False,
    "paymentFormSelected": False,
    "catalogsSelected": False,
    "botsSelected": False,
    "registrationsSelected": False,
    "printDocumentsSelected": False,
    "costsSelected": False,
    "requirementsSelected": False,
    "resultsSelected": False,
    "activityConditionsSelected": False,
    "registrationLawsSelected": False,
    "serviceLocationsSelected": False,
    "serviceTutorialsSelected": False,
    "serviceTranslationsSelected": False,
    "copyService": False,
}


def _unwrap_export(export_data: dict[str, Any]) -> dict[str, Any]:
    """Handle two-level nesting in export response."""
    if "service" in export_data and isinstance(export_data["service"], dict):
        return export_data["service"]
    return export_data


async def workflow_graph(
    service_id: str | int,
    instance: str | None = None,
) -> dict[str, Any]:
    """Return the approval workflow graph with typed transitions.

    Builds a directed graph from role status destinations in a single
    export call. Each role shows forward, corrections, reject, and
    user-defined transition targets.

    Args:
        service_id: BPA service UUID.
        instance: Instance profile name (from instance_list).

    Returns:
        dict with chain, terminal_roles, orphaned_roles, stats.
    """
    if not service_id:
        raise ToolError(
            "'service_id' is required. Use 'service_list' to find valid service IDs."
        )

    try:
        async with BPAClient.for_instance(instance) as client:
            # Verify service exists and get name
            try:
                service_data = await client.get(
                    "/service/{id}",
                    path_params={"id": service_id},
                    resource_type="service",
                    resource_id=str(service_id),
                )
            except BPANotFoundError:
                raise ToolError(
                    f"Service '{service_id}' not found. "
                    "Use 'service_list' to see available services."
                )

            # Fetch export with roles only
            export_data, _ = await client.download_service(
                str(service_id),
                options=_WORKFLOW_EXPORT_OPTIONS,
            )
    except ToolError:
        raise
    except BPAClientError as e:
        raise translate_error(e, resource_type="service")

    data = _unwrap_export(export_data)
    roles = data.get("roles", [])

    result = _build_workflow_graph(roles)
    result["service_id"] = str(service_id)
    result["service_name"] = service_data.get("name", data.get("name", ""))
    return result


# ---------------------------------------------------------------------------
# FormIO schema parsing
# ---------------------------------------------------------------------------


def _parse_form_components(
    schema: str | list[dict[str, Any]] | dict[str, Any] | None,
) -> dict[str, dict[str, Any]]:
    """Build field_key -> {label, type, values} lookup from FormIO schema.

    Walks the component tree recursively, handling panels, columns,
    and other nested layouts.

    Args:
        schema: FormIO schema — stringified JSON, components list, or full schema dict.

    Returns:
        Dict mapping component key to {label, type, values}.
    """
    if not schema:
        return {}

    # Parse stringified JSON
    if isinstance(schema, str):
        try:
            schema = json.loads(schema)
        except (json.JSONDecodeError, TypeError):
            return {}

    # Extract components list
    if isinstance(schema, dict):
        components = schema.get("components", [])
    elif isinstance(schema, list):
        components = schema
    else:
        return {}

    lookup: dict[str, dict[str, Any]] = {}
    _walk_components(components, lookup)
    return lookup


def _walk_components(
    components: list[dict[str, Any]],
    lookup: dict[str, dict[str, Any]],
) -> None:
    """Recursively walk FormIO component tree."""
    for comp in components:
        if not isinstance(comp, dict):
            continue
        key = comp.get("key")
        if key:
            # Extract values for select/radio/selectboxes
            values: list[dict[str, str]] = []
            if comp.get("data", {}).get("values"):
                values = comp["data"]["values"]
            elif comp.get("values"):
                values = comp["values"]

            lookup[key] = {
                "label": comp.get("label", ""),
                "type": comp.get("type", ""),
                "values": values,
            }

        # Recurse into nested components
        if comp.get("components"):
            _walk_components(comp["components"], lookup)
        # Handle columns layout
        if comp.get("columns"):
            for col in comp["columns"]:
                if isinstance(col, dict) and col.get("components"):
                    _walk_components(col["components"], lookup)
        # Handle table layout (rows is a list of lists; skip int from textarea)
        rows = comp.get("rows")
        if rows and isinstance(rows, list):
            for row in rows:
                if isinstance(row, list):
                    for cell in row:
                        if isinstance(cell, dict) and cell.get("components"):
                            _walk_components(cell["components"], lookup)


# ---------------------------------------------------------------------------
# Form lookup builder
# ---------------------------------------------------------------------------

_FORM_PAGES = [
    "guideFormPage",
    "applicantFormPage",
    "documentForm",
    "sendFileForm",
    "paymentForm",
]


def _build_form_lookup(
    export_data: dict[str, Any],
) -> dict[str, dict[str, Any]]:
    """Build field_key -> {label, type, values} lookup from all form pages.

    Parses formSchema from each of the 5 form pages in the export.

    Args:
        export_data: Unwrapped service export data.

    Returns:
        Merged lookup across all forms.
    """
    lookup: dict[str, dict[str, Any]] = {}
    for page_key in _FORM_PAGES:
        page = export_data.get(page_key)
        if not isinstance(page, dict):
            continue
        schema = page.get("formSchema")
        if schema:
            page_lookup = _parse_form_components(schema)
            lookup.update(page_lookup)
    return lookup


# ---------------------------------------------------------------------------
# Branch detection
# ---------------------------------------------------------------------------

_ESTIMATED_PATHS_CAP = 10_000


def _build_det_to_components(
    behaviours: list[dict[str, Any]],
) -> dict[str, list[str]]:
    """Build reverse index: determinant_id -> [componentKey].

    Parses jsonDeterminants from each effect to extract determinant IDs.
    Structure: [{"type": "AND|OR", "items": [{"determinantId": "<uuid>"}]}]
    """
    index: dict[str, list[str]] = defaultdict(list)
    for behaviour in behaviours:
        comp_key = behaviour.get("componentKey", "")
        for effect in behaviour.get("effects") or []:
            jd_str = effect.get("jsonDeterminants")
            if not jd_str or jd_str in ("null", "[]"):
                continue
            try:
                parsed = json.loads(jd_str) if isinstance(jd_str, str) else jd_str
            except (json.JSONDecodeError, TypeError):
                continue
            for group in parsed:
                for item in group.get("items") or []:
                    det_id = item.get("determinantId")
                    if det_id:
                        index[det_id].append(comp_key)
    return dict(index)


def _build_branch_map(
    determinants: list[dict[str, Any]],
    behaviours: list[dict[str, Any]],
    form_lookup: dict[str, dict[str, Any]],
) -> dict[str, Any]:
    """Detect form branching fields from determinants and behaviours.

    Args:
        determinants: serviceDeterminants from export.
        behaviours: componentBehaviours from export.
        form_lookup: field_key -> {label, type, values} from form schemas.

    Returns:
        Dict with branches, conditional_fields, stats.
    """
    # Filter out grid determinants and those without selectValue
    eligible = [
        d
        for d in determinants
        if not d.get("determinantInsideGrid")
        and d.get("selectValue")
        and d.get("operator") == "EQUAL"
    ]

    # Build reverse index
    det_to_components = _build_det_to_components(behaviours)

    # Group by targetFormFieldKey
    by_field: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for d in eligible:
        by_field[d["targetFormFieldKey"]].append(d)

    branches: list[dict[str, Any]] = []
    conditional_fields: list[dict[str, Any]] = []
    all_affected: set[str] = set()

    for field_key, field_dets in sorted(by_field.items()):
        # Get form info
        form_info = form_lookup.get(field_key, {})
        field_label = form_info.get("label") or field_key
        field_type = form_info.get("type", "")

        if len(field_dets) >= 2:
            # Potential branch: multiple values
            paths: list[dict[str, Any]] = []
            all_shows_sets: list[set[str]] = []
            for d in field_dets:
                shows = det_to_components.get(d["oldId"], [])
                shows_set = set(shows)
                if not shows_set:
                    continue
                all_affected.update(shows_set)
                all_shows_sets.append(shows_set)

                # Resolve value label from form schema
                value_label = d["selectValue"]
                for v in form_info.get("values") or []:
                    if v.get("value") == d["selectValue"]:
                        value_label = v.get("label", d["selectValue"])
                        break

                paths.append(
                    {
                        "value": d["selectValue"],
                        "label": value_label,
                        "determinant_name": d["name"],
                        "shows": sorted(shows),
                    }
                )

            # Only a branch if different values show different components
            if len(paths) >= 2 and len(set(frozenset(s) for s in all_shows_sets)) > 1:
                branches.append(
                    {
                        "field_key": field_key,
                        "field_label": field_label,
                        "field_type": field_type,
                        "paths": paths,
                    }
                )
        elif len(field_dets) == 1:
            # Single determinant = conditional field
            d = field_dets[0]
            shows = det_to_components.get(d["oldId"], [])
            if shows:
                all_affected.update(shows)
                conditional_fields.append(
                    {
                        "field_key": field_key,
                        "field_label": field_label,
                        "condition": d["selectValue"],
                        "determinant_name": d["name"],
                        "shows": sorted(shows),
                    }
                )

    # Compute estimated_paths
    estimated: int | str = 1
    for b in branches:
        if isinstance(estimated, int):
            estimated = estimated * len(b["paths"])
        if isinstance(estimated, int) and estimated > _ESTIMATED_PATHS_CAP:
            estimated = f"> {_ESTIMATED_PATHS_CAP}"
            break

    return {
        "branches": branches,
        "conditional_fields": conditional_fields,
        "stats": {
            "total_determinants": len(determinants),
            "branching_fields": len(branches),
            "conditional_fields": len(conditional_fields),
            "total_affected_components": len(all_affected),
            "estimated_paths": estimated,
        },
    }


# ---------------------------------------------------------------------------
# service_branches public tool
# ---------------------------------------------------------------------------

_BRANCHES_EXPORT_OPTIONS = {
    "serviceSelected": True,
    "rolesSelected": False,
    "determinantsSelected": True,
    "guideFormSelected": True,
    "applicantFormSelected": True,
    "sendFileFormSelected": True,
    "paymentFormSelected": True,
    "catalogsSelected": False,
    "botsSelected": False,
    "registrationsSelected": False,
    "printDocumentsSelected": False,
    "costsSelected": False,
    "requirementsSelected": False,
    "resultsSelected": False,
    "activityConditionsSelected": False,
    "registrationLawsSelected": False,
    "serviceLocationsSelected": False,
    "serviceTutorialsSelected": False,
    "serviceTranslationsSelected": False,
    "copyService": False,
}


async def service_branches(
    service_id: str | int,
    instance: str | None = None,
) -> dict[str, Any]:
    """Return form branching fields and the components each value shows.

    Identifies fields where different values create different form paths
    (show/hide different panels). Single export call.

    Args:
        service_id: BPA service UUID.
        instance: Instance profile name (from instance_list).

    Returns:
        dict with branches, conditional_fields, stats.
    """
    if not service_id:
        raise ToolError(
            "'service_id' is required. Use 'service_list' to find valid service IDs."
        )

    try:
        async with BPAClient.for_instance(instance) as client:
            try:
                service_data = await client.get(
                    "/service/{id}",
                    path_params={"id": service_id},
                    resource_type="service",
                    resource_id=str(service_id),
                )
            except BPANotFoundError:
                raise ToolError(
                    f"Service '{service_id}' not found. "
                    "Use 'service_list' to see available services."
                )

            export_data, _ = await client.download_service(
                str(service_id),
                options=_BRANCHES_EXPORT_OPTIONS,
            )
    except ToolError:
        raise
    except BPAClientError as e:
        raise translate_error(e, resource_type="service")

    data = _unwrap_export(export_data)

    # Parse form schemas from all 5 pages
    form_lookup = _build_form_lookup(data)

    # Extract determinants and behaviours
    determinants = data.get("serviceDeterminants", [])
    behaviours = data.get("componentBehaviours", [])

    result = _build_branch_map(determinants, behaviours, form_lookup)
    result["service_id"] = str(service_id)
    result["service_name"] = service_data.get("name", data.get("name", ""))
    return result


# ---------------------------------------------------------------------------
# Registration
# ---------------------------------------------------------------------------


def register_service_insights_tools(mcp: Any) -> None:
    """Register service insight tools with the MCP server."""
    from mcp_eregistrations_bpa.tools.annotations import READ

    mcp.tool(annotations=READ)(workflow_graph)
    mcp.tool(annotations=READ)(service_branches)
