from __future__ import annotations

import math

from ...models import (
    BoundaryCondition,
    ConcreteNode,
    ConcreteRelation,
    ConcreteTopology,
    ReducerParameters,
    ReducerRequirement,
    StageDerived,
    StageParameters,
    StageRequirement,
)


MOTION_MEMBERS = {"sun", "ring", "carrier"}


def stage_component_id(stage_id: str, local_role: str) -> str:
    return f"{stage_id}_{local_role}"


def local_role(component_id: str) -> str:
    return component_id.split("_", 1)[1] if "_" in component_id else component_id


def stage_boundary(stage: StageRequirement) -> BoundaryCondition:
    fixed_members = (
        [stage_component_id(stage.stage_id, stage.fixed)]
        if stage.fixed is not None
        else []
    )
    return BoundaryCondition(
        input=stage_component_id(stage.stage_id, stage.input),
        fixed=fixed_members[0] if len(fixed_members) == 1 else None,
        output=stage_component_id(stage.stage_id, stage.output),
        fixed_members=fixed_members,
    )


def cascade_boundary(requirement: ReducerRequirement) -> BoundaryCondition:
    if not requirement.stages or len(requirement.stages) != 2:
        raise ValueError("simple_2k_h_cascade requires exactly two stages")
    fixed_members = []
    for stage in requirement.stages:
        if stage.fixed is not None:
            fixed_members.append(stage_component_id(stage.stage_id, stage.fixed))
    return BoundaryCondition(
        input=stage_component_id(requirement.stages[0].stage_id, requirement.stages[0].input),
        fixed=fixed_members[0] if len(fixed_members) == 1 else None,
        output=stage_component_id(requirement.stages[1].stage_id, requirement.stages[1].output),
        fixed_members=fixed_members,
    )


def stage_geometry(parameters: StageParameters) -> dict[str, float]:
    sun_pitch_radius = parameters.module_mm * parameters.z_sun / 2.0
    planet_pitch_radius = parameters.module_mm * parameters.z_planet / 2.0
    ring_pitch_radius = parameters.module_mm * parameters.z_ring / 2.0
    center_distance = parameters.module_mm * (parameters.z_sun + parameters.z_planet) / 2.0
    planet_outer_diameter = parameters.module_mm * (
        parameters.z_planet + 2.0 * parameters.addendum_coeff
    )
    neighbor_distance = 2.0 * center_distance * math.sin(math.pi / parameters.planet_count)
    ring_outer_diameter = (
        parameters.module_mm * parameters.z_ring
        + 2.0 * parameters.module_mm * (parameters.addendum_coeff + parameters.clearance_coeff)
        + 2.0 * parameters.ring_rim_thickness_mm
    )
    return {
        "sun_pitch_radius_mm": sun_pitch_radius,
        "planet_pitch_radius_mm": planet_pitch_radius,
        "ring_pitch_radius_mm": ring_pitch_radius,
        "sun_pitch_diameter_mm": sun_pitch_radius * 2.0,
        "planet_pitch_diameter_mm": planet_pitch_radius * 2.0,
        "ring_pitch_diameter_mm": ring_pitch_radius * 2.0,
        "center_distance_mm": center_distance,
        "planet_orbit_radius_mm": center_distance,
        "ring_outer_diameter_mm": ring_outer_diameter,
        "planet_outer_diameter_mm": planet_outer_diameter,
        "planet_neighbor_distance_mm": neighbor_distance,
        "planet_neighbor_clearance_mm": neighbor_distance - planet_outer_diameter,
    }


def stage_derived(
    *,
    parameters: StageParameters,
    ratio: float,
    direction: str,
    speed_mode: str,
    z_offset_mm: float,
) -> StageDerived:
    return StageDerived(
        stage_id=parameters.stage_id,
        ratio=abs(ratio),
        direction=direction,
        speed_mode=speed_mode,
        z_offset_mm=z_offset_mm,
        **stage_geometry(parameters),
    )


def stage_parameter_by_id(parameters: ReducerParameters, stage_id: str) -> StageParameters:
    if not parameters.stages:
        raise ValueError("cascade parameters are missing stages")
    for stage in parameters.stages:
        if stage.stage_id == stage_id:
            return stage
    raise ValueError(f"missing cascade stage parameters: {stage_id}")


def _stage_nodes(stage: StageParameters) -> list[ConcreteNode]:
    nodes = [
        ConcreteNode(
            component_id=stage_component_id(stage.stage_id, "sun"),
            template_node_id=stage_component_id(stage.stage_id, "sun"),
            kind="gear",
            role="sun",
            gear_type="external",
            repeat_index=None,
            stage_id=stage.stage_id,
            local_role="sun",
        ),
        ConcreteNode(
            component_id=stage_component_id(stage.stage_id, "ring"),
            template_node_id=stage_component_id(stage.stage_id, "ring"),
            kind="gear",
            role="ring",
            gear_type="internal",
            repeat_index=None,
            stage_id=stage.stage_id,
            local_role="ring",
        ),
        ConcreteNode(
            component_id=stage_component_id(stage.stage_id, "carrier"),
            template_node_id=stage_component_id(stage.stage_id, "carrier"),
            kind="carrier",
            role="carrier",
            gear_type=None,
            repeat_index=None,
            stage_id=stage.stage_id,
            local_role="carrier",
        ),
    ]
    for index in range(1, stage.planet_count + 1):
        nodes.append(
            ConcreteNode(
                component_id=f"{stage.stage_id}_planet_{index}",
                template_node_id=stage_component_id(stage.stage_id, "planet"),
                kind="gear",
                role="planet",
                gear_type="external",
                repeat_index=index,
                stage_id=stage.stage_id,
                local_role="planet",
            )
        )
    return nodes


def _rel(
    *,
    relation_id: str,
    relation_type: str,
    source: str,
    target: str,
    source_template_node: str,
    target_template_node: str,
    stage_id: str | None = None,
    carrier: str | None = None,
    mesh_id: str | None = None,
) -> ConcreteRelation:
    return ConcreteRelation(
        relation_id=relation_id,
        template_relation_id=relation_id,
        relation_type=relation_type,
        source_component_id=source,
        target_component_id=target,
        source_template_node=source_template_node,
        target_template_node=target_template_node,
        carrier_component_id=carrier,
        mesh_id=mesh_id,
        stage_id=stage_id,
    )


def _stage_relations(stage: StageParameters, requirement_stage: StageRequirement) -> list[ConcreteRelation]:
    sid = stage.stage_id
    sun = stage_component_id(sid, "sun")
    ring = stage_component_id(sid, "ring")
    carrier = stage_component_id(sid, "carrier")
    relations: list[ConcreteRelation] = []
    for index in range(1, stage.planet_count + 1):
        planet = f"{sid}_planet_{index}"
        relations.extend(
            [
                _rel(
                    relation_id=f"{sid}_sun_to_planet_{index}_external_mesh",
                    relation_type="external_mesh",
                    source=sun,
                    target=planet,
                    source_template_node=sun,
                    target_template_node=stage_component_id(sid, "planet"),
                    carrier=carrier,
                    mesh_id=f"{sid}_sun_planet",
                    stage_id=sid,
                ),
                _rel(
                    relation_id=f"{sid}_ring_to_planet_{index}_internal_mesh",
                    relation_type="internal_mesh",
                    source=ring,
                    target=planet,
                    source_template_node=ring,
                    target_template_node=stage_component_id(sid, "planet"),
                    carrier=carrier,
                    mesh_id=f"{sid}_ring_planet",
                    stage_id=sid,
                ),
                _rel(
                    relation_id=f"{sid}_planet_{index}_pin_revolute",
                    relation_type="revolute_joint",
                    source=planet,
                    target=carrier,
                    source_template_node=stage_component_id(sid, "planet"),
                    target_template_node=carrier,
                    carrier=carrier,
                    stage_id=sid,
                ),
            ]
        )

    for local in ["sun", "ring", "carrier"]:
        component = stage_component_id(sid, local)
        relations.append(
            _rel(
                relation_id=f"{sid}_{local}_main_axis_coaxial",
                relation_type="coaxial",
                source=component,
                target="main_axis",
                source_template_node=component,
                target_template_node="main_axis",
                stage_id=sid,
            )
        )
    if requirement_stage.fixed is not None:
        fixed_component = stage_component_id(sid, requirement_stage.fixed)
        relations.append(
            _rel(
                relation_id=f"{sid}_{requirement_stage.fixed}_housing_fixed",
                relation_type="fixed",
                source=fixed_component,
                target="housing",
                source_template_node=fixed_component,
                target_template_node="housing",
                stage_id=sid,
            )
        )
    return relations


def build_cascade_topology(
    requirement: ReducerRequirement,
    parameters: ReducerParameters,
) -> ConcreteTopology:
    if not requirement.stages or not parameters.stages:
        raise ValueError("cascade topology requires requirement.stages and parameters.stages")
    stage_requirements = {stage.stage_id: stage for stage in requirement.stages}
    nodes = [
        ConcreteNode(
            component_id="housing",
            template_node_id="housing",
            kind="housing",
            role="housing",
            gear_type=None,
            repeat_index=None,
            local_role="housing",
        ),
        ConcreteNode(
            component_id="main_axis",
            template_node_id="main_axis",
            kind="axis",
            role="main_axis",
            gear_type=None,
            repeat_index=None,
            local_role="main_axis",
        ),
    ]
    relations: list[ConcreteRelation] = []
    for stage in parameters.stages:
        nodes.extend(_stage_nodes(stage))
        relations.extend(_stage_relations(stage, stage_requirements[stage.stage_id]))

    first_requirement, second_requirement = requirement.stages
    source = stage_component_id(first_requirement.stage_id, first_requirement.output)
    target = stage_component_id(second_requirement.stage_id, second_requirement.input)
    relations.append(
        _rel(
            relation_id=f"{source}_rigid_to_{target}",
            relation_type="rigid",
            source=source,
            target=target,
            source_template_node=source,
            target_template_node=target,
            stage_id="cascade",
        )
    )
    return ConcreteTopology(
        template_id="simple_2kh_cascade_runtime_v0_1",
        family="simple_2k_h_cascade",
        nodes=nodes,
        relations=relations,
    )