parser.py

from __future__ import annotations

import re
from dataclasses import dataclass, asdict


@dataclass
class PromptSpec:
    object_type: str = "cargo_hauler"
    scale: str = "small"
    hull_style: str = "boxy"
    engine_count: int = 2
    wing_span: float = 0.2
    cargo_ratio: float = 0.38
    cockpit_ratio: float = 0.18
    fin_height: float = 0.0
    landing_gear: bool = True
    asymmetry: float = 0.0
    notes: str = ""

    def to_dict(self) -> dict:
        return asdict(self)


TYPE_KEYWORDS = {
    "fighter": "fighter",
    "combat": "fighter",
    "interceptor": "fighter",
    "shuttle": "shuttle",
    "freighter": "freighter",
    "hauler": "cargo_hauler",
    "cargo": "cargo_hauler",
    "transport": "cargo_hauler",
    "dropship": "dropship",
    "drone": "drone",
}

STYLE_KEYWORDS = {
    "boxy": "boxy",
    "industrial": "boxy",
    "hard-surface": "boxy",
    "rounded": "rounded",
    "sleek": "sleek",
    "streamlined": "sleek",
    "brutalist": "boxy",
}

SCALE_KEYWORDS = {
    "tiny": "small",
    "small": "small",
    "compact": "small",
    "medium": "medium",
    "mid-size": "medium",
    "large": "large",
    "heavy": "large",
    "huge": "large",
}

VALID_OBJECT_TYPES = {"cargo_hauler", "fighter", "shuttle", "freighter", "dropship", "drone"}
VALID_SCALES = {"small", "medium", "large"}
VALID_HULL_STYLES = {"boxy", "rounded", "sleek"}


def _clamp(value: float, low: float, high: float) -> float:
    return max(low, min(high, value))


def merge_prompt_specs(primary: PromptSpec, secondary: PromptSpec) -> PromptSpec:
    merged = PromptSpec(**primary.to_dict())

    if secondary.object_type in VALID_OBJECT_TYPES:
        merged.object_type = secondary.object_type
    if secondary.scale in VALID_SCALES:
        merged.scale = secondary.scale
    if secondary.hull_style in VALID_HULL_STYLES:
        merged.hull_style = secondary.hull_style

    merged.engine_count = int(_clamp(secondary.engine_count, 1, 6))
    merged.wing_span = float(_clamp(secondary.wing_span, 0.0, 0.6))
    merged.cargo_ratio = float(_clamp(secondary.cargo_ratio, 0.0, 0.65))
    merged.cockpit_ratio = float(_clamp(secondary.cockpit_ratio, 0.10, 0.30))
    merged.fin_height = float(_clamp(secondary.fin_height, 0.0, 0.3))
    merged.landing_gear = bool(secondary.landing_gear)
    merged.asymmetry = float(_clamp(secondary.asymmetry, 0.0, 0.2))
    merged.notes = secondary.notes or primary.notes

    if merged.object_type in {"fighter", "drone"}:
        merged.cargo_ratio = min(merged.cargo_ratio, 0.20)
        if merged.hull_style == "boxy":
            merged.hull_style = "sleek"

    return merged


def parse_prompt(prompt: str) -> PromptSpec:
    text = prompt.lower().strip()
    spec = PromptSpec(notes=prompt.strip())

    for key, value in TYPE_KEYWORDS.items():
        if key in text:
            spec.object_type = value
            break

    for key, value in STYLE_KEYWORDS.items():
        if key in text:
            spec.hull_style = value
            break

    for key, value in SCALE_KEYWORDS.items():
        if key in text:
            spec.scale = value
            break

    if any(word in text for word in ["wing", "wings"]):
        spec.wing_span = 0.42 if spec.object_type == "fighter" else 0.28
    if any(word in text for word in ["no wings", "wingless"]):
        spec.wing_span = 0.0

    if any(word in text for word in ["cargo bay", "cargo hold", "container", "freight"]):
        spec.cargo_ratio = 0.48

    if any(word in text for word in ["big cockpit", "large cockpit", "glass nose"]):
        spec.cockpit_ratio = 0.24
    if any(word in text for word in ["small cockpit", "tiny cockpit"]):
        spec.cockpit_ratio = 0.13

    if any(word in text for word in ["fin", "tail", "vertical stabilizer"]):
        spec.fin_height = 0.18 if spec.object_type != "fighter" else 0.12

    if any(word in text for word in ["hover", "hovercraft", "antigrav"]):
        spec.landing_gear = False

    if spec.object_type in {"fighter", "drone"}:
        spec.engine_count = 1 if "single engine" in text else 2
        spec.cargo_ratio = min(spec.cargo_ratio, 0.18)
        spec.hull_style = "sleek"
    elif spec.object_type in {"cargo_hauler", "freighter", "dropship"}:
        spec.engine_count = 4 if any(x in text for x in ["4 engine", "four engine", "quad engine"]) else 2
        spec.hull_style = "boxy" if spec.hull_style == "sleek" else spec.hull_style

    numeric_engine = re.search(r"(\d+)\s*(?:engine|engines)", text)
    if numeric_engine:
        spec.engine_count = max(1, min(6, int(numeric_engine.group(1))))

    if any(word in text for word in ["asymmetric", "uneven", "offset"]):
        spec.asymmetry = 0.12

    return spec