state.py

# ─────────────────────────────────────────────
#  state.py  –  pure game logic, no Gradio
# ─────────────────────────────────────────────
from __future__ import annotations

import math
import random
from typing import Any

from game_config import (
    BASE_REWARD,
    BASE_TRAIN_COST,
    CHAT_TEMP_MAX,
    CHAT_TEMP_MIN,
    CLOUD_BASE_REVENUE,
    CLOUD_QUALITY_MULT,
    CLOUD_UNLOCK_COST,
    MAX_QUALITY_REWARD_MULT,
    QUALITY_EXPONENT,
    REWARD_NOISE_PCT,
    STARTING_MONEY,
    TRAIN_COST_QUALITY_SCALE,
    UPGRADES,
)

# ── Helpers ──────────────────────────────────


def _max_quality() -> float:
    return sum(upg["levels"][-1]["quality_score"] for upg in UPGRADES.values())


def _quality_ratio(state: dict) -> float:
    total = sum(
        UPGRADES[k]["levels"][state["upgrades"][k]]["quality_score"] for k in UPGRADES
    )
    mx = _max_quality()
    return total / mx if mx > 0 else 0.0


def _total_quality_score(state: dict) -> int:
    return sum(
        UPGRADES[k]["levels"][state["upgrades"][k]]["quality_score"] for k in UPGRADES
    )


# ── State factory ────────────────────────────


def new_game() -> dict:
    return {
        "money": STARTING_MONEY,
        "upgrades": {k: 0 for k in UPGRADES},  # current tier index per upgrade
        "trained": False,  # has trained at least once
        "cloud_unlocked": False,
        "cloud_active": False,
        "runs": 0,
        "total_earned": 0.0,
        "log": [],  # list of log-line strings
    }


# ── Actions ──────────────────────────────────


def buy_upgrade(state: dict, upgrade_key: str) -> tuple[dict, str]:
    """Buy the next tier of an upgrade. Returns (new_state, message)."""
    state = _copy(state)
    upg = UPGRADES[upgrade_key]
    current = state["upgrades"][upgrade_key]
    max_tier = len(upg["levels"]) - 1

    if current >= max_tier:
        return state, f"[WARN] {upg['label']} is already maxed."

    next_tier = current + 1
    cost = upg["levels"][next_tier]["cost"]

    if state["money"] < cost:
        return (
            state,
            f"[FAIL] Not enough funds. Need ${cost:.0f}, have ${state['money']:.2f}.",
        )

    state["money"] -= cost
    state["upgrades"][upgrade_key] = next_tier
    tier_name = upg["levels"][next_tier]["name"]
    msg = f"[OK]   {upg['label']} → {tier_name}  (−${cost:.0f})"
    state["log"].append(msg)
    return state, msg


def train_model(state: dict) -> tuple[dict, str]:
    """Run one training pass. Returns (new_state, multiline log)."""
    state = _copy(state)
    qr = _quality_ratio(state)
    qs = _total_quality_score(state)

    # Cost
    cost = BASE_TRAIN_COST + TRAIN_COST_QUALITY_SCALE * qs
    if state["money"] < cost:
        msg = f"[FAIL] Insufficient funds for training run. Need ${cost:.2f}, have ${state['money']:.2f}."
        state["log"].append(msg)
        return state, msg

    # Reward
    raw_mult = 1.0 + (MAX_QUALITY_REWARD_MULT - 1.0) * (qr**QUALITY_EXPONENT)
    noise = 1.0 + random.uniform(-REWARD_NOISE_PCT, REWARD_NOISE_PCT)
    reward = BASE_REWARD * raw_mult * noise

    state["money"] += reward - cost
    state["runs"] += 1
    state["total_earned"] += reward
    state["trained"] = True

    lines = [
        f"[RUN #{state['runs']:04d}] ──────────────────────",
        f"  Cost:      −${cost:.2f}",
        f"  Reward:    +${reward:.2f}",
        f"  Net:        ${reward - cost:+.2f}",
        f"  Balance:   ${state['money']:.2f}",
        f"  Quality:   {qs}/{_max_quality():.0f}  ({qr * 100:.1f}%)",
    ]
    for l in lines:
        state["log"].append(l)
    return state, "\n".join(lines)


def unlock_cloud(state: dict) -> tuple[dict, str]:
    state = _copy(state)
    if state["cloud_unlocked"]:
        return state, "[WARN] Cloud Inference already unlocked."
    if state["money"] < CLOUD_UNLOCK_COST:
        return (
            state,
            f"[FAIL] Need ${CLOUD_UNLOCK_COST:.0f} to unlock Cloud Inference. Have ${state['money']:.2f}.",
        )
    state["money"] -= CLOUD_UNLOCK_COST
    state["cloud_unlocked"] = True
    state["cloud_active"] = True
    msg = f"[OK]   Cloud Inference unlocked & active.  (−${CLOUD_UNLOCK_COST:.0f})"
    state["log"].append(msg)
    return state, msg


def cloud_tick(state: dict) -> tuple[dict, float]:
    """Called periodically when cloud is active. Returns (new_state, revenue_this_tick)."""
    if not state["cloud_active"]:
        return state, 0.0
    state = _copy(state)
    qs = _total_quality_score(state)
    revenue = CLOUD_BASE_REVENUE + CLOUD_QUALITY_MULT * qs
    state["money"] += revenue
    state["total_earned"] += revenue
    return state, revenue


def toggle_cloud(state: dict) -> tuple[dict, str]:
    state = _copy(state)
    if not state["cloud_unlocked"]:
        return state, "[FAIL] Cloud Inference not unlocked yet."
    state["cloud_active"] = not state["cloud_active"]
    status = "ACTIVE" if state["cloud_active"] else "PAUSED"
    msg = f"[OK]   Cloud Inference → {status}"
    state["log"].append(msg)
    return state, msg


# ── Chat helpers ─────────────────────────────


def chat_temperature(state: dict) -> float:
    qr = _quality_ratio(state)
    return CHAT_TEMP_MAX + (CHAT_TEMP_MIN - CHAT_TEMP_MAX) * qr


def quality_ratio(state: dict) -> float:
    return _quality_ratio(state)


def status_lines(state: dict) -> str:
    qs = _total_quality_score(state)
    mx = int(_max_quality())
    qr = _quality_ratio(state)
    temp = chat_temperature(state)
    lines = [
        f"  balance    : ${state['money']:.2f}",
        f"  quality    : {qs}/{mx}  ({qr * 100:.1f}%)",
        f"  temperature: {temp:.2f}",
        f"  runs       : {state['runs']}",
        f"  cloud      : {'ACTIVE' if state['cloud_active'] else ('UNLOCKED/PAUSED' if state['cloud_unlocked'] else 'LOCKED')}",
    ]
    return "\n".join(lines)


def upgrade_status(state: dict) -> dict[str, dict]:
    """Returns per-upgrade display info."""
    result = {}
    for k, upg in UPGRADES.items():
        cur = state["upgrades"][k]
        max_tier = len(upg["levels"]) - 1
        next_cost = upg["levels"][cur + 1]["cost"] if cur < max_tier else None
        result[k] = {
            "label": upg["label"],
            "icon": upg["icon"],
            "description": upg["description"],
            "current_tier": cur,
            "current_name": upg["levels"][cur]["name"],
            "max_tier": max_tier,
            "next_cost": next_cost,
            "maxed": cur >= max_tier,
        }
    return result


# ── Internal ─────────────────────────────────


def _copy(state: dict) -> dict:
    import copy

    return copy.deepcopy(state)