from abc import ABC, abstractmethod
from typing import Any, Tuple, Dict, Optional
import briefcase_ai as bf
import uuid

class Environment(ABC):
    """
    Gymnasium-like environment for agentic tasks with automated Briefcase tracking.
    """
    def __init__(self, study_id: str, name: Optional[str] = None):
        self.study_id = study_id
        self.name = name or self.__class__.__name__
        self.current_step = 0
        self.trace_id = str(uuid.uuid4())

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Flush any pending traces or clean up resources
        bf.shutdown()

    @abstractmethod
    def reset(self) -> Any:
        """Resets the environment and returns initial observation."""
        self.current_step = 0
        self.trace_id = str(uuid.uuid4())
        
        # Capture reset event
        dec = bf.Decision(f"{self.name}_reset")
        dec.add_tag("trace_id", self.trace_id)
        bf.save(dec)
        return self._reset()

    @abstractmethod
    def _reset(self) -> Any:
        pass

    def step(self, action: Any) -> Tuple[Any, float, bool, Dict]:
        """
        Executes an action, captures the decision, and returns results.
        """
        self.current_step += 1
        
        # 1. Execute task-specific step logic
        observation, reward, terminated, info = self._step(action)
        
        # 2. Automatically capture the step
        dec = bf.Decision(f"{self.name}_step")
        dec.add_tag("trace_id", self.trace_id)
        dec.add_tag("step", str(self.current_step))
        dec.add_input(bf.Input("action", action, "any"))
        
        # Capture Score
        score = bf.Scorecard()
        score.add_score("step_reward", float(reward), weight=1.0)
        dec.with_scorecard(score)
        
        # Capture output
        output = bf.Output("observation", observation, "any")
        output.with_confidence(info.get("confidence", 1.0))
        dec.add_output(output)
        
        bf.save(dec)

        return observation, reward, terminated, info

    @abstractmethod
    def _step(self, action: Any) -> Tuple[Any, float, bool, Dict]:
        pass
