from dataclasses import dataclass

from openai import OpenAI
from pydantic import BaseModel

import config
from config import SERVOS, HOME, clamp_pose
from presets import EXAMPLES

MAX_STEPS = 16
MIN_MS, MAX_MS = 80, 6000


class Step(BaseModel):
    s1: int
    s2: int
    s3: int
    s4: int
    ms: int


class Reply(BaseModel):
    summary: str
    steps: list[Step]


@dataclass
class Plan:
    summary: str
    steps: list  # [(pose, ms), ...]


def _fmt(steps):
    return "\n".join(f"  {list(pose)} over {ms}ms" for pose, ms in steps)


def _instructions():
    servos = "\n".join(
        f"- {s.key} ({s.name}): {s.role}. range {s.lo}-{s.hi}, home {s.home}"
        for s in SERVOS
    )
    examples = "\n\n".join(f'"{name}":\n{_fmt(steps)}' for name, steps in EXAMPLES.items())
    return f"""You drive a small 4-servo robot arm. Turn the user's spoken command into keyframes.

Servos (angles in degrees):
{servos}

Each step is an absolute pose [s1, s2, s3, s4] reached by moving all servos together over `ms` milliseconds.
Repeat the same pose with a duration to hold still.

Rules:
- Stay inside each servo's range.
- {MIN_MS}-{MAX_MS} ms per step, at most {MAX_STEPS} steps. ~4 degrees per 30 ms looks natural; go faster only for snappy moves.
- Unless the user asks to hold a position, finish back at home {list(HOME)}.
- Relative commands ("a bit more left", "lower") start from the current pose.
- If the command has nothing to do with moving the arm, return no steps and say why in the summary.
- summary: a few words describing the motion, in the user's language.

Reference moves:

{examples}"""


class Planner:
    def __init__(self):
        if not config.OPENAI_API_KEY:
            raise RuntimeError("OPENAI_API_KEY is not set")
        self._client = OpenAI(api_key=config.OPENAI_API_KEY)

    def plan(self, text, current_pose):
        kwargs = {}
        if config.OPENAI_REASONING_EFFORT:
            kwargs["reasoning"] = {"effort": config.OPENAI_REASONING_EFFORT}

        resp = self._client.responses.parse(
            model=config.OPENAI_MODEL,
            instructions=_instructions(),
            input=f"Current pose: {list(current_pose)}\nCommand: {text}",
            text_format=Reply,
            timeout=30,
            **kwargs,
        )
        reply = resp.output_parsed
        if reply is None:
            raise RuntimeError("model returned nothing usable")

        steps = [
            (clamp_pose((s.s1, s.s2, s.s3, s.s4)), max(MIN_MS, min(MAX_MS, s.ms)))
            for s in reply.steps[:MAX_STEPS]
        ]
        return Plan(reply.summary.strip(), steps)