import json
from uuid import uuid4
from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

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

import tasks

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

    def reset(self, task: str = "easy") -> GitCommitReviewerObservation:
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._reset_count += 1
        
        self.current_task_commits = tasks.get_task_data(task)
        self.current_idx = 0
        self.total_rewards = 0.0
        
        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
        )

    def step(self, action: GitCommitReviewerAction) -> GitCommitReviewerObservation:
        self._state.step_count += 1
        
        if self.current_idx >= len(self.current_task_commits):
            # Already done, just return terminal state
            return GitCommitReviewerObservation(
                commit_message="DONE",
                commit_idx=self.current_idx,
                total_commits=len(self.current_task_commits),
                done=True,
                reward=0.0
            )
            
        gt = self.current_task_commits[self.current_idx]["ground_truth"]
        
        # Grading Logic!
        reward = 0.0
        
        # 1. Score matching (up to 0.45 weight) - closer is better
        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 (weight 0.2)
        if action.type_valid == gt["type_valid"]:
            reward += 0.2
            
        # 3. Summary length proximity (weight 0.15)
        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 (weight 0.2)
        # We simply check if the correct type is somehow referenced, or if they found > 0 issues when there were issues
        if gt["issues"] and len(action.issues) > 0:
            reward += 0.1
        elif not gt["issues"] and len(action.issues) == 0:
            reward += 0.1
            
        # Give remaining 0.1 if suggestion is reasonably close to standard Conventional Commits
        if ":" in action.suggestion and len(action.suggestion.split(":")[0]) <= 20:
            reward += 0.1
            
        # Ensure reward is between 0 and 1
        reward = max(0.0, min(1.0, reward))
        self.total_rewards += reward
        
        self.current_idx += 1
        done = (self.current_idx >= len(self.current_task_commits))
        
        next_commit = ""
        if not done:
            next_commit = self.current_task_commits[self.current_idx]["commit_message"]
            
        # Optional: Meta openenv clamping
        # A single partial step clamp: (We only really need final score clamped, but reward clamped ensures log safety)
        
        return GitCommitReviewerObservation(
            commit_message=next_commit if not done else "DONE",
            commit_idx=self.current_idx,
            total_commits=len(self.current_task_commits),
            done=done,
            reward=reward,
            metadata={"step": self._state.step_count}
        )

    @property
    def state(self) -> State:
        return self._state