"""
Step 1: Import data and create annotation project
"""

from rm_gallery.core.data.annotation.annotation import create_annotator
from rm_gallery.core.data.load.base import create_loader
import rm_gallery.core.data     # Core strategy registration
import rm_gallery.gallery.data  # Extension strategy registration

# Replace with your actual API Token obtained from Label Studio
API_TOKEN = ""

# Step 1.1: Configure data loading parameters
load_config = {
    "path": "../../../data/reward-bench-2/data/test-00000-of-00001.parquet",  # Replace with actual data path
    "limit": 1000,  # Limit the number of records loaded to avoid loading too much data for initial testing
}

# Step 1.2: Create data loader
loader = create_loader(
    name="rewardbench2",           # Dataset identifier name
    load_strategy_type="local",    # Use local file loading strategy
    data_source="rewardbench2",    # Specify data source format converter
    config=load_config             # Pass loading configuration parameters
)

# Step 1.3: Execute data loading
print("Loading data...")
dataset = loader.run()
print(f"Data loading completed, {len(dataset.datasamples)} records total")

# Step 1.4: Configure annotation project parameters
annotation_config = {
    "server_url": "http://localhost:8080",           # Label Studio service address
    "api_token": API_TOKEN,                          # API access token
    "project_title": "RM Gallery Quality Annotation", # Project display name
    "template_name": "rewardbench2",                 # Use built-in RewardBench2 template
    "project_description": "Data quality annotation project based on RewardBench2 template"
}

# Step 1.5: Create annotation module instance
annotation_module = create_annotator(
    name="rm_gallery_annotation",
    **annotation_config
)

# Step 1.6: Create annotation project and import data
print("Creating annotation project...")
result = annotation_module.run(dataset, create_new_project=True)

if result:
    project_url = f"{result.metadata['annotation_server_url']}/projects/{result.metadata['annotation_project_id']}"
    print(f"✅ Annotation project created successfully!")
    print(f"🌐 Project URL: {project_url}")
    print(f"📊 Project ID: {result.metadata['annotation_project_id']}")
else:
    print("❌ Failed to create annotation project, please check configuration and network connection")

"""
Step 2: Perform data annotation

After the project is created successfully, you can:
1. Visit the project URL output above
2. Login using admin@rmgallery.com / RM-Gallery
3. Annotate data in the annotation interface
4. After annotation is complete, run the next step to export annotation results

Note: In actual usage, you need to manually complete the annotation work, then run the export code below.
"""

# This is a placeholder for annotation operations
# Actual annotation work needs to be completed in the Label Studio web interface
print("📝 Please complete data annotation work in the Label Studio interface")
print("💡 After annotation is complete, run the code below to export results")

"""
Step 3: Export annotation results
"""

from rm_gallery.core.data.annotation.annotation import create_annotator
from rm_gallery.core.data.export import create_exporter
import rm_gallery.core.data     # Core strategy registration
import rm_gallery.gallery.data  # Extension strategy registration

# Use the same API Token as when creating the project
API_TOKEN = ""

# Step 3.1: Recreate annotation module instance
annotation_module = create_annotator(
    name="rm_gallery_annotation",
    template_name="rewardbench2",
    api_token=API_TOKEN
)

# Step 3.2: Set project ID (obtained from Step 1 output)
annotation_module.project_id = 3  # Replace with actual project_id

# Step 3.3: Export annotation data from Label Studio
print("Exporting annotation data...")
try:
    annotated_dataset = annotation_module.export_annotations_to_dataset()
    print(f"✅ Annotation data exported successfully, {len(annotated_dataset.datasamples)} annotated records total")
except Exception as e:
    print(f"❌ Export failed: {e}")
    annotated_dataset = None

# Step 3.4: Configure file exporter
if annotated_dataset:
    exporter = create_exporter(
        name="annotation_exporter",
        config={
            "output_dir": "./exports",          # Export file storage directory
            "formats": ["jsonl"]        # Support multiple export formats
        }
    )

    # Step 3.5: Execute data export
    print("Saving annotation results to file...")
    export_result = exporter.run(annotated_dataset)

    if export_result:
        print("✅ Annotation results saved to ./exports directory")
        print(f"📊 Dataset info: {annotated_dataset.name}")
        print(f"📁 Contains {len(annotated_dataset.datasamples)} annotation data")
    else:
        print("❌ File export failed")

    # Display partial data preview
    if annotated_dataset.datasamples:
        print("\n📋 Data preview:")
        sample = annotated_dataset.datasamples[0]
        print(f"  - Sample ID: {sample.unique_id}")
        print(f"  - Annotation status: {sample.metadata.get('annotation_status', 'unknown')}")
        print(f"  - Output count: {len(sample.output) if sample.output else 0}")

from rm_gallery.core.data.annotation.template import BaseAnnotationTemplate, AnnotationTemplateRegistry

@AnnotationTemplateRegistry.register("custom_template")
class CustomAnnotationTemplate(BaseAnnotationTemplate):
    @property
    def label_config(self) -> str:
        """
        Define Label Studio annotation interface configuration
        Using Label Studio's XML configuration syntax
        """
        return """
        <View>
            <Text name="question" value="$question"/>
            <Choices name="quality" toName="question" choice="single-radio">
                <Choice value="excellent" background="green"/>
                <Choice value="good" background="blue"/>
                <Choice value="fair" background="yellow"/>
                <Choice value="poor" background="red"/>
            </Choices>
            <Rating name="score" toName="question" maxRating="10" />
            <TextArea name="comments" toName="question"
                     placeholder="Please enter evaluation reason..." rows="3"/>
        </View>
        """

    def process_annotations(self, annotation_data):
        """
        Process annotation data obtained from Label Studio
        Convert raw annotation data to structured format
        """
        processed_data = {
            "quality_rating": annotation_data.get("choices", {}).get("quality", {}).get("choices", []),
            "numerical_score": annotation_data.get("rating", {}).get("score", {}).get("rating", 0),
            "textual_feedback": annotation_data.get("textarea", {}).get("comments", {}).get("text", [""])[0]
        }
        return processed_data

# Import custom annotation template
from rm_gallery.gallery.data.annotation.custom_template import CustomAnnotationTemplate

annotation_module = create_annotation_module(
    name="custom_annotation_project",
    template_name="custom_template",  # Use your registered template name
    # ... other configuration parameters
)

# Annotated DataSample structure example
data_sample = DataSample(
    unique_id="sample_001",
    input=[...],  # Original input data
    output=[...], # Original output data (if any)
    metadata={
        "annotation_status": "completed",
        "annotation_project_id": 123,
        "annotator_id": "user@example.com"
    }
)

# The label field in each output contains annotation results
output.label = {
    "annotation_data": {
        "ratings": {...},      # Rating data
        "choices": {...},      # Choice data
        "text_areas": {...}    # Text input data
    },
    "processed": {...}         # Structured data processed by template
}