"""Minimum-sufficient concrete-role planner.

Families are candidate pools, never actors.  The planner uses a deterministic greedy
set-cover with token and independence constraints and records both selected and skipped roles.
"""
from __future__ import annotations

import os
from typing import Any

import yaml

from .budget_planner import estimate_plan, estimate_role, max_selected_roles, workflow_budget
from .coverage_model import (
    CAPABILITY_ROLE_HINTS,
    RISK_ROLE_HINTS,
    required_coverage,
    role_coverage,
    tokens,
)
from .lens_policy import candidate_family_errors, normalize_family_ids, required_lenses
from .task_graph import build_task_graph

ROOT = os.environ.get("CLAUDE_PROJECT_DIR") or os.path.dirname(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
)
REG = os.path.join(ROOT, "org-os", "00-role-registry")
ARTIFACT_REGISTRY = os.path.join(ROOT, "org-os", "06-agent-work", "generated", "artifact-registry.yaml")
CONTRACTS = os.path.join(ROOT, "org-os", "06-agent-work", "workflow-contracts.yaml")
SCORECARD = os.path.join(REG, "role-selection-scorecard.yaml")
EXECUTION_POLICY = os.path.join(ROOT, "org-os", "06-agent-work", "execution-policy.yaml")


def _load(path: str) -> dict[str, Any]:
    try:
        return yaml.safe_load(open(path, encoding="utf-8")) or {}
    except Exception:
        return {}


def _registries() -> tuple[dict[str, Any], dict[str, Any], dict[str, Any], dict[str, Any], dict[str, list[str]]]:
    roles_doc = _load(os.path.join(REG, "roles.yaml")).get("role-registry", {}) or {}
    families_doc = _load(os.path.join(REG, "capability-families.yaml")).get("capability-families", {}) or {}
    profiles_doc = _load(os.path.join(REG, "role-profiles.yaml")).get("role-profiles", {}) or {}
    artifacts = _load(ARTIFACT_REGISTRY).get("artifact-registry", {}) or {}
    contracts = _load(CONTRACTS).get("workflow-contracts", {}) or {}
    roles = {str(item["role-id"]): item for item in roles_doc.get("roles", []) or [] if item.get("role-id")}
    families = {str(item["family-id"]): item for item in families_doc.get("families", []) or [] if item.get("family-id")}
    profiles = {str(item["role-id"]): item for item in profiles_doc.get("profiles", []) or [] if item.get("role-id")}
    return roles, families, profiles, artifacts, contracts.get("role-capabilities", {}) or {}


def _candidate_families(profile: dict[str, Any], families: dict[str, Any]) -> list[str]:
    has_explicit = "candidate-families" in profile or "candidate-family" in profile
    explicit = profile.get("candidate-families") if "candidate-families" in profile else profile.get("candidate-family")
    if has_explicit:
        # Unknown ids are intentionally retained here and rejected by the caller;
        # silently dropping them used to turn a malformed explicit plan into an
        # unrelated inferred plan.
        return normalize_family_ids(explicit)
    signal_tokens = tokens(profile.get("signals")) | tokens(profile.get("objective"))
    ranked = []
    for family_id, family in families.items():
        haystack = tokens(family.get("invocation-triggers")) | tokens(family_id)
        overlap = len(signal_tokens & haystack)
        if overlap:
            ranked.append((-overlap, family_id))
    return [family_id for _, family_id in sorted(ranked)[:5]]


def _score(role: dict[str, Any], family: dict[str, Any], coverage: set[str], required: set[str], profile: dict[str, Any]) -> dict[str, int]:
    signal_tokens = tokens(profile.get("signals")) | tokens(profile.get("objective"))
    keyword_hits = len({value.split(":", 1)[1] for value in coverage if value.startswith("keyword:")} & signal_tokens)
    if ("owner" in required
            and family.get("lead-role-id") == role.get("role-id")):
        relevance = 3
    elif f"owner:{family['family-id']}" in required:
        relevance = 3
    elif "owner" in required:
        relevance = min(3, 1 + keyword_hits)
    else:
        relevance = min(3, keyword_hits)
    risk_coverage = min(3, len({value for value in required & coverage if value.startswith("risk:")}))
    evidence_need = 3 if any(value.startswith("artifact:") for value in required & coverage) else 0
    decision_authority = 3 if "authority" in required & coverage else 0
    implementation_impact = 3 if "implementation" in required & coverage else (1 if role.get("is-execution-agent") else 0)
    evidence = profile.get("already-available-evidence", []) or profile.get("existing-evidence", []) or []
    duplicate_penalty = 3 if f"role:{role['role-id']}" in set(evidence) else 0
    total = relevance + risk_coverage + evidence_need + decision_authority + implementation_impact - duplicate_penalty
    return {
        "relevance": relevance,
        "risk-coverage": risk_coverage,
        "evidence-need": evidence_need,
        "decision-authority": decision_authority,
        "implementation-impact": implementation_impact,
        "duplicate-penalty": duplicate_penalty,
        "total": total,
    }


def select_minimum_sufficient_roles(profile: dict[str, Any]) -> dict[str, Any]:
    roles, families, profiles, artifacts, role_capabilities = _registries()
    tier = str(profile.get("tier") or "standard").lower()
    stage = str(profile.get("workflow-stage") or profile.get("stage") or "") or None
    scorecard = _load(SCORECARD).get("role-selection-scorecard", {}) or {}
    execution_policy = _load(EXECUTION_POLICY).get("execution-policy", {}) or {}
    family_ids = _candidate_families(profile, families)
    has_explicit = "candidate-families" in profile or "candidate-family" in profile
    family_errors = (candidate_family_errors(
        family_ids, tier=tier, mode=str(profile.get("mode") or "converge"))
        if has_explicit else [])
    if family_errors:
        empty_estimate = estimate_plan([], tier, stage)
        plan = {
            "version": 1,
            "workflow-id": profile.get("workflow-id"),
            "tier": tier,
            "workflow-stage": stage,
            "candidate-families": family_ids,
            "selected": {"owner": None, "contributors": [], "reviewers": []},
            "skipped": [],
            "coverage": {"required": [], "already-covered": [], "covered": [], "missing": []},
            "estimated-tokens": empty_estimate,
            "budget": {"max-total": workflow_budget(tier, profile.get("token-budget")),
                       "within-budget": True},
            "status": "blocked",
            "errors": family_errors,
        }
        plan["task-graph"] = build_task_graph(plan)
        return {"selection-plan": plan}
    if profile.get("auto-expand-candidates", True):
        family_by_role = {role_id: family_id for family_id, family in families.items()
                          for role_id in family.get("member-role-ids", []) or []}
        artifact_kinds = artifacts.get("artifact-kinds", {}) or {}
        extra_roles = set()
        for kind in profile.get("required-artifacts", []) or []:
            extra_roles |= set((artifact_kinds.get(kind) or {}).get("producer-roles", []) or [])
        risks = profile.get("risks", []) or []
        if isinstance(risks, dict):
            risks = [key for key, value in risks.items() if value]
        for risk in risks:
            extra_roles |= RISK_ROLE_HINTS.get(str(risk).lower(), set())
        capabilities = profile.get("required-capabilities", []) or []
        if isinstance(capabilities, str):
            capabilities = [capabilities]
        for capability in capabilities:
            extra_roles |= CAPABILITY_ROLE_HINTS.get(str(capability).strip().lower(), set())
        if profile.get("authority-required"):
            extra_roles |= {role_id for role_id, role in roles.items() if role.get("is-decision-maker")}
        for role_id in sorted(extra_roles):
            family_id = family_by_role.get(role_id)
            if family_id and family_id not in family_ids:
                family_ids.append(family_id)
    required = required_coverage(profile, family_ids)
    required |= {f"lens:{lens}" for lens in required_lenses(
        family_ids, tier=tier, mode=str(profile.get("mode") or "converge"))}
    existing = {str(item) for item in profile.get("already-available-evidence", []) or profile.get("existing-evidence", []) or []}
    # Evidence may discharge an evidence/artifact need, but it cannot stand in
    # for assigning a concrete capability or divergent lens carrier.
    non_delegable = {item for item in required
                     if item.startswith("capability:") or item.startswith("lens:")}
    effective_existing = existing - non_delegable
    uncovered = required - effective_existing
    excluded_roles = {str(value).upper() for value in profile.get("excluded-role-ids", []) or []}
    candidates = []
    for family_id in family_ids:
        family = families[family_id]
        for role_id in family.get("member-role-ids", []) or []:
            if str(role_id).upper() in excluded_roles:
                continue
            role = roles.get(str(role_id))
            if not role:
                continue
            coverage = role_coverage(role, family, profiles.get(str(role_id)), artifacts, role_capabilities)
            candidates.append({
                "role": role,
                "family": family,
                "coverage": coverage,
                "score": _score(role, family, coverage, required, profile),
            })

    selected: list[dict[str, Any]] = []
    limit = int(profile.get("max-selected-roles") or max_selected_roles(tier))
    budget = workflow_budget(tier, profile.get("token-budget"))
    while uncovered and candidates and len(selected) < limit:
        ranked = []
        for candidate in candidates:
            gain = candidate["coverage"] & uncovered
            if not gain:
                continue
            cost = estimate_role(candidate["role"], tier, stage)["total"]
            ranked.append((
                -(len(gain) * 100000 + candidate["score"]["total"] * 1000 - cost),
                candidate["role"]["role-id"], candidate, gain,
            ))
        if not ranked:
            break
        _, _, chosen, gain = sorted(ranked, key=lambda value: (value[0], value[1]))[0]
        tentative = selected + [chosen]
        if estimate_plan([item["role"] for item in tentative], tier, stage)["total"] > budget:
            break
        selected.append(chosen)
        candidates.remove(chosen)
        uncovered -= gain

    # Independent review is a relational constraint: a producer cannot review its own output.
    producer_ids = {item["role"]["role-id"] for item in selected
                    if "implementation" in item["coverage"] or any(v.startswith("artifact:") for v in item["coverage"] & required)}
    needs_review = "independent-review" in required
    if needs_review and not any(item["role"]["role-id"] not in producer_ids and "independent-review" in item["coverage"] for item in selected):
        reviewer_pool = []
        for family_id, family in families.items():
            if not family.get("audit-capable"):
                continue
            for role_id in family.get("member-role-ids", []) or []:
                if role_id in producer_ids or role_id not in roles:
                    continue
                coverage = role_coverage(roles[role_id], family, profiles.get(role_id), artifacts, role_capabilities)
                risk_gain = len(coverage & required)
                reviewer_pool.append((-risk_gain, role_id, {"role": roles[role_id], "family": family,
                                      "coverage": coverage,
                                      "score": _score(roles[role_id], family, coverage, required, profile)}))
        if reviewer_pool and len(selected) < limit:
            reviewer = sorted(reviewer_pool)[0][2]
            if estimate_plan([item["role"] for item in selected + [reviewer]], tier, stage)["total"] <= budget:
                selected.append(reviewer)
                uncovered.discard("independent-review")

    selected_ids = {item["role"]["role-id"] for item in selected}
    owner_items = [item for item in selected if "owner" in item["coverage"] & required
                   or any(value.startswith("owner:") for value in item["coverage"] & required)]
    owner = owner_items[0] if owner_items else (selected[0] if selected else None)
    reviewers = [item for item in selected if item is not owner and "independent-review" in item["coverage"]]
    contributors = [item for item in selected if item is not owner and item not in reviewers]

    def public(item: dict[str, Any], assignment: str) -> dict[str, Any]:
        return {
            "role-id": item["role"]["role-id"],
            "family-id": item["family"]["family-id"],
            "assignment": assignment,
            "coverage": sorted(item["coverage"] & required),
            "score": item["score"],
            "reason": "minimum marginal coverage under token and independence constraints",
        }

    all_candidate_ids = [role_id for family_id in family_ids
                         for role_id in families[family_id].get("member-role-ids", []) or []]
    skipped = []
    for role_id in all_candidate_ids:
        if role_id in selected_ids:
            continue
        skipped.append({
            "role-id": role_id,
            "reason": "coverage already satisfied by a lower-cost or higher-gain concrete role",
        })
    estimate = estimate_plan([item["role"] for item in selected], tier, stage)
    covered = required - uncovered
    result = {
        "selection-plan": {
            "version": 1,
            "workflow-id": profile.get("workflow-id"),
            "tier": tier,
            "workflow-stage": stage,
            "candidate-families": family_ids,
            "selected": {
                "owner": public(owner, "owner") if owner else None,
                "contributors": [public(item, "contributor") for item in contributors],
                "reviewers": [public(item, "independent-reviewer") for item in reviewers],
            },
            "skipped": skipped,
            "coverage": {
                "required": sorted(required),
                "already-covered": sorted(required & effective_existing),
                "covered": sorted(covered),
                "missing": sorted(uncovered),
            },
            "estimated-tokens": estimate,
            "budget": {"max-total": budget, "within-budget": estimate["total"] <= budget},
            "status": "ready" if not uncovered and estimate["total"] <= budget else "blocked",
            "policy": {
                "score-formula": scorecard.get("total-score-formula"),
                "decision-thresholds": scorecard.get("decision-thresholds"),
                "max-concurrent-role-agents": ((execution_policy.get("wave") or {}).get("max-concurrent-role-agents") or 5),
                "sources": ["role-selection-scorecard.yaml", "execution-policy.yaml", "governance-tiers.yaml"],
            },
        }
    }
    result["selection-plan"]["task-graph"] = build_task_graph(result["selection-plan"])
    return result


def resolve_family(family_id: str, signals: list[str] | None = None, tier: str = "standard") -> dict[str, Any] | None:
    profile = {
        "candidate-families": [family_id],
        "signals": signals or [],
        "tier": tier,
        "independent-review-required": False,
        "max-selected-roles": 1 if tier in {"light", "standard"} else 3,
    }
    plan = select_minimum_sufficient_roles(profile)["selection-plan"]
    owner = (plan.get("selected") or {}).get("owner")
    if not owner:
        return None
    selected = [owner] + (plan["selected"].get("contributors") or [])
    workers = [item["role-id"] for item in selected]
    _, families, _, _, _ = _registries()
    family = families.get(family_id)
    return {
        "requested-family": family_id,
        "resolved-workers": workers,
        "primary-worker": workers[0],
        "available-workers": list((family or {}).get("member-role-ids", []) or []),
        "routing-reason": "minimum-sufficient-coverage",
        "collaboration-default": (family or {}).get("collaboration-default"),
        "selection-plan": plan,
    }