"""
Inference script for OpenEnv Store environment.

Uses an LLM (via OpenAI-compatible client) to play the kirana-store simulation.
Emits structured [START], [STEP], [END] logs for evaluation scoring.

Required env vars:
    API_BASE_URL  – LLM API endpoint
    MODEL_NAME    – model identifier
    HF_TOKEN      – API key
"""

from __future__ import annotations

import asyncio
import json
import os
import sys
import time
from typing import Any

from openai import AsyncOpenAI

from client import StoreEnv
from models import (
    OrderRequest,
    PriceUpdate,
    StoreAction,
    StoreObservation,
    dump_model,
)
from dotenv import load_dotenv
load_dotenv()

API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3.5-397B-A17B")
HF_TOKEN = os.getenv("HF_TOKEN")

BENCHMARK = "store"

TASKS: list[dict[str, Any]] = [
    {
        "name": "store_short_run",
        "description": "Run the store profitably over 5 decision windows (15 days).",
        "max_steps": 5,
        "success_threshold": 0.5,
    },
    {
        "name": "store_medium_run",
        "description": "Run the store profitably over 10 decision windows (30 days).",
        "max_steps": 10,
        "success_threshold": 0.55,
    },
    {
        "name": "store_long_run",
        "description": (
            "Run the store profitably over 15 decision windows (45 days). "
            "Requires sustained capital growth and careful long-term planning."
        ),
        "max_steps": 15,
        "success_threshold": 0.6,
    },
]

ENV_URL = os.environ.get("ENV_URL", "http://localhost:8000")


def log_start(*, task: str, env: str, model: str) -> None:
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(
    *,
    step: int,
    action: dict[str, Any] | str,
    reward: float,
    done: bool,
    error: str | None,
) -> None:
    done_str = "true" if done else "false"
    error_str = error if error else "null"
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} "
        f"done={done_str} error={error_str}",
        flush=True,
    )


def log_end(
    *,
    success: bool,
    steps: int,
    score: float,
    rewards: list[float],
) -> None:
    success_str = "true" if success else "false"
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={success_str} steps={steps} score={score:.2f} "
        f"rewards={rewards_str}",
        flush=True,
    )


def _format_observation(obs: StoreObservation) -> str:
    """Build a concise text summary of the observation for the LLM prompt."""
    store = obs.store
    lines = [
        "=== Store Snapshot ===",
        f"Day: {store.day}/{store.episode_horizon_days}",
        f"Bank Balance: ₹{store.bank_balance_inr:,.2f}",
        f"Customer Satisfaction: {store.customer_satisfaction}",
        f"Order Processing Capacity: {store.order_processing_capacity}",
        f"Daily Expense: ₹{store.current_daily_expense_inr:,.2f}",
        f"Cumulative Profit: ₹{store.cumulative_profit_inr:,.2f}",
        f"Active Recurring Cost: ₹{store.active_recurring_cost_inr:,.2f}",
        f"Active Action IDs: {store.active_action_ids}",
        "",
        "=== Current Stock (top items by value) ===",
    ]
    sorted_stock = sorted(obs.stock, key=lambda s: s.stock_value_inr, reverse=True)
    for item in sorted_stock[:15]:
        margin = (
            ((item.selling_price_inr - item.buying_price_inr) / item.selling_price_inr) * 100
            if item.selling_price_inr > 0
            else 0
        )
        lines.append(
            f"  {item.product_name}: qty={item.quantity_on_hand:.1f}, "
            f"buy=₹{item.buying_price_inr}, sell=₹{item.selling_price_inr}, "
            f"demand={item.demand_score}, margin={margin:.1f}%, "
            f"value=₹{item.stock_value_inr:.0f}"
        )

    low_stock = [s for s in obs.stock if 0 < s.quantity_on_hand < 15]
    if low_stock:
        lines.append("")
        lines.append("=== Low Stock Alert ===")
        for item in low_stock:
            lines.append(f"  {item.product_name}: qty={item.quantity_on_hand:.1f}")

    if obs.pending_orders:
        lines.append("")
        lines.append("=== Pending Orders ===")
        for po in obs.pending_orders:
            lines.append(
                f"  {po.order_id}: {po.product_name} x{po.quantity} from {po.seller_name}, "
                f"arrives in {po.days_remaining}d"
            )

    lines.append("")
    lines.append("=== Marketing Actions (inactive only) ===")
    for ma in obs.marketing_actions:
        if not ma.is_active:
            lines.append(
                f"  ID {ma.action_id}: {ma.description} | "
                f"cost=₹{ma.cost_inr}, recurring=₹{ma.recurring_cost_inr}/day, "
                f"satisfaction_delta={ma.customer_satisfaction_delta:+d}, "
                f"capacity_delta={ma.order_processing_capacity_delta:+d}"
            )

    if obs.metadata:
        rb = obs.metadata.get("reward_breakdown")
        if rb:
            lines.append("")
            lines.append("=== Last Reward Breakdown ===")
            lines.append(f"  Starting Capital: ₹{rb.get('starting_capital_inr', 'N/A')}")
            lines.append(f"  Ending Capital: ₹{rb.get('ending_capital_inr', 'N/A')}")
            lines.append(f"  Difference: ₹{rb.get('difference_inr', 'N/A')}")
            lines.append(f"  Recurring Costs: ₹{rb.get('recurring_costs_paid_inr', 'N/A')}")

    return "\n".join(lines)


SYSTEM_PROMPT = """\
You are a store management AI agent controlling a neighborhood kirana (grocery) store in India.

## Simulation Mechanics

Each step you submit ONE bundled action for the next 3-day decision window:
1. orders: list of supplier orders [{"product": "...", "seller": "Seller X", "quantity": N}]
2. prices: list of price updates [{"product": "...", "price": N}] (in-stock products only)
3. action_ids: list of marketing/operational action IDs to activate

## Ground-Truth Formulas

DAILY ORDERS PROCESSED = order_processing_capacity × (customer_satisfaction / 100)

For each in-stock product with quantity > 0:
  margin_pct = ((selling_price - buying_price) / selling_price) × 100
  product_selling_score = max(0, demand_score - margin_pct × 3)

TOTAL_SELLING_SCORE = sum of all product_selling_scores
Each product's daily sales = (product_selling_score / TOTAL_SELLING_SCORE) × DAILY_ORDERS_PROCESSED
  (capped at available quantity)

REWARD = 0.5 + (ending_capital - starting_capital - recurring_costs_6days) / starting_capital
  where capital = bank_balance + total_inventory_value

CUSTOMER SATISFACTION UPDATE after price/stock changes:
  new_satisfaction = (1 + (score_after - score_before) / score_after) × base_satisfaction
  Capped at [1, 100].

## Seller Profiles
- Seller A: Cheapest on staples but OVERCOMMITS — deliveries often 1-2 days late
- Seller B: Reliable baseline, matches quoted ETA consistently
- Seller C: Quotes conservatively, often arrives EARLY — best for urgent fresh restocks
- Seller D: Competitive on snacks/beverages but execution is uneven

## Key Constraints
- Episode: 30 days, each step commits 3 days, reward horizon = 6 days
- Bank balance < 0 → episode ends immediately
- Marketing actions are DURABLE: one-time cost + permanent recurring daily cost
- Starting cash is tight (₹100k) relative to inventory (~₹240k)

## Response Format
You MUST respond with ONLY valid JSON:
{
  "reasoning": "brief explanation of your decision",
  "orders": [{"product": "...", "seller": "Seller X", "quantity": N}, ...],
  "prices": [{"product": "...", "price": N}, ...],
  "action_ids": [N, ...]
}
No extra text outside JSON. Product names must exactly match the catalog.
"""


def _build_task_system_prompt(task: dict[str, Any]) -> str:
    max_steps = task["max_steps"]
    threshold = task["success_threshold"]
    goal = (
        f"## Task: {task['name']}\n\n"
        f"You have {max_steps} decision windows. "
        f"Target average reward per step ≥ {threshold:.1f}.\n\n"
        "STRATEGY PRIORITIES:\n"
        "1. Activate zero-recurring marketing early: IDs 1, 10, 6 (+75 capacity, no daily cost)\n"
        "2. Restock high-demand cheap items: Biscuits (demand 90), Bottled Water (demand 89), "
        "Eggs (demand 84), Bread (demand 86) — use Seller B/C for reliability\n"
        "3. Use Seller A for bulk staples (Rice, Oil, Spices) only when stock cover is safe\n"
        "4. Protect anchor item prices: do NOT raise Milk/Bread/Eggs/Sugar/Bottled Water\n"
        "5. NEVER reorder Kitchen Appliances (dead stock, demand 28, margin ~35%)\n"
        "6. NEVER activate ID 15 (social ads: ₹1,200/day for +0 capacity)\n"
        "7. Keep bank_balance > ₹15k safety buffer at all times\n"
        "8. After foundations, consider WhatsApp ordering (ID 4) and delivery partner (ID 7)"
    )
    return SYSTEM_PROMPT + f"\n\n{goal}"


async def get_model_action(
    client: AsyncOpenAI,
    step: int,
    observation_text: str,
    last_reward: float,
    history: list[str],
    task_system_prompt: str,
) -> dict[str, Any]:
    """Query the LLM for the next action."""
    history_text = "\n".join(history[-5:]) if history else "No history yet."

    user_msg = (
        f"Step {step} | Last reward: {last_reward:+.4f}\n\n"
        f"Recent history:\n{history_text}\n\n"
        f"Current observation:\n{observation_text}\n\n"
        "Decide your next action. Respond with ONLY valid JSON."
    )

    response = await client.chat.completions.create(
        model=MODEL_NAME,
        messages=[
            {"role": "system", "content": task_system_prompt},
            {"role": "user", "content": user_msg},
        ],
        temperature=0.3,
        max_tokens=4096,
    )

    content = response.choices[0].message.content or "{}"
    content = content.strip()

    print(f"[DEBUG] Raw LLM response (first 500 chars): {content[:500]}", flush=True)

    think_end = content.find("</think>")
    if think_end != -1:
        content = content[think_end + len("</think>"):].strip()

    if content.startswith("```"):
        content = content.split("\n", 1)[-1]
        if content.endswith("```"):
            content = content[: -len("```")]
        content = content.strip()

    return json.loads(content)


def _parse_action(raw: dict[str, Any]) -> StoreAction:
    """Convert the LLM JSON response into a StoreAction."""
    orders = [
        OrderRequest(
            product=o["product"],
            seller=o["seller"],
            quantity=int(o["quantity"]),
        )
        for o in raw.get("orders", [])
    ]
    prices = [
        PriceUpdate(product=p["product"], price=int(p["price"]))
        for p in raw.get("prices", [])
    ] or None
    action_ids = [int(a) for a in raw.get("action_ids", [])]
    return StoreAction(orders=orders, prices=prices, action_ids=action_ids)


async def run_task(
    llm_client: AsyncOpenAI,
    task: dict[str, Any],
) -> float:
    """Run a single task episode and return the clamped score."""
    task_name = task["name"]
    max_steps = task["max_steps"]
    threshold = task["success_threshold"]
    task_system_prompt = _build_task_system_prompt(task)

    history: list[str] = []
    rewards: list[float] = []
    steps_taken = 0
    score = 0.0
    success = False

    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)

    env = StoreEnv(base_url=ENV_URL)

    try:
        result = await env.reset()
        obs = result.observation
        obs_text = _format_observation(obs)
        last_reward = 0.0

        for step in range(1, max_steps + 1):
            if obs.done:
                break

            raw_action: dict[str, Any] = {}
            action_str = ""
            error = None

            try:
                raw_action = await get_model_action(
                    llm_client,
                    step,
                    obs_text,
                    last_reward,
                    history,
                    task_system_prompt,
                )
                action = _parse_action(raw_action)
                action_str = json.dumps(dump_model(action, exclude_none=True), separators=(",", ":"))
            except Exception as e:
                print(f"[DEBUG] Action parse error at step {step}: {e}", flush=True)
                action = StoreAction()
                action_str = "{}"
                error = str(e)

            try:
                result = await env.step(action)
            except Exception as e:
                print(f"[DEBUG] Step error at step {step}: {e}", flush=True)
                action = StoreAction()
                action_str = "{}"
                error = str(e)
                try:
                    result = await env.step(action)
                except Exception as e2:
                    print(f"[DEBUG] Pass-step also failed: {e2}", flush=True)
                    error = str(e2)
                    break

            obs = result.observation
            reward = result.reward or 0.0
            done = result.done

            rewards.append(reward)
            steps_taken = step
            obs_text = _format_observation(obs)
            last_reward = reward

            log_step(
                step=step,
                action=action_str,
                reward=reward,
                done=done,
                error=error,
            )

            reasoning = raw_action.get("reasoning", "")
            history.append(
                f"Step {step}: reward={reward:+.4f}, done={done}, "
                f"reasoning={reasoning!r}"
            )

            if done:
                break

        score = sum(rewards) / (max_steps * 1.0) if max_steps > 0 else 0.0
        score = min(max(score, 0.0), 1.0)
        success = score >= threshold

    finally:
        try:
            await env.close()
        except Exception as e:
            print(f"[DEBUG] env.close() error (container cleanup): {e}", flush=True)
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)

    return score


async def main() -> None:
    llm_client = AsyncOpenAI(
        base_url=API_BASE_URL,
        api_key=HF_TOKEN,
    )

    print(f"[DEBUG] Starting inference with model={MODEL_NAME}", flush=True)
    print(f"[DEBUG] API_BASE_URL={API_BASE_URL}", flush=True)
    print(f"[DEBUG] ENV_URL={ENV_URL}", flush=True)
    print(f"[DEBUG] Tasks: {[t['name'] for t in TASKS]}", flush=True)

    all_scores: list[float] = []

    for task in TASKS:
        print(f"\n[DEBUG] === Running task: {task['name']} ===", flush=True)
        t0 = time.time()
        task_score = await run_task(llm_client, task)
        elapsed = time.time() - t0
        all_scores.append(task_score)
        print(
            f"[DEBUG] Task {task['name']} completed: score={task_score:.4f}, "
            f"elapsed={elapsed:.1f}s",
            flush=True,
        )

    avg_score = sum(all_scores) / len(all_scores) if all_scores else 0.0
    print(f"\n[DEBUG] All task scores: {all_scores}", flush=True)
    print(f"[DEBUG] Average score: {avg_score:.4f}", flush=True)


if __name__ == "__main__":
    asyncio.run(main())