src/tool.py

"""Random encounter generator: input handling, generation, rendering, export."""

import html
import random
import tempfile
from pathlib import Path

from .shared import ai_client
from .shared.ratelimit import RateLimiter
from .prompts import SCHEMA_SPEC, SYSTEM_PROMPT, DIFFICULTY_GUIDE, build_user_prompt
from .fixtures import FIXTURE_FALSE_LANTERNS

limiter = RateLimiter()

ENVIRONMENTS = [
    "Forest", "Mountains", "Swamp", "Desert", "Arctic", "Coast", "Open Sea",
    "Grassland", "Underdark", "Urban Streets", "Dungeon", "Graveyard", "Planar",
]
LEVELS = [str(n) for n in range(1, 21)]
PARTY_SIZES = [str(n) for n in range(2, 8)]
FLAVORS = ["Combat", "Social", "Exploration", "Mixed"]
DIFFICULTIES = list(DIFFICULTY_GUIDE.keys())  # Easy, Medium, Hard, Deadly
TIMES = ["Day", "Night", "Dawn/Dusk"]

MAX_FIELD = 300
MAX_NOTES = 800


def randomize():
    """Fill the core inputs with a random rollable combination."""
    return (
        random.choice(ENVIRONMENTS),
        random.choice(LEVELS),
        random.choice(FLAVORS),
        random.choice(TIMES),
    )


def _clip(value: str, limit: int) -> str:
    return (value or "").strip()[:limit]


def generate(
    environment, level, size, flavor, difficulty, time_of_day, notes,
    request=None,
):
    """Returns (markdown, plain_text, download_path, error_message)."""
    allowed, message = limiter.check(request)
    if not allowed:
        return None, None, None, message

    difficulty = difficulty if difficulty in DIFFICULTY_GUIDE else "Medium"
    level = level if str(level) in LEVELS else ""
    size = size if str(size) in PARTY_SIZES else ""
    flavor = flavor if flavor in FLAVORS else ""
    time_of_day = time_of_day if time_of_day in TIMES else ""
    fields = {
        "Environment": _clip(environment, MAX_FIELD),
        "Party level": str(level),
        "Party size": str(size),
        "Encounter flavor": flavor,
        "Difficulty feel": difficulty,
        "Time of day": time_of_day,
        "DM notes": _clip(notes, MAX_NOTES),
    }
    user_prompt = build_user_prompt(fields) + "\n\n" + DIFFICULTY_GUIDE[difficulty]

    try:
        data = ai_client.generate_json(
            SYSTEM_PROMPT, user_prompt, SCHEMA_SPEC,
            fixture=FIXTURE_FALSE_LANTERNS,
        )
    except ai_client.AIUnavailable as exc:
        return None, None, None, str(exc)

    md = render_markdown(data)
    path = Path(tempfile.mkdtemp()) / "dnd-encounter.md"
    path.write_text(md, encoding="utf-8")
    return md, md, str(path), None


def render_markdown(d: dict) -> str:
    return f"""## {d['title']}

*{d['encounter_type']}*

### The Setup
{d['setup']}

### Composition
{d['composition']}

### Why They're Here
{d['motivation']}

### Terrain & Features
{d['terrain_and_features']}

### Complication
{d['complication']}

### Scaling It
{d['scaling_notes']}

### Outcome & Treasure
{d['outcome_and_treasure']}

### Follow-Up Hook
{d['followup_hook']}
"""


def render_error(message: str) -> str:
    return (
        '<div class="lf-output" role="alert" style="border-color:#A7343B;">'
        f"<p>{html.escape(message)}</p></div>"
    )