"""Deterministic role-instance model assignment and host binding."""
from __future__ import annotations

from dataclasses import dataclass
from types import MappingProxyType
from typing import Literal, Mapping, Sequence

from .application.resolve_assignment import (
    resolve_assignment_runner,
    resolve_lead_provider_id,
    resolve_lead_provider_from_specs,
)
from .domain.host import HostDescriptor, HostSessionContext
from .domain.provider import (
    ExecutionNormalization,
    HostModelBinding,
    LeadProviderAssignment,
)
from .domain.role import normalize_role
from .domain.worker_exec import WorkerWriteCapability
from .model_defaults import (
    ModelDefaultError,
    ModelDefaultScopes,
    default_candidates,
    default_scope,
)
from .model_pool import ModelPool, ResolvedModel
from .ports.host_model import (
    HostModelBindingError,
    HostModelBindingPort,
    HostModelBindingRequest,
)
from .role_requirements import RoleProfile, RoleRequirement


TerminalBackend = Literal[
    "cli-wrapper",
    "cmux-pane",
    "mixed",
]
_TERMINAL_BACKENDS = frozenset({
    "cli-wrapper",
    "cmux-pane",
    "mixed",
})


@dataclass(frozen=True)
class AssignmentEnvironment:
    """Immutable host and dispatch facts consumed by assignment resolution."""

    host_descriptor: HostDescriptor
    host_model_port: HostModelBindingPort
    native_worker_write_capability: WorkerWriteCapability
    terminal_backend: TerminalBackend
    execution_resolutions: Mapping[
        tuple[str, str],
        ExecutionNormalization,
    ]

    def __post_init__(self) -> None:
        if self.terminal_backend not in _TERMINAL_BACKENDS:
            allowed = ", ".join(sorted(_TERMINAL_BACKENDS))
            raise ValueError(
                f"unknown terminal backend {self.terminal_backend!r}. "
                f"Allowed values: {allowed}"
            )
        resolutions: dict[tuple[str, str], ExecutionNormalization] = {}
        for key, resolution in self.execution_resolutions.items():
            if (
                not isinstance(key, tuple)
                or len(key) != 2
                or not all(isinstance(value, str) and value for value in key)
                or not isinstance(resolution, ExecutionNormalization)
            ):
                raise ValueError(
                    "execution resolutions require (model_ref, role) keys and "
                    "ExecutionNormalization values"
                )
            resolutions[key] = resolution
        object.__setattr__(
            self,
            "execution_resolutions",
            MappingProxyType(resolutions),
        )


@dataclass(frozen=True)
class AssignmentContext:
    pool: ModelPool
    environment: AssignmentEnvironment


@dataclass(frozen=True)
class RoleCountSelection:
    role: str
    count: int


@dataclass(frozen=True)
class RoleModelSelection:
    role: str
    model_refs: tuple[str, ...]


@dataclass(frozen=True)
class RoleInstance:
    role: str
    duty_id: str
    ordinal: int


@dataclass(frozen=True)
class CandidateAssessment:
    role: str
    ordinal: int
    model_ref: str
    scope: str
    available: bool
    reason: str = ""


@dataclass(frozen=True)
class UnsatisfiedConstraint:
    role: str
    constraint: str
    required: int
    available: int


@dataclass(frozen=True)
class ResolvedAssignment:
    instance: RoleInstance
    model_ref: str | None
    provider_id: str
    model_id: str
    display_name: str
    host_runtime: str
    entry_mode: Literal["current-session", "new-session"]
    binding: HostModelBinding | None

    @property
    def role(self) -> str:
        return self.instance.role

    @property
    def duty_id(self) -> str:
        return self.instance.duty_id

    @property
    def ordinal(self) -> int:
        return self.instance.ordinal

    @property
    def display(self) -> str:
        return self.display_name

    @property
    def execution(self) -> str:
        return self.binding.resolved_execution_value if self.binding else ""

    @property
    def execution_label(self) -> str:
        return f"{self.role}-{self.provider_id}-{self.model_id}-{self.ordinal}"


@dataclass(frozen=True)
class AssignmentPlan:
    assignments: tuple[ResolvedAssignment, ...]
    candidates: tuple[CandidateAssessment, ...] = ()

    def by_role(self, role: str) -> tuple[ResolvedAssignment, ...]:
        return tuple(row for row in self.assignments if row.role == role)


class AssignmentResolutionError(ValueError):
    """Raised with all evaluated candidates and failed constraints."""

    def __init__(
        self,
        message: str,
        *,
        candidates: Sequence[CandidateAssessment] = (),
        unsatisfied_constraints: Sequence[UnsatisfiedConstraint] = (),
    ) -> None:
        self.candidates = tuple(candidates)
        self.unsatisfied_constraints = tuple(unsatisfied_constraints)
        super().__init__(message)


def _validate_environment(
    host: HostSessionContext,
    environment: AssignmentEnvironment,
) -> None:
    if environment.host_descriptor.id != host.host_id:
        raise AssignmentResolutionError(
            "assignment environment host does not match the session context: "
            f"{environment.host_descriptor.id!r} != {host.host_id!r}"
        )


def resolve_assignments(
    *,
    profile: RoleProfile,
    role_counts: Mapping[str, int],
    role_models: Mapping[str, tuple[str, ...]],
    scopes: ModelDefaultScopes,
    pool: ModelPool,
    host: HostSessionContext,
    environment: AssignmentEnvironment,
) -> AssignmentPlan:
    """Resolve the first complete assignment in deterministic DFS order."""
    _validate_environment(host, environment)
    requirements = _static_requirements(profile)
    instances = _role_instances(requirements, role_counts)
    _validate_model_selections(instances, role_models)
    current_leader = _current_leader_assignment(
        instances,
        role_models,
        pool,
        host,
        environment,
    )
    options, assessments = _assignment_options(
        instances,
        role_models,
        scopes,
        pool,
        host,
        profile,
        requirements,
        current_leader,
        environment,
    )
    resolved = _first_complete_assignment(options, requirements)
    if resolved is None:
        constraints = _unmet_constraints(options, requirements)
        raise AssignmentResolutionError(
            "no complete role assignment satisfies the profile constraints",
            candidates=assessments,
            unsatisfied_constraints=constraints,
        )
    return AssignmentPlan(tuple(resolved), tuple(assessments))


def resolve_model_assignment(
    *,
    instance: RoleInstance,
    model_ref: str,
    pool: ModelPool,
    host: HostSessionContext,
    environment: AssignmentEnvironment,
    entry_mode: Literal["current-session", "new-session"] = "new-session",
) -> ResolvedAssignment:
    """Resolve one canonical model through the selected host execution path."""
    _validate_environment(host, environment)
    model = pool.resolve(model_ref)
    availability = pool.availability(
        model_ref,
        instance.role,
        host.host_id,
        entry_mode,
    )
    if not availability.available:
        raise AssignmentResolutionError(
            f"model {model_ref!r} is unavailable: {availability.reason}"
        )
    binding = _first_supported_binding(
        instance,
        model,
        pool,
        environment,
    )
    return _resolved_assignment(instance, model, host, entry_mode, binding)


def resolve_dispatch_assignment(
    *,
    context: AssignmentContext,
    host_runtime: str,
    role: str,
    duty_id: str,
    provider: str,
    model: str = "",
) -> ResolvedAssignment:
    """Resolve one legacy provider/model input through the shared core."""
    canonical_role = normalize_role(role)
    pool = context.pool
    selected = resolve_model_selection(pool, provider, canonical_role, model)
    host = HostSessionContext(
        host_id=host_runtime,
        entry_mode="spawn-process",
        available_functions=frozenset(),
        interaction_surface="application-dispatch",
    )
    return resolve_model_assignment(
        instance=RoleInstance(canonical_role, duty_id, 1),
        model_ref=str(selected.model_ref),
        pool=pool,
        host=host,
        environment=context.environment,
    )


def resolve_context_lead_provider(
    *,
    context: AssignmentContext,
    requested_provider: str,
) -> LeadProviderAssignment:
    """Resolve a lead provider from the same immutable assignment facts."""
    host = context.environment.host_descriptor
    provider_id = resolve_lead_provider_id(
        host=host,
        requested_provider=requested_provider,
    )
    provider = context.pool.resolve_provider(provider_id)
    return resolve_lead_provider_from_specs(
        host=host,
        requested_provider=requested_provider,
        provider=provider,
    )


def resolve_model_selection(
    pool: ModelPool,
    provider: str,
    role: str,
    raw_model: str,
) -> ResolvedModel:
    """Resolve one explicit or provider-default catalog model without binding."""
    if raw_model:
        return pool.resolve_alias(provider, raw_model)
    selected = pool.default_candidate(provider, role)
    if selected is None:
        raise AssignmentResolutionError(
            f"provider {provider!r} has no default model for role {role!r}"
        )
    return selected


def _static_requirements(profile: RoleProfile) -> tuple[RoleRequirement, ...]:
    leader = RoleRequirement("leader", 1, 1, 1, "lead")
    return (leader, *(row for row in profile.roles if not row.dynamic))


def _role_instances(
    requirements: Sequence[RoleRequirement],
    role_counts: Mapping[str, int],
) -> tuple[RoleInstance, ...]:
    allowed = {row.role for row in requirements}
    unknown = tuple(sorted(set(role_counts) - allowed))
    if unknown:
        raise AssignmentResolutionError(
            f"role counts contain roles outside the profile: {', '.join(unknown)}"
        )
    instances: list[RoleInstance] = []
    for requirement in requirements:
        fixed = requirement.min_count == requirement.max_count
        if requirement.role in role_counts and fixed:
            raise AssignmentResolutionError(
                f"role {requirement.role!r} has a fixed quantity and cannot "
                "receive an explicit count"
            )
        count = role_counts.get(requirement.role, requirement.recommended_count)
        if count < requirement.min_count or count > requirement.max_count:
            raise AssignmentResolutionError(
                f"role {requirement.role!r} count must be in "
                f"{requirement.min_count}..{requirement.max_count}: {count}"
            )
        instances.extend(
            RoleInstance(requirement.role, requirement.duty, ordinal)
            for ordinal in range(1, count + 1)
        )
    return tuple(instances)


def _validate_model_selections(
    instances: Sequence[RoleInstance],
    role_models: Mapping[str, tuple[str, ...]],
) -> None:
    counts = {
        role: sum(instance.role == role for instance in instances)
        for role in {instance.role for instance in instances}
    }
    unknown = tuple(sorted(set(role_models) - counts.keys()))
    if unknown:
        raise AssignmentResolutionError(
            f"role models contain roles outside the profile: {', '.join(unknown)}"
        )
    for role, models in role_models.items():
        if len(models) > counts[role]:
            raise AssignmentResolutionError(
                f"role {role!r} has {len(models)} models for {counts[role]} instances"
            )


def _current_leader_assignment(
    instances: Sequence[RoleInstance],
    role_models: Mapping[str, tuple[str, ...]],
    pool: ModelPool,
    host: HostSessionContext,
    environment: AssignmentEnvironment,
) -> ResolvedAssignment | None:
    if host.entry_mode != "current-session":
        return None
    leader = next(row for row in instances if row.role == "leader")
    attestation = host.current_model
    model = _validated_current_model(pool, host)
    explicit = role_models.get("leader", ())
    if explicit and (model is None or explicit[0] != str(model.model_ref)):
        return None
    if model is None:
        provider_id = (
            attestation.provider_id
            or environment.host_descriptor.native_provider_id
        )
        return ResolvedAssignment(
            leader,
            None,
            provider_id,
            "unknown",
            "unknown",
            host.host_id,
            "current-session",
            None,
        )
    binding = HostModelBinding(
        runner="current-session",
        catalog_execution_value=model.execution_value,
        resolved_execution_value=model.execution_value,
        host_model_value=attestation.observed_model,
        binding_fidelity=attestation.level,
        worker_write_capability=None,
    )
    return _resolved_assignment(leader, model, host, "current-session", binding)


def _validated_current_model(
    pool: ModelPool,
    host: HostSessionContext,
) -> ResolvedModel | None:
    attestation = host.current_model
    if attestation.level == "unknown":
        return None
    model_ref = attestation.normalized_model_ref or ""
    try:
        model = pool.resolve(model_ref)
    except ValueError as exc:
        action = "update the provider adapter catalog or start a registered leader model"
        raise AssignmentResolutionError(
            "unregistered current-session model attestation: "
            f"provider={attestation.provider_id!r}, "
            f"model={attestation.observed_model!r}, source={attestation.source!r}; "
            f"{action}"
        ) from exc
    if model.provider_id != attestation.provider_id:
        raise AssignmentResolutionError(
            "current-session model attestation provider does not match its model ref"
        )
    if attestation.level == "exact" and model.version_kind != "pinned":
        raise AssignmentResolutionError(
            "exact current-session attestation must reference a pinned model"
        )
    if attestation.level == "channel" and model.version_kind != "channel":
        raise AssignmentResolutionError(
            "channel current-session attestation must reference a channel model"
        )
    availability = pool.availability(
        str(model.model_ref),
        "leader",
        host.host_id,
        "current-session",
    )
    if not availability.available:
        raise AssignmentResolutionError(
            "current-session model is unavailable for leader: "
            f"{availability.reason}"
        )
    return model


def _assignment_options(
    instances: Sequence[RoleInstance],
    role_models: Mapping[str, tuple[str, ...]],
    scopes: ModelDefaultScopes,
    pool: ModelPool,
    host: HostSessionContext,
    profile: RoleProfile,
    requirements: Sequence[RoleRequirement],
    current_leader: ResolvedAssignment | None,
    environment: AssignmentEnvironment,
) -> tuple[list[tuple[ResolvedAssignment, ...]], list[CandidateAssessment]]:
    options: list[tuple[ResolvedAssignment, ...]] = []
    assessments: list[CandidateAssessment] = []
    for instance in instances:
        if current_leader is not None and instance.role == "leader":
            options.append((current_leader,))
            continue
        rows, evaluated = _instance_options(
            instance,
            role_models,
            scopes,
            pool,
            host,
            profile,
            environment,
        )
        options.append(rows)
        assessments.extend(evaluated)
    if any(not rows for rows in options):
        raise AssignmentResolutionError(
            "one or more role instances have no executable model candidates",
            candidates=assessments,
            unsatisfied_constraints=_unmet_constraints(options, requirements),
        )
    return options, assessments


def _instance_options(
    instance: RoleInstance,
    role_models: Mapping[str, tuple[str, ...]],
    scopes: ModelDefaultScopes,
    pool: ModelPool,
    host: HostSessionContext,
    profile: RoleProfile,
    environment: AssignmentEnvironment,
) -> tuple[tuple[ResolvedAssignment, ...], list[CandidateAssessment]]:
    explicit = role_models.get(instance.role, ())
    if instance.ordinal <= len(explicit):
        candidates, scope = (explicit[instance.ordinal - 1],), "explicit"
    else:
        candidates, scope = _rotated_defaults(instance, scopes, pool)
    resolved: list[ResolvedAssignment] = []
    assessments: list[CandidateAssessment] = []
    for model_ref in candidates:
        try:
            assignment = resolve_model_assignment(
                instance=instance,
                model_ref=model_ref,
                pool=pool,
                host=host,
                environment=environment,
            )
            resolved.append(assignment)
            assessments.append(_assessment(instance, model_ref, scope, True, ""))
        except (AssignmentResolutionError, ValueError, RuntimeError) as exc:
            assessments.append(
                _assessment(instance, model_ref, scope, False, str(exc))
            )
    if scope != "bundled" and any(not row.available for row in assessments):
        return (), assessments
    return tuple(resolved), assessments


def _rotated_defaults(
    instance: RoleInstance,
    scopes: ModelDefaultScopes,
    pool: ModelPool,
) -> tuple[tuple[str, ...], str]:
    try:
        candidates = default_candidates(instance.role, scopes)
    except ModelDefaultError as exc:
        raise AssignmentResolutionError(str(exc)) from exc
    scope = default_scope(instance.role, scopes)
    if not candidates:
        candidates = tuple(
            str(model.model_ref) for model in pool.default_candidates(instance.role)
        )
        scope = "bundled"
    if not candidates:
        raise AssignmentResolutionError(
            f"role {instance.role!r} has no default model candidates"
        )
    offset = (instance.ordinal - 1) % len(candidates)
    return (*candidates[offset:], *candidates[:offset]), scope or "bundled"


def _assessment(
    instance: RoleInstance,
    model_ref: str,
    scope: str,
    available: bool,
    reason: str,
) -> CandidateAssessment:
    return CandidateAssessment(
        instance.role,
        instance.ordinal,
        model_ref,
        scope,
        available,
        reason,
    )


def _first_supported_binding(
    instance: RoleInstance,
    model: ResolvedModel,
    pool: ModelPool,
    environment: AssignmentEnvironment,
) -> HostModelBinding:
    provider = pool.resolve_provider(model.provider_id)
    errors: list[str] = []
    for runner in _runner_candidates(instance, environment, provider):
        try:
            resolve_assignment_runner(
                host=environment.host_descriptor,
                provider=provider,
                role=instance.role,
                requested_runner=runner,
            )
            binding = _binding_for_runner(
                instance,
                model,
                provider,
                environment,
                runner,
            )
            _validate_binding_fidelity(model, binding)
            return binding
        except (ValueError, RuntimeError, AttributeError) as exc:
            errors.append(f"{runner}: {exc}")
    raise AssignmentResolutionError(
        f"model {model.model_ref!r} has no executable binding: {'; '.join(errors)}"
    )


def _runner_candidates(
    instance: RoleInstance,
    environment: AssignmentEnvironment,
    provider,
) -> tuple[str, ...]:
    if (
        instance.role != "leader"
        and environment.terminal_backend == "cmux-pane"
    ):
        return ("cli-wrapper",) if provider.wrapper else ()
    runners: list[str] = []
    if provider.provider == environment.host_descriptor.native_provider_id:
        runners.append("native-session")
    if provider.wrapper:
        runners.append("cli-wrapper")
    return tuple(runners or ("cli-wrapper",))


def _binding_for_runner(
    instance: RoleInstance,
    model: ResolvedModel,
    provider,
    environment: AssignmentEnvironment,
    runner: str,
) -> HostModelBinding:
    execution = model.execution_value
    if runner == "cli-wrapper":
        resolution = environment.execution_resolutions.get(
            (str(model.model_ref), instance.role)
        )
        if resolution is None and provider.execution_normalization_required:
            raise AssignmentResolutionError(
                f"missing execution snapshot for model {model.model_ref!r} "
                f"and role {instance.role!r}"
            )
        if resolution is not None:
            if resolution.error:
                raise AssignmentResolutionError(resolution.error)
            execution = resolution.resolved_execution_value
    capability = _worker_write_capability(
        instance,
        provider,
        environment,
        runner,
    )
    request = HostModelBindingRequest(
        host_runtime=environment.host_descriptor.id,
        provider=model.provider_id,
        model_ref=str(model.model_ref),
        version_kind=model.version_kind,
        model_execution_value=model.execution_value,
        resolved_execution_value=execution,
        runner=runner,
        worker_write_capability=capability,
    )
    try:
        return environment.host_model_port.resolve(request)
    except HostModelBindingError as exc:
        raise AssignmentResolutionError(str(exc)) from exc


def _worker_write_capability(instance, provider, environment, runner):
    if instance.role == "leader":
        return None
    if runner == "native-session":
        return environment.native_worker_write_capability
    if provider.cli_worker_write_capability is None:
        raise AssignmentResolutionError(
            f"provider {provider.provider!r} has no worker execution strategy"
        )
    return provider.cli_worker_write_capability


def _validate_binding_fidelity(
    model: ResolvedModel,
    binding: HostModelBinding,
) -> None:
    allowed = {"exact"} if model.version_kind == "pinned" else {"exact", "channel"}
    if binding.binding_fidelity not in allowed:
        raise AssignmentResolutionError(
            f"model {model.model_ref!r} cannot bind with fidelity "
            f"{binding.binding_fidelity!r}"
        )


def _resolved_assignment(instance, model, host, entry_mode, binding):
    return ResolvedAssignment(
        instance=instance,
        model_ref=str(model.model_ref),
        provider_id=model.provider_id,
        model_id=model.model_ref.model_id,
        display_name=model.display_name,
        host_runtime=host.host_id,
        entry_mode=entry_mode,
        binding=binding,
    )


def _first_complete_assignment(
    options: Sequence[Sequence[ResolvedAssignment]],
    requirements: Sequence[RoleRequirement],
) -> list[ResolvedAssignment] | None:
    chosen: list[ResolvedAssignment] = []

    def visit(index: int) -> list[ResolvedAssignment] | None:
        if index == len(options):
            return list(chosen) if not _constraint_failures(chosen, requirements) else None
        for assignment in options[index]:
            chosen.append(assignment)
            if (complete := visit(index + 1)) is not None:
                return complete
            chosen.pop()
        return None

    return visit(0)


def _constraint_failures(
    assignments: Sequence[ResolvedAssignment],
    requirements: Sequence[RoleRequirement],
) -> tuple[UnsatisfiedConstraint, ...]:
    failures: list[UnsatisfiedConstraint] = []
    for requirement in requirements:
        role_rows = [row for row in assignments if row.role == requirement.role]
        if not role_rows:
            continue
        refs = {row.model_ref for row in role_rows}
        if len(refs) < len(role_rows):
            failures.append(UnsatisfiedConstraint(
                requirement.role,
                "uniqueModelRefs",
                len(role_rows),
                len(refs),
            ))
    return tuple(failures)


def _unmet_constraints(
    options: Sequence[Sequence[ResolvedAssignment]],
    requirements: Sequence[RoleRequirement],
) -> tuple[UnsatisfiedConstraint, ...]:
    failures: list[UnsatisfiedConstraint] = []
    for requirement in requirements:
        role_groups = [
            candidates
            for candidates in options
            if candidates and candidates[0].role == requirement.role
        ]
        if not role_groups:
            continue
        worker_count = len(role_groups)
        unique_refs = {
            row.model_ref
            for candidates in role_groups
            for row in candidates
        }
        if len(unique_refs) < worker_count:
            failures.append(UnsatisfiedConstraint(
                requirement.role,
                "uniqueModelRefs",
                worker_count,
                len(unique_refs),
            ))
    return tuple(failures)