#!/usr/bin/env python3
"""
Simple test script to verify the environment works correctly.
"""

from server.doc_quality_env_environment import DocQualityEnvironment
from models import DocQualityAction

print("=" * 70)
print("Doc Quality Environment - Local Test")
print("=" * 70)

# Create environment
env = DocQualityEnvironment()

# Reset
obs = env.reset()
print(f"\n✓ Environment reset successfully")
print(f"  Task: {obs.task_name}")
print(f"  Difficulty: {obs.task_difficulty}")
print(f"  Doc preview: {obs.current_doc[:100]}...")
print(f"  Step: {obs.step_count}/{obs.max_steps}")

# Take a step
action = DocQualityAction(
    action_type="identify_issue",
    content="Missing response format documentation",
    issue_category="completeness"
)
obs = env.step(action)
print(f"\n✓ Step 1 executed successfully")
print(f"  Reward: {obs.reward}")
print(f"  Feedback: {obs.feedback}")
print(f"  Issues found: {len(obs.issues_identified)}")

# Take another step
action = DocQualityAction(
    action_type="suggest_improvement",
    content="Add a section documenting all response fields and their types",
    issue_category="completeness"
)
obs = env.step(action)
print(f"\n✓ Step 2 executed successfully")
print(f"  Reward: {obs.reward}")
print(f"  Total issues found: {len(obs.issues_identified)}")

# Rate quality
action = DocQualityAction(
    action_type="rate_quality",
    content="0.5",
    issue_category="overall"
)
obs = env.step(action)
print(f"\n✓ Step 3 (rating) executed successfully")
print(f"  Reward: {obs.reward}")
print(f"  Episode complete: {obs.done}")
print(f"  Final quality score: {obs.quality_score}")

print("\n" + "=" * 70)
print("✓ All tests passed!")
print("=" * 70)