from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, Field


class PartNode(BaseModel):
    id: str
    name: str
    parent_id: str | None = None
    children: list[str] = Field(default_factory=list)


class AxisCandidate(BaseModel):
    id: str
    origin: list[float] = Field(default_factory=lambda: [0.0, 0.0, 0.0])
    direction: list[float] = Field(default_factory=lambda: [0.0, 0.0, 1.0])
    kind: str = "unknown"
    score: float = 0.0
    note: str = ""


class LinkDef(BaseModel):
    name: str
    part_names: list[str] = Field(default_factory=list)
    mesh: str | None = None


class JointDef(BaseModel):
    name: str
    joint_type: Literal["revolute", "prismatic", "fixed", "continuous"] = "revolute"
    parent: str
    child: str
    origin_xyz: list[float] = Field(default_factory=lambda: [0.0, 0.0, 0.0])
    origin_rpy: list[float] = Field(default_factory=lambda: [0.0, 0.0, 0.0])
    axis: list[float] = Field(default_factory=lambda: [0.0, 0.0, 1.0])
    lower: float = -3.14
    upper: float = 3.14
    effort: float = 100.0
    velocity: float = 1.0
    rationale: str = ""


class RobotDraft(BaseModel):
    name: str = "robot"
    profile: str = "generic"
    links: list[LinkDef] = Field(default_factory=list)
    joints: list[JointDef] = Field(default_factory=list)
    notes: list[str] = Field(default_factory=list)
    raw_parts: list[str] = Field(default_factory=list)
    axis_candidates: list[AxisCandidate] = Field(default_factory=list)
    llm_raw: dict[str, Any] | None = None


class ParseResult(BaseModel):
    job_id: str
    filename: str
    root_name: str
    parts: list[PartNode]
    product_names: list[str]
    stats: dict[str, Any] = Field(default_factory=dict)


class ProposeRequest(BaseModel):
    """Zhipu propose: prefer client ``part_names``; optional legacy ``job_id``."""

    job_id: str | None = None
    profile: Literal["a7", "generic"] = "generic"
    robot_name: str = "robot"
    total_mass_kg: float | None = None
    extra_hint: str = ""
    part_names: list[str] = Field(default_factory=list)
    solid_names: list[str] = Field(default_factory=list)
    solids: list[dict[str, Any]] = Field(default_factory=list)


class ExportRequest(BaseModel):
    job_id: str
    draft: RobotDraft
    include_meshes: bool = True