workflow_twin/environment.py

from __future__ import annotations

import hashlib
import random
import uuid
from typing import Any

from workflow_twin.core.config import LEVEL_CONFIGS, LevelConfig
from workflow_twin.core.dynamics import (
    advance_time,
    apply_action,
    check_sla,
    compute_sla_penalty,
    is_terminal,
    pick_next_ticket,
    severity_priority,
)
from workflow_twin.core.entities import Agent, SystemState, Ticket
from workflow_twin.levels import apply_level1, apply_level2, apply_level3, apply_level4, apply_level5
from workflow_twin.models import Action, Observation, RewardSignal


class WorkflowTwinEnv:
    def __init__(self, level: int = 1, seed: int = 42, embedding_dim: int = 16) -> None:
        if level not in LEVEL_CONFIGS:
            raise ValueError(f"Unsupported level {level}. Expected one of {sorted(LEVEL_CONFIGS)}")

        self.level = level
        self.config: LevelConfig = LEVEL_CONFIGS[level]
        self.seed = seed
        self.embedding_dim = embedding_dim
        self._rng = random.Random(seed)
        self._state: SystemState | None = None
        self._arrival_seed = 0
        self.task_definition = {
            "objective": "minimize SLA violations",
            "max_steps": self.config.max_steps,
            "difficulty": f"level_{level}",
        }

    def _build_embedding(self, summary: str, severity: str) -> list[float]:
        key = f"{summary}|{severity}|{self.embedding_dim}".encode("utf-8")
        digest = hashlib.sha256(key).digest()
        raw = [((digest[i % len(digest)] / 127.5) - 1.0) for i in range(self.embedding_dim)]

        norm = sum(v * v for v in raw) ** 0.5
        if norm > 0:
            raw = [v / norm for v in raw]
        return raw

    def _sample_ticket(self, idx: int) -> Ticket:
        severity = self._rng.choices(["low", "medium", "high"], weights=[0.45, 0.35, 0.2], k=1)[0]
        summary = f"Workflow ticket {idx} ({severity})"
        is_approval = self.level >= 3 and self._rng.random() < 0.35
        approval_chain = ["manager", "legal"] if is_approval else []

        dependencies: list[str] = []
        if self.level >= 3 and idx > 0 and self._rng.random() < 0.25:
            dependencies = [f"T-{max(0, idx - 1)}"]

        created_at = float(idx)
        deadline = created_at + (8.0 if severity == "high" else 12.0 if severity == "medium" else 18.0)

        return Ticket(
            id=f"T-{idx}",
            summary=summary,
            severity=severity,
            embedding=self._build_embedding(summary, severity),
            created_at=created_at,
            deadline=deadline if self.config.enable_sla else None,
            priority=severity_priority(severity),
            max_attempts=5 if severity == "high" else 4,
            ticket_type="approval" if is_approval else "support",
            approval_chain=approval_chain,
            dependencies=dependencies,
        )

    def _build_agents(self) -> list[Agent]:
        agents: list[Agent] = []
        for i in range(self.config.agent_count):
            if i == 0:
                agents.append(Agent(id=f"agent-{i}", role="support_agent", agent_type="support"))
            elif i == 1 and self.level >= 2:
                agents.append(
                    Agent(
                        id=f"agent-{i}",
                        role="manager_agent",
                        agent_type="manager",
                        approval_authority=["manager"],
                        hourly_cost=130.0,
                    )
                )
            elif i == 2 and self.level >= 3:
                agents.append(
                    Agent(
                        id=f"agent-{i}",
                        role="legal_agent",
                        agent_type="legal",
                        approval_authority=["legal"],
                        hourly_cost=180.0,
                    )
                )
            else:
                agents.append(Agent(id=f"agent-{i}", role="support_agent", agent_type="support", hourly_cost=110.0))
        return agents

    def reset(self) -> Observation:
        initial_tickets = [self._sample_ticket(i) for i in range(self.config.initial_tickets)]
        state = SystemState(
            episode_id=str(uuid.uuid4()),
            waiting_queue=initial_tickets,
            agents=self._build_agents(),
            memory_budget=self.config.memory_budget,
        )

        state.current_ticket = pick_next_ticket(state)
        self._state = state
        self._arrival_seed = 0
        return self._to_observation()

    def _to_observation(self) -> Observation:
        if self._state is None:
            raise RuntimeError("Environment not initialized. Call reset().")

        ticket = self._state.current_ticket
        waiting_preview = [
            {"id": t.id, "severity": t.severity, "priority": t.priority, "status": t.status}
            for t in self._state.waiting_queue[:5]
        ]

        agents_available = sum(1 for agent in self._state.agents if self._state.current_time >= agent.failed_until)

        return Observation(
            ticket_id=ticket.id if ticket else None,
            ticket_status=ticket.status if ticket else None,
            attempts_used=ticket.attempts_used if ticket else 0,
            attempts_remaining=(ticket.max_attempts - ticket.attempts_used) if ticket else 0,
            severity=ticket.severity if ticket else "low",
            summary=ticket.summary if ticket else "",
            embedding=ticket.embedding if ticket else None,
            queue_length=len(self._state.waiting_queue),
            waiting_tickets=waiting_preview,
            agents_available=agents_available,
            current_time=self._state.current_time,
            level=self.level,
            memory_used=self._state.memory_used,
            memory_budget=self._state.memory_budget,
        )

    def _process_action(self, action: Action) -> dict[str, Any]:
        assert self._state is not None
        info: dict[str, Any] = {}

        if action.action_type == "compress_memory":
            if hasattr(self._state, "memory_used"):
                self._state.memory_used = int(self._state.memory_used * 0.7)
            info["manual_compression"] = True
            return info

        ticket = self._state.current_ticket

        if ticket is None:
            self._state.current_ticket = pick_next_ticket(self._state)
            info["idle"] = True
            return info

        if action.action_type == "assign":
            if action.agent_id:
                ticket.assigned_agent = action.agent_id
                info["assigned_agent"] = action.agent_id
            ticket = apply_action(ticket, "assign")
        elif action.action_type == "reprioritize":
            if action.priority is not None:
                ticket.priority = max(0, min(2, action.priority))
                info["reprioritized"] = ticket.priority
            ticket = apply_action(ticket, "reprioritize")
        elif action.action_type == "skip":
            self._state.waiting_queue.append(ticket)
            self._state.current_ticket = pick_next_ticket(self._state)
            info["skipped"] = ticket.id
            return info
        else:
            ticket = apply_action(ticket, action.action_type)

        if ticket.ticket_type == "approval" and ticket.status == "in_progress" and ticket.current_approval_step < len(ticket.approval_chain):
            required = ticket.approval_chain[ticket.current_approval_step]
            if action.action_type == "escalate":
                ticket.approvals_received.append(required)
                ticket.current_approval_step += 1
                info["approval_progress"] = ticket.current_approval_step

        if is_terminal(ticket):
            self._state.completed_tickets.append(ticket)
            self._state.current_ticket = pick_next_ticket(self._state)
            info["completed"] = ticket.id
        else:
            self._state.current_ticket = ticket

        return info

    def _compute_reward(self, previous_state: SystemState, action: Action) -> RewardSignal:
        assert self._state is not None

        reward = 0.0
        reasons: list[str] = []

        resolved_delta = len(self._state.completed_tickets) - len(previous_state.completed_tickets)
        if resolved_delta > 0:
            reward += 10.0 * resolved_delta
            reasons.append("resolved")

        if action.action_type in {"triage", "respond", "assign", "reprioritize", "skip"}:
            reward -= 0.5
            reasons.append("step_cost")

        if action.action_type == "escalate":
            reward += 0.3
            reasons.append("escalation")

        if self.config.enable_sla:
            penalties = 0.0
            for ticket in self._state.completed_tickets:
                if not check_sla(ticket, self._state.current_time):
                    penalties += compute_sla_penalty(ticket, self._state.current_time)
            if penalties > 0:
                reward -= penalties / 100.0
                reasons.append("sla")

            time_delta = self._state.current_time - previous_state.current_time
            cost_delta = self._state.total_cost - previous_state.total_cost
            if time_delta > 0 and cost_delta > 0:
                reward -= cost_delta / 250.0
                reasons.append("cost")

        queue_pressure = 0.1 * len(self._state.waiting_queue)
        reward -= queue_pressure
        reasons.append("queue")

        if self.config.enable_memory and self._state.memory_used > self._state.memory_budget:
            overflow = self._state.memory_used - self._state.memory_budget
            reward -= overflow / 1000.0
            reasons.append("memory")

        value = max(-1.0, min(1.0, reward / 5.0))
        return RewardSignal(value=value, reason="+".join(reasons))

    def _apply_level_hooks(self) -> None:
        assert self._state is not None
        apply_level1(self._state)
        if self.level >= 2:
            apply_level2(self._state)
        if self.level >= 3:
            apply_level3(self._state)
        if self.level >= 4:
            apply_level4(self._state, self.config.arrival_rate, self._arrival_seed, self._rng)
            self._arrival_seed += 1
        if self.level >= 5:
            apply_level5(self._state)

    def step(self, action: Action | dict[str, Any]) -> tuple[Observation, float, bool, dict[str, Any]]:
        if self._state is None:
            self.reset()

        if isinstance(action, dict):
            action = Action(**action)

        assert self._state is not None
        previous_state = self._state.model_copy(deep=True)

        advance_time(self._state, 1.0)
        action_info = self._process_action(action)
        self._apply_level_hooks()

        self._state.step_count += 1
        self._state.done = (
            self._state.current_ticket is None
            and len(self._state.waiting_queue) == 0
            and len(self._state.active_tickets) == 0
        ) or self._state.step_count >= self.config.max_steps

        reward = self._compute_reward(previous_state, action)
        if self._state.done:
            resolved = len(self._state.completed_tickets)
            total = max(resolved + len(self._state.waiting_queue) + (1 if self._state.current_ticket else 0), 1)
            self._state.score = min(1.0, resolved / total)

        observation = self._to_observation()
        info = {
            "episode_id": self._state.episode_id,
            "step_count": self._state.step_count,
            "level": self.level,
            "task": self.task_definition,
            "queue": {
                "waiting": len(self._state.waiting_queue),
                "completed": len(self._state.completed_tickets),
                "active": len(self._state.active_tickets),
            },
            "sla_violations": self._state.sla_violations,
            "total_cost": round(self._state.total_cost, 2),
            "memory": {
                "used": self._state.memory_used,
                "budget": self._state.memory_budget,
            },
            "action": action_info,
            "score": self._state.score,
        }

        return observation, reward.value, self._state.done, info

    def state(self) -> dict[str, Any]:
        if self._state is None:
            return {"initialized": False, "level": self.level}
        return {
            "episode_id": self._state.episode_id,
            "step": self._state.step_count,
            "level": self.level,
            "task": self.task_definition,
            "queue": {
                "current_ticket": self._state.current_ticket.id if self._state.current_ticket else None,
                "waiting": len(self._state.waiting_queue),
                "active": len(self._state.active_tickets),
                "completed": len(self._state.completed_tickets),
            },
            "metrics": {
                "time": self._state.current_time,
                "score": self._state.score,
                "total_cost": self._state.total_cost,
                "sla_violations": self._state.sla_violations,
                "memory_used": self._state.memory_used,
                "memory_budget": self._state.memory_budget,
            },
            "critical_events": self._state.critical_events[-20:],
            "agents": [agent.model_dump() for agent in self._state.agents],
        }