server/environment.py

"""OrgOS environment — the single stateful RL environment object."""

import uuid
from typing import Dict, Optional

from models import OrgOSAction, OrgOSObservation, OrgOSState, RewardBreakdown
from server.apps.jira import JiraApp
from server.apps.zendesk import ZendeskApp
from server.apps.salesforce import SalesforceApp
from server.apps.workday import WorkdayApp
from server.business_rules import BusinessRuleEngine
from server.data_generator import generate_episode_data
from server.schema_drift import SchemaDriftEngine
from server.workflow_engine import WorkflowEngine


class OrgOSEnvironment:
    MAX_STEPS = {"A": 15, "B": 20, "C": 18}
    WORKFLOWS  = ["A", "B", "C"]

    def __init__(self):
        self._drift    = SchemaDriftEngine(seed=42)
        self._rules    = BusinessRuleEngine()
        self._workflow = WorkflowEngine()
        self._apps: Dict[str, object] = {
            "jira":       JiraApp(self._drift),
            "zendesk":    ZendeskApp(self._drift),
            "salesforce": SalesforceApp(self._drift),
            "workday":    WorkdayApp(self._drift),
        }
        self._episode_num          = 0
        self._episode_id           = ""
        self._workflow_id          = "A"
        self._step_count           = 0
        self._last_score           = 0.001
        self._policy_drift_applied = False

        # Reward component trackers
        self._wf_score     = 0.0   # workflow completion
        self._rule_score   = 0.0   # compliance — earned +0.10 per successful action
        self._schema_score = 0.0   # schema adaptation successes
        self._efficiency   = 0.0   # efficiency — earned +0.10 per successful action
        self._policy_score = 0.0   # policy drift handling bonus

    # ------------------------------------------------------------------
    # OpenEnv core API
    # ------------------------------------------------------------------

    def reset(self, workflow_id: Optional[str] = None) -> OrgOSObservation:
        self._episode_num += 1
        self._episode_id   = str(uuid.uuid4())
        self._workflow_id  = workflow_id or self.WORKFLOWS[(self._episode_num - 1) % 3]
        self._step_count   = 0
        self._last_score   = 0.001
        self._rule_score   = 0.0
        self._wf_score     = 0.0
        self._schema_score = 0.0
        self._efficiency   = 0.0
        self._policy_score = 0.0
        self._policy_drift_applied = False

        # Sample schema versions for this episode
        self._drift.sample_for_episode(self._episode_num)

        # Possibly activate policy drift (every 3rd episode)
        self._rules = BusinessRuleEngine()
        if self._episode_num % 3 == 0:
            self._rules.apply_policy_drift("sla_tighten")
            self._policy_drift_applied = True

        # Load fresh synthetic data into each app
        records = generate_episode_data(self._workflow_id, seed=42 + self._episode_num)
        for app_name, app in self._apps.items():
            app.initialize(records[app_name])

        # Start workflow tracking
        self._workflow.start(self._workflow_id)

        return self._build_obs(
            reward=0.001,
            done=False,
            message="Episode started. Study the workflow goal and schema hints before acting.",
        )

    def step(self, action: OrgOSAction) -> OrgOSObservation:
        self._step_count += 1
        old_score    = self._last_score
        extra_penalty = 0.0

        # 1. Validate app exists
        if action.app not in self._apps:
            return self._build_obs(
                reward=-0.05,
                done=False,
                message=f"Unknown app '{action.app}'. Valid apps: {list(self._apps)}",
            )

        # 2. Business rule check (RBAC, approvals)
        agent_role = self._workflow.get_role()
        ctx        = {"agent_role": agent_role, "manager_approved": False}
        allowed, reason, rule_penalty = self._rules.check_action(action, ctx)
        if not allowed:
            self._rule_score = max(0.0, self._rule_score - 0.08)
            extra_penalty    = rule_penalty
            return self._build_obs(
                reward=extra_penalty,
                done=False,
                message=f"Rule violation: {reason}",
            )

        # 3. Execute on app
        result = self._apps[action.app].execute(action.operation, action.args)

        # 4. Check schema drift FIRST — apps return success:False when schema_error is set
        if result.get("schema_error"):
            self._efficiency -= 0.02
            return self._build_obs(
                reward=-0.20,
                done=False,
                message=(
                    f"Stale schema: field '{result['schema_error']}' is no longer valid. "
                    "Check schema_hints for the current field name. "
                    f"Hint: {result.get('message', '')}"
                ),
            )

        if not result.get("success"):
            self._efficiency -= 0.02   # penalize failed/no-op actions
            return self._build_obs(
                reward=-0.01,
                done=False,
                message=result.get("message", "Operation failed"),
            )

        # Schema adaptation bonus (agent used correct drifted field name)
        if result.get("schema_adapted"):
            self._schema_score = min(1.0, self._schema_score + 0.10)
            self._policy_score = min(1.0, self._policy_score + 0.05)

        # Earn compliance for every successful compliant action,
        # normalised so completing all workflow steps earns exactly 1.0.
        total_steps = max(1, len(self._workflow._steps))
        earn_rate   = 1.0 / total_steps
        self._rule_score = min(1.0, self._rule_score + earn_rate)

        # 5. Re-evaluate workflow completion
        old_wf_score     = self._wf_score
        self._wf_score   = self._workflow.evaluate(self._apps)
        wf_advanced      = self._wf_score > old_wf_score

        # Earn efficiency ONLY when a new workflow step was just completed.
        # This penalises padding: closing random tickets, repeating ops, etc.
        if wf_advanced:
            self._efficiency = min(1.0, self._efficiency + earn_rate)

        # 6. SLA check (only if a ticket was touched)
        sla_ok, sla_pen = self._rules.check_sla(
            result.get("ticket", {}),
            self._step_count * 10,    # 10 min per step — P1 breaches at step 24 (step 12 under policy drift)
        )
        if not sla_ok:
            extra_penalty   += sla_pen
            self._rule_score = max(0.0, self._rule_score - 0.05)

        # 7. Compute composite score
        new_score = self._compute_score()
        delta     = new_score - old_score + extra_penalty
        self._last_score = max(0.001, min(0.999, new_score))

        # 8. Terminal condition
        done = (
            self._wf_score >= 0.95
            or self._step_count >= self.MAX_STEPS[self._workflow_id]
        )
        if done and self._wf_score >= 0.95:
            delta += 0.20   # terminal completion bonus

        return self._build_obs(
            reward=delta,
            done=done,
            message=result.get("message", "OK"),
        )

    # ------------------------------------------------------------------
    # State endpoint
    # ------------------------------------------------------------------

    def state(self) -> OrgOSState:
        return OrgOSState(
            episode_id           = self._episode_id,
            workflow_id          = self._workflow_id,
            schema_versions      = self._drift._versions,
            step_count           = self._step_count,
            max_steps            = self.MAX_STEPS.get(self._workflow_id, 15),
            rule_violation_count = len(self._rules._violation_log),
            workflow_completion  = self._wf_score,
            rule_compliance_rate = self._rule_score,
            policy_drift_active  = self._policy_drift_applied,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _compute_score(self) -> float:
        raw = (
            0.30 * self._wf_score   +
            0.25 * self._rule_score +
            0.20 * self._schema_score +
            0.15 * self._efficiency +
            0.10 * self._policy_score
        )
        return max(0.001, min(0.999, raw))

    def _build_obs(self, reward: float, done: bool, message: str) -> OrgOSObservation:
        """Construct a fully-populated observation from current environment state."""
        # Per-app state previews
        app_states = {
            name: app.get_state_view(max_rows=3)
            for name, app in self._apps.items()
        }

        flat_hints = self._drift.get_hints()

        # # Schema hints (partial — agent must probe to discover full mapping)
        # schema_hints = self._drift.get_hints()
        # # Flatten to dot-notation: {"jira.priority": "severity", ...}
        # flat_hints: Dict[str, str] = {}
        # for app_name, field_map in schema_hints.items():
        #     for canonical, drifted in field_map.items():
        #         if canonical != drifted:
        #             flat_hints[f"{app_name}.{canonical}"] = drifted

        # Workflow progress
        completed_steps = self._workflow.get_completed()
        pending_steps   = self._workflow.get_pending()
        workflow_goal   = self._workflow.get_goal()

        # Reward breakdown snapshot
        breakdown = RewardBreakdown(
            workflow_completion   = self._wf_score,
            rule_compliance       = self._rule_score,
            schema_adaptation     = self._schema_score,
            efficiency            = self._efficiency,
            policy_drift_handling = self._policy_score,
        )

        return OrgOSObservation(
            done              = done,
            reward            = round(float(reward), 6),
            current_score     = round(float(self._last_score),4),
            workflow_id       = self._workflow_id,
            step_count        = self._step_count,
            app_states        = app_states,
            workflow_goal     = workflow_goal,
            completed_steps   = completed_steps,
            pending_steps     = pending_steps,
            schema_hints      = flat_hints,
            active_rules      = self._rules.get_active_rules_summary(),
            rule_violations   = self._rules.get_violations_this_step(),
            reward_breakdown  = breakdown,
            message           = message,
        )