import json
from uuid import uuid4
from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State
from typing import List, Dict, Optional

try:
    from ..models import GitCommitReviewerAction, GitCommitReviewerObservation
except ImportError:
    from models import GitCommitReviewerAction, GitCommitReviewerObservation

import tasks

# Global cache to handle stateless server environments
# Shared across instances
GLOBAL_SESSION_STORE = {}

class GitCommitReviewerEnvironment(Environment):
    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._reset_count = 0
        self.current_task_commits = []
        self.current_idx = 0
        self.total_rewards = 0.0
        self.last_task = "easy"

    def reset(self, task: str = "easy") -> GitCommitReviewerObservation:
        self.last_task = str(task).lower()
        self._state = State(episode_id=self._state.episode_id, step_count=0) # Keep same EP ID if possible or new one
        self._state.episode_id = str(uuid4()) # Generate new unique session
        self._reset_count += 1
        
        self.current_task_commits = tasks.get_task_data(self.last_task)
        if not self.current_task_commits:
            self.current_task_commits = tasks.get_task_data("easy")

        self.current_idx = 0
        self.total_rewards = 0.0
        
        # Store initial state in global cache
        GLOBAL_SESSION_STORE[self._state.episode_id] = {
            "current_idx": self.current_idx,
            "total_rewards": self.total_rewards,
            "commits": self.current_task_commits,
            "task": self.last_task,
            "step_count": self._state.step_count
        }

        return GitCommitReviewerObservation(
            commit_message=self.current_task_commits[0]["commit_message"],
            commit_idx=0,
            total_commits=len(self.current_task_commits),
            done=False,
            reward=0.0,
            session_id=self._state.episode_id,
            metadata={"session_id": self._state.episode_id}
        )

    def step(self, action: GitCommitReviewerAction) -> GitCommitReviewerObservation:
        # Try to restore from global cache if session matches
        sid = getattr(action, "session_id", None)
        if sid and sid in GLOBAL_SESSION_STORE:
            cached = GLOBAL_SESSION_STORE[sid]
            # Only restore if current instance is uninitialized or mismatched
            if not self.current_task_commits or self._state.episode_id != sid:
                self.current_task_commits = cached["commits"]
                self.current_idx = cached["current_idx"]
                self.total_rewards = cached["total_rewards"]
                self._state.episode_id = sid
                self._state.step_count = cached["step_count"]
                self.last_task = cached["task"]

        self._state.step_count += 1
        
        if not self.current_task_commits or self.current_idx >= len(self.current_task_commits):
            return GitCommitReviewerObservation(
                commit_message="DONE",
                commit_idx=self.current_idx,
                total_commits=len(self.current_task_commits) if self.current_task_commits else 0,
                done=True,
                reward=0.0,
                session_id=self._state.episode_id,
                metadata={"session_id": self._state.episode_id}
            )
            
        gt = self.current_task_commits[self.current_idx]["ground_truth"]
        
        # Grading Logic
        reward = 0.0
        
        # 1. Score matching
        score_diff = abs(action.score - gt["score"])
        if score_diff <= 0.05: reward += 0.45
        elif score_diff <= 0.15: reward += 0.3
        elif score_diff <= 0.3: reward += 0.15
        
        # 2. Type Valid matching
        if action.type_valid == gt["type_valid"]: reward += 0.2
            
        # 3. Summary length proximity
        length_diff = abs(action.summary_length - gt["summary_length"])
        if length_diff == 0: reward += 0.15
        elif length_diff <= 2: reward += 0.05
            
        # 4. Suggestion/Issues quality
        if gt["issues"] and len(action.issues) > 0: reward += 0.1
        elif not gt["issues"] and len(action.issues) == 0: reward += 0.1
            
        if action.suggestion and ":" in action.suggestion and len(action.suggestion.split(":")[0]) <= 20:
            reward += 0.1
            
        reward = max(0.01, min(0.99, reward)) # Clamp for OpenEnv safety
        self.total_rewards += reward
        
        self.current_idx += 1
        done = (self.current_idx >= len(self.current_task_commits))
        
        # Update global cache for this session
        GLOBAL_SESSION_STORE[self._state.episode_id] = {
            "current_idx": self.current_idx,
            "total_rewards": self.total_rewards,
            "commits": self.current_task_commits,
            "task": self.last_task,
            "step_count": self._state.step_count
        }

        return GitCommitReviewerObservation(
            commit_message=self.current_task_commits[self.current_idx]["commit_message"] if not done else "DONE",
            commit_idx=self.current_idx,
            total_commits=len(self.current_task_commits),
            done=done,
            reward=reward,
            session_id=self._state.episode_id,
            metadata={
                "step": self._state.step_count, 
                "total_rewards": self.total_rewards, 
                "session_id": self._state.episode_id,
                "reward": reward
            }
        )

    @property
    def state(self) -> dict:
        # Crucial: Check global cache to ensure we return the latest state
        # even if this is a newly created instance for a stateless request.
        eid = self._state.episode_id
        current_rewards = self.total_rewards
        current_steps = self._state.step_count
        
        if eid in GLOBAL_SESSION_STORE:
            cached = GLOBAL_SESSION_STORE[eid]
            current_rewards = max(current_rewards, cached.get("total_rewards", 0))
            current_steps = max(current_steps, cached.get("step_count", 0))

        return {
            "episode_id": eid,
            "step_count": current_steps,
            "total_rewards": current_rewards,
            "current_idx": self.current_idx,
            "total_commits": len(self.current_task_commits),
            "is_done": self.current_idx >= len(self.current_task_commits) if self.current_task_commits else True,
            "task": self.last_task
        }