"""
PII-Scrub-Assistant-v1 -- Environment
======================================
3 tasks, programmatic graders (0.0-1.0), partial-progress reward.
"""
from __future__ import annotations
import random
from typing import Optional
from models import ScrubObservation, ScrubState


TASKS = {
    "task_1": {
        "difficulty": "easy",
        "instruction": "Redact the phone number in the following text. Replace any PII with [REDACTED].",
        "original_text": "Hi, please call me back at 555-867-5309 regarding my order.",
        "pii_items": {"phone": ["555-867-5309"]},
        "order_ids": [],
    },
    "task_2": {
        "difficulty": "medium",
        "instruction": "Redact all PII (names and email addresses) in the following paragraph. Replace each piece of PII with [REDACTED].",
        "original_text": "Dear Support,\n\nMy name is Johnathan Doe and I recently purchased your premium plan. Could you send the invoice to johndoe92@example.com? I need it for my expense report.\n\nThanks,\nJohnathan Doe",
        "pii_items": {"name": ["Johnathan Doe"], "email": ["johndoe92@example.com"]},
        "order_ids": [],
    },
    "task_3": {
        "difficulty": "hard",
        "instruction": "Redact all PII (names and email addresses) but DO NOT redact Order IDs (e.g. #12345) or System IDs. Replace each piece of PII with [REDACTED].",
        "original_text": "Hello,\n\nI'm Alice Johnson and I need help with order #98432. My email is alice.johnson@corpmail.net and my secondary contact is alice_j@inbox.org. Please also check order #10271.\n\nRegards,\nAlice Johnson",
        "pii_items": {"name": ["Alice Johnson"], "email": ["alice.johnson@corpmail.net", "alice_j@inbox.org"]},
        "order_ids": ["#98432", "#10271"],
    },
}


def _all_pii_removed(redacted, pii_items):
    return {t: all(v not in redacted for v in vs) for t, vs in pii_items.items()}

def _order_ids_preserved(redacted, order_ids):
    return all(oid in redacted for oid in order_ids)

def _count_over_redactions(original, redacted, pii_items):
    expected = sum(original.count(v) for vs in pii_items.values() for v in vs)
    return max(0, redacted.count("[REDACTED]") - expected)

def _structural_damage(original, redacted, pii_items):
    tag = "[REDACTED]"
    delta = sum(original.count(v) * (len(tag) - len(v)) for vs in pii_items.values() for v in vs)
    exp_len = len(original) + delta
    return abs(len(redacted) - exp_len) > max(20, int(exp_len * 0.15))


def compute_reward(original, redacted, pii_items, order_ids):
    parts, reward, n_ok = [], 0.0, 0
    for t, ok in _all_pii_removed(redacted, pii_items).items():
        if ok:
            reward += 0.5; n_ok += 1; parts.append(f"[OK] {t} redacted")
        else:
            parts.append(f"[FAIL] {t} NOT fully redacted")
    over = _count_over_redactions(original, redacted, pii_items)
    if over > 0:
        reward -= over * 0.2; parts.append(f"[FAIL] Over-redacted ({over}), -{over*0.2:.1f}")
    if order_ids:
        if not _order_ids_preserved(redacted, order_ids):
            reward -= 0.3; parts.append("[FAIL] Order/System ID(s) redacted (-0.3)")
        else:
            parts.append("[OK] Order/System IDs preserved")
    if _structural_damage(original, redacted, pii_items):
        reward -= 1.0; parts.append("[FAIL] Structural damage (-1.0)")
    total = len(pii_items)
    base = n_ok / total if total else 1.0
    ded = over * 0.05
    if order_ids and not _order_ids_preserved(redacted, order_ids): ded += 0.2
    if _structural_damage(original, redacted, pii_items): ded += 0.3
    score = max(0.0, min(1.0, base - ded))
    return reward, score, " | ".join(parts)


class ScrubEnvironment:
    def __init__(self):
        self._task_id = self._task = self._player_id = self._session_id = None
        self._step_count = 0; self._done = True
        self._last_score = self._last_reward = None

    def _validate(self, pid, sid):
        if self._player_id and pid != self._player_id:
            raise PermissionError(f"Player mismatch: expected {self._player_id!r}")
        if self._session_id and sid != self._session_id:
            raise PermissionError(f"Session mismatch: expected {self._session_id!r}")

    def reset(self, player_id, session_id, task_id=None):
        if task_id is None: task_id = random.choice(list(TASKS))
        if task_id not in TASKS: raise ValueError(f"Unknown task: {task_id}")
        self._player_id, self._session_id = player_id, session_id
        self._task_id, self._task = task_id, TASKS[task_id]
        self._step_count, self._done = 0, False
        self._last_score = self._last_reward = None
        return ScrubObservation(task_id=task_id, original_text=self._task["original_text"], instruction=self._task["instruction"])

    def step(self, player_id, session_id, redacted_text):
        if self._done or not self._task: raise RuntimeError("Call /reset first.")
        self._validate(player_id, session_id)
        self._step_count += 1
        reward, score, feedback = compute_reward(self._task["original_text"], redacted_text, self._task["pii_items"], self._task["order_ids"])
        self._done, self._last_score, self._last_reward = True, score, reward
        return ScrubObservation(task_id=self._task_id, original_text=self._task["original_text"], instruction=self._task["instruction"], score=score, reward=reward, done=True, feedback=feedback)

    def state(self):
        return ScrubState(current_task_id=self._task_id, player_id=self._player_id, session_id=self._session_id, step_count=self._step_count, done=self._done, last_score=self._last_score, last_reward=self._last_reward, original_text=self._task["original_text"] if self._task else None)