import re
from typing import Any, Dict, List

from rm_gallery.core.data.annotation.template import (
    AnnotationTemplateRegistry,
    BaseAnnotationTemplate,
)
from rm_gallery.core.data.schema import DataSample


@AnnotationTemplateRegistry.register("asio_finance")
class AsioFinanceAnnotationTemplate(BaseAnnotationTemplate):
    """Reward Bench 2 annotation template implementation for 4-way comparison"""

    def __init__(self, name: str):
        super().__init__(name)

    @property
    def label_config(self) -> str:
        """Return the Label Studio XML configuration for reward bench 2 evaluation (4-way comparison)"""
        return """
<View>
  <!-- Sample Information -->
  <!-- Input Messages -->
  <Header value="问题(包含历史对话)"/>
  <Paragraphs name="input_dialogue" value="$input_messages" layout="dialogue" nameKey="role" textKey="content" />
  
  <!-- ["asio-fin","qwen-dr", "kimi-chat", "tongyi", "gpt5-dr", "doubao-dr"] -->
  <Header value="各模型回答"/>
  
  <Collapse bordered="true">
    <Panel value="asio_old">
        <View>
          <HyperText name="answer_asio_old" value="$answer_asio_old" />
        </View>
    </Panel>
  </Collapse>

    <Panel value="asio_qwen3_max">
        <View>
          <HyperText name="answer_asio_qwen3_max" value="$answer_asio_qwen3_max" />
        </View>
    </Panel>
  </Collapse>

  <Collapse bordered="true">
    <Panel value="kimi_chat">
        <View>
          <HyperText name="answer_kimi_chat" value="$answer_kimi_chat" />
        </View>
    </Panel>
  </Collapse>

  <Collapse bordered="true">
    <Panel value="tongyi_dr">
        <View>
          <HyperText name="answer_tongyi_dr" value="$answer_tongyi_dr" />
        </View>
    </Panel>
  </Collapse>

  <Collapse bordered="true">
    <Panel value="gpt5_dr">
        <View>
          <HyperText name="answer_gpt5_dr" value="$answer_gpt5_dr" />
        </View>
    </Panel>
  </Collapse>

  <Collapse bordered="true">
    <Panel value="doubao_dr">
        <View>
          <HyperText name="answer_doubao-dr" value="$answer_doubao_dr" />
        </View>
    </Panel>
  </Collapse>

  

  <Header value="大模型评测得分"/>
  <HyperText name="dimension_scores_json" value="$dimension_scores" />

  <Header value="大模型评分理由"/>
  <HyperText name="dimension_reason_json" value="$dimension_reasons" />

  <!-- Step 1: Best Answer Selection -->
  <View>
    <Header value="选择最优回答"/>
    <Choices name="best_answer" toName="dimension_reason_json" choice="single" title="🏆 Best Answer">
      <Choice value="asio_old" />
      <Choice value="asio_qwen3_max" />
      <Choice value="doubao_dr" />
      <Choice value="gpt5_dr" />
      <Choice value="kimi_chat" />
      <Choice value="qwen_dr" />
      <Choice value="tongyi_dr" />
      <Choice value="all-equal" />
    </Choices>
  </View>

  <!-- Step 4: Additional Comments -->
  <View>
    <Header value="说出你的理由"/>
    <TextArea name="additional_comments" toName="dimension_scores_json" placeholder="因为..." title="Additional Comments"/>
  </View>

</View>
"""

    def process_annotations(self, annotation_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process annotation data specific to reward bench 2 evaluation (4-way comparison)

        Args:
            annotation_data: Generic annotation data with ratings, choices, text_areas

        Returns:
            Processed data structured for reward bench 2 evaluation
        """
        processed = {
            "best_answer": None,
            "answer_rankings": {},
            "answer_ratings": {},
            "ranking_order": [],
            "quality_comparison": {},
            "comments": "",
            "preference": None,
        }

        # Extract best answer selection (Step 1)
        if "best_answer" in annotation_data.get("choices", {}):
            best_answer_choices = annotation_data["choices"]["best_answer"]["choices"]
            if best_answer_choices:
                processed["best_answer"] = best_answer_choices[0]
                processed["preference"] = best_answer_choices[0]

        

        # Extract additional comments (Step 4)
        if "additional_comments" in annotation_data.get("text_areas", {}):
            processed["comments"] = annotation_data["text_areas"][
                "additional_comments"
            ]["text"]

        return processed

    def validate_annotation_data(self, annotation_data: Dict[str, Any]) -> bool:
        """
        Validate annotation data for reward bench 2 evaluation

        Args:
            annotation_data: Annotation data to validate

        Returns:
            True if valid, False otherwise
        """
        # Check if required fields are present
        required_sections = ["choices", "ratings"]
        for section in required_sections:
            if section not in annotation_data:
                return False

        # Check if best answer is selected
        if "best_answer" not in annotation_data.get("choices", {}):
            return False
        return True
    def process_annotations_v2(self, origin_sample:DataSample, annotation_datas):
        
        # Generic annotation data structure - collect all results by type
        annotation_data = {
            "ratings": {},  # All rating fields
            "choices": {},  # All choice fields
            "text_areas": {},  # All text area fields
            # "raw_results": ann_results,  # Keep original results for reference
        }
        for annotation in annotation_datas:
            ann_results = annotation["annotations"][0].get("result", [])

            # Process annotation results generically
            for result in ann_results:
                field_name = result.get("from_name", "")
                value = result.get("value", {})
                result_type = result.get("type", "")

                # Handle different control types generically
                if "rating" in value:
                    # Rating controls (star ratings, etc.)
                    annotation_data["ratings"][field_name] = {
                        "rating": value.get("rating", 0),
                        "type": result_type,
                    }

                elif "choices" in value or "choice" in value:
                    # Choice controls (radio, checkbox, etc.)
                    choices = value.get("choices", [])
                    if not choices and "choice" in value:
                        choices = [value.get("choice")]
                    if not annotation_data["choices"].get(field_name):
                        annotation_data["choices"][field_name] = []
                    annotation_data["choices"][field_name].extend(value["choices"])

                elif "text" in value:
                    # Text controls (textarea, text input, etc.)
                    text = value.get("text", [])
                    if not isinstance(text, list):
                        text = [text]
                    if not annotation_data["text_areas"].get(field_name):
                        annotation_data["text_areas"][field_name] = []
                    annotation_data["text_areas"][field_name].extend(text)
        for output in origin_sample.output:
          answer = output.answer
          answer.label["preference"] = "rejected"
          model_name = output.answer.label.get("model_name", "")
          if model_name in annotation_data["choices"]["best_answer"]:
              answer.label["preference"] = "chosen"
              if annotation_data["text_areas"].get("additional_comments"):
                answer.label["reason"] = annotation_data["text_areas"]["additional_comments"]
        return origin_sample
    
    def _markdown_to_html(self, markdown_text: str) -> str:
        """Convert basic markdown to HTML"""
        if not markdown_text:
            return ""
        
        html = markdown_text
        
        # Convert headers
        html = re.sub(r'^### (.*?)$', r'<h3>\1</h3>', html, flags=re.MULTILINE)
        html = re.sub(r'^## (.*?)$', r'<h2>\1</h2>', html, flags=re.MULTILINE)
        html = re.sub(r'^# (.*?)$', r'<h1>\1</h1>', html, flags=re.MULTILINE)
        
        # Convert bold text
        html = re.sub(r'\*\*(.*?)\*\*', r'<strong>\1</strong>', html)
        
        # Convert italic text
        html = re.sub(r'\*(.*?)\*', r'<em>\1</em>', html)
        
        # Convert line breaks
        html = html.replace('\n', '<br>')
        
        # Convert tables (basic support)
        lines = html.split('<br>')
        in_table = False
        result_lines = []
        
        for line in lines:
            if '|' in line and not line.strip().startswith('|'):
                if not in_table:
                    result_lines.append('<table style="border-collapse: collapse; width: 100%;">')
                    in_table = True
                # Skip separator lines
                if not re.match(r'^\s*\|[\s\-\|]+\|\s*$', line):
                    cells = [cell.strip() for cell in line.split('|') if cell.strip()]
                    if cells:
                        result_lines.append('<tr>')
                        for cell in cells:
                            result_lines.append(f'<td style="border: 1px solid #ddd; padding: 8px;">{cell}</td>')
                        result_lines.append('</tr>')
            else:
                if in_table:
                    result_lines.append('</table>')
                    in_table = False
                result_lines.append(line)
        
        if in_table:
            result_lines.append('</table>')
        
        html = '<br>'.join(result_lines)
        
        return html

    def _json_to_html(self, json_data: Dict, title: str = "数据详情") -> str:
        """Convert JSON data to structured HTML format"""
        if not json_data:
            return f'<div style="padding: 10px; color: #666;">暂无{title}</div>'
        
        html_parts = [f'<div style="font-family: Arial, sans-serif; margin: 10px 0;">']
        
        for key, value in json_data.items():
            if isinstance(value, dict):
                # Create collapsible section for nested objects
                html_parts.append(f'''
                <details style="margin: 8px 0; border: 1px solid #e0e0e0; border-radius: 6px; background-color: #f9f9f9;">
                    <summary style="padding: 12px; background-color: #e8f4fd; cursor: pointer; font-weight: bold; color: #2c5aa0; border-radius: 6px 6px 0 0;">
                        📊 {key}
                    </summary>
                    <div style="padding: 12px; background-color: #ffffff; border-radius: 0 0 6px 6px;">
                        {self._format_dict_content(value)}
                    </div>
                </details>
                ''')
            else:
                # Simple key-value pair
                html_parts.append(f'''
                <div style="margin: 6px 0; padding: 8px; background-color: #f0f8ff; border-left: 4px solid #4a90e2; border-radius: 4px;">
                    <strong style="color: #2c5aa0;">{key}:</strong> 
                    <span style="color: #333; font-family: monospace;">{value}</span>
                </div>
                ''')
        
        html_parts.append('</div>')
        return ''.join(html_parts)
    
    def _format_dict_content(self, data: Dict) -> str:
        """Format dictionary content as HTML"""
        html_parts = []
        
        for key, value in data.items():
            # Color code based on value
            if isinstance(value, (int, float)):
                if value >= 0.8:
                    color = "#28a745"  # Green for high scores
                elif value >= 0.6:
                    color = "#ffc107"  # Yellow for medium scores
                else:
                    color = "#dc3545"  # Red for low scores
            else:
                color = "#333"
            
            html_parts.append(f'''
            <div style="margin: 4px 0; padding: 6px; background-color: #f8f9fa; border-radius: 4px; display: flex; justify-content: space-between; align-items: center;">
                <span style="font-weight: 500; color: #495057;">{key}:</span>
                <span style="font-family: monospace; font-weight: bold; color: {color}; background-color: #ffffff; padding: 2px 8px; border-radius: 4px; border: 1px solid #dee2e6;">
                    {value}
                </span>
            </div>
            ''')
        
        return ''.join(html_parts)

    def _reasons_to_html(self, reasons_data: Dict, title: str = "评分理由") -> str:
        """Convert dimension reasons to structured HTML format with collapsible sections"""
        if not reasons_data:
            return f'<div style="padding: 10px; color: #666;">暂无{title}</div>'
        
        html_parts = [f'<div style="font-family: Arial, sans-serif; margin: 10px 0;">']
        
        for dimension, comparisons in reasons_data.items():
            if not comparisons:
                continue
                
            # Create main dimension section
            html_parts.append(f'''
            <details style="margin: 10px 0; border: 1px solid #d1ecf1; border-radius: 8px; background-color: #f8f9fa;">
                <summary style="padding: 15px; background-color: #d1ecf1; cursor: pointer; font-weight: bold; color: #0c5460; border-radius: 8px 8px 0 0; font-size: 16px;">
                    📝 {dimension}
                </summary>
                <div style="padding: 15px; background-color: #ffffff; border-radius: 0 0 8px 8px;">
                    {self._format_comparison_reasons(comparisons)}
                </div>
            </details>
            ''')
        
        html_parts.append('</div>')
        return ''.join(html_parts)
    
    def _format_comparison_reasons(self, comparisons: Dict) -> str:
        """Format comparison reasons with nested collapsible sections"""
        html_parts = []
        
        for comparison_key, reason_text in comparisons.items():
            if not reason_text:
                continue
                
            # Truncate long text for summary
            summary_text = reason_text[:100] + "..." if len(reason_text) > 100 else reason_text
            
            html_parts.append(f'''
            <details style="margin: 8px 0; border: 1px solid #e9ecef; border-radius: 6px; background-color: #f8f9fa;">
                <summary style="padding: 12px; background-color: #e9ecef; cursor: pointer; font-weight: 500; color: #495057; border-radius: 6px 6px 0 0;">
                    <span style="color: #6c757d; font-size: 12px; margin-right: 8px;">🔍</span>
                    <strong>{comparison_key}:</strong>
                    <span style="color: #6c757d; font-size: 13px; margin-left: 8px;">{summary_text}</span>
                </summary>
                <div style="padding: 12px; background-color: #ffffff; border-radius: 0 0 6px 6px; border-left: 4px solid #17a2b8;">
                    <div style="line-height: 1.6; color: #333; font-size: 14px; white-space: pre-wrap; word-wrap: break-word;">
                        {reason_text}
                    </div>
                </div>
            </details>
            ''')
        
        return ''.join(html_parts)

    def _convert_to_tasks(self, data_samples: List[DataSample]):
        tasks = []
        
        for i, sample in enumerate(data_samples):

            ##临时策略
            _models = [output.answer.label["model_name"] for output in sample.output]
            # if "asio_0912" not in _models:
            #     continue
            models = set()
            task_data = {
                "id": sample.unique_id,
            }

            # Process input messages - List[ChatMessage] -> dialogue format
            input_dialogue = []
            if sample.input:
                for msg in sample.input:
                    input_dialogue.append({"role": msg.role, "content": msg.content})
            task_data["input_messages"] = input_dialogue

            # Process output - List[DataOutput] -> dialogue format with multiple answers
            output_dialogue = {}
            answer_count = 0
            output_dialogue["default"] = "[No output available]"
            if sample.output:
                for output in sample.output:
                    # Process answer (Step)
                    if output.answer:
                        answer_count += 1
                        # Convert markdown content to HTML
                        # Remove leading and trailing quotes if present
                        content = output.answer.content.strip()
                        if content.startswith('"') and content.endswith('"'):
                            content = content[1:-1]
                        # Unescape any escaped quotes
                        content = content.replace('\\"', '"').replace('\\n', '\n')
                        html_content = self._markdown_to_html(content)
                        model = f"{output.answer.label['model_name']}"
                        models.add(model)
                        output_dialogue[model] = html_content

            dimension_details = sample.metadata.get("dimension_details", {})
            dimension_reasons = {}
            for dimension,details in dimension_details.items():
                dimension_reasons[dimension] = {}
                comparison_results = details.get("comparison_results", {})
                for vs,response in comparison_results.items():
                    if response.get("response",{}).get("reason"):
                        dimension_reasons[dimension][vs] = response["response"]["reason"]

            # Create HTML content with collapsible answers
            
            task_data["output_messages"] = output_dialogue
            # task_data["answer_demo"] = output_dialogue["qwen-dr"]
            for model in models:
                answer = output_dialogue.get(model, output_dialogue["default"])
                key = model
                task_data[f"answer_{key.replace('-', '_')}"] = answer
            task_data["answer_count"] = str(answer_count)
            # Convert dimension_scores to HTML format
            dimension_scores_data = sample.metadata.get("dimension_scores", {})
            task_data["dimension_scores"] = self._json_to_html(dimension_scores_data, "维度评分")
            # Convert dimension_reasons to HTML format
            task_data["dimension_reasons"] = self._reasons_to_html(dimension_reasons, "评分理由")
            tasks.append({"data": task_data})

        return tasks