# chaplain_feedback_ui.py
"""
Gradio UI components for Chaplain Feedback & Tagging System.

Provides interface components for displaying classification flows,
collecting chaplain feedback, and displaying error patterns.

Requirements: 1.5, 2.3, 3.3, 4.1, 5.1, 5.3, 6.1, 6.3, 8.1, 8.2, 8.3, 10.1, 10.2, 10.3
"""

from __future__ import annotations

import gradio as gr
from typing import List, Dict, Tuple, Optional, Any
from dataclasses import dataclass

from src.core.chaplain_models import (
    ClassificationFlowResult,
    DistressIndicator,
    FollowUpQuestion,
    TaggingRecord,
    CLASSIFICATION_SUBCATEGORIES,
    QUESTION_ISSUE_TYPES,
    REFERRAL_ISSUE_TYPES,
)


class ChaplainFeedbackUIComponents:
    """Manages Gradio UI components for chaplain feedback system."""
    
    # Color mappings for classification badges
    BADGE_COLORS = {
        "red": "🔴",
        "yellow": "🟡",
        "green": "🟢",
    }
    
    BADGE_LABELS = {
        "red": "RED - Severe Distress",
        "yellow": "YELLOW - Potential Distress",
        "green": "GREEN - No Distress",
    }
    
    # Severity color codes for indicators
    SEVERITY_COLORS = {
        "red": "#ea9999",      # Red from definitions document
        "yellow": "#ffe599",   # Yellow from definitions document
    }
    
    @staticmethod
    def create_classification_flow_display() -> Tuple[gr.Component, gr.Component, gr.Component, gr.Component]:
        """
        Create ClassificationFlowDisplay component.
        
        Displays RED/YELLOW/GREEN flow results with all generated content.
        
        Returns:
            Tuple of (classification_badge, explanation, content_section, indicators_section) components
            
        Requirements: 1.5, 2.3, 3.3
        """
        classification_badge = gr.Markdown(
            value="🔄 Loading classification...",
            label="Classification Result",
        )
        
        explanation = gr.Markdown(
            value="",
            label="Explanation",
        )
        
        content_section = gr.Markdown(
            value="",
            label="Generated Content",
        )
        
        indicators_section = gr.Markdown(
            value="",
            label="Detected Indicators",
        )
        
        return classification_badge, explanation, content_section, indicators_section
    
    @staticmethod
    def render_classification_flow(
        flow_result: ClassificationFlowResult,
    ) -> Tuple[str, str, str, str]:
        """
        Render complete classification flow result.
        
        Args:
            flow_result: ClassificationFlowResult with all flow data
            
        Returns:
            Tuple of (badge, explanation, content, indicators) markdown strings
        """
        # Classification badge
        badge_emoji = ChaplainFeedbackUIComponents.BADGE_COLORS.get(flow_result.classification, "❓")
        badge_label = ChaplainFeedbackUIComponents.BADGE_LABELS.get(flow_result.classification, "UNKNOWN")
        confidence_pct = int(round(flow_result.confidence * 100))
        badge = f"## {badge_emoji} {badge_label}\n\n**Confidence:** {confidence_pct}%"
        
        # Explanation
        explanation = f"### Explanation\n\n{flow_result.explanation}"
        
        # Generated content based on classification
        content = ""
        if flow_result.classification == "red":
            content = ChaplainFeedbackUIComponents._render_red_flow_content(flow_result)
        elif flow_result.classification == "yellow":
            content = ChaplainFeedbackUIComponents._render_yellow_flow_content(flow_result)
        elif flow_result.classification == "green":
            content = ChaplainFeedbackUIComponents._render_green_flow_content(flow_result)
        
        # Indicators
        indicators = ChaplainFeedbackUIComponents._render_indicators(flow_result.indicators)
        
        return badge, explanation, content, indicators
    
    @staticmethod
    def _render_red_flow_content(flow_result: ClassificationFlowResult) -> str:
        """Render RED flow content (permission check + referral message)."""
        content = "### 🔴 RED FLAG - Severe Distress Detected\n\n"
        
        if flow_result.permission_check_message:
            content += "#### Patient Permission Check\n\n"
            content += f"{flow_result.permission_check_message}\n\n"
        
        if flow_result.consent_status:
            content += f"**Consent Status:** {flow_result.consent_status}\n\n"
        
        if flow_result.referral_message and flow_result.consent_status == "granted":
            content += "#### Referral Message for Spiritual Care Team\n\n"
            content += f"{flow_result.referral_message}\n\n"
        elif flow_result.consent_status == "declined":
            content += "**Status:** No further action - patient declined spiritual support referral\n\n"
        
        return content
    
    @staticmethod
    def _render_yellow_flow_content(flow_result: ClassificationFlowResult) -> str:
        """Render YELLOW flow content (follow-up questions + re-evaluation)."""
        content = "### 🟡 YELLOW FLAG - Potential Distress\n\n"
        
        if flow_result.follow_up_questions:
            content += "#### Follow-Up Questions\n\n"
            for i, question in enumerate(flow_result.follow_up_questions, 1):
                content += f"**Question {i}:** {question.question_text}\n\n"
                content += f"*Purpose:* {question.purpose}\n\n"
        
        if flow_result.patient_responses:
            content += "#### Patient Responses\n\n"
            for i, response in enumerate(flow_result.patient_responses, 1):
                content += f"**Response {i}:** {response}\n\n"
        
        if flow_result.re_evaluation_result:
            content += f"#### Re-Evaluation Result\n\n"
            if flow_result.re_evaluation_result == "red":
                content += "🔴 **Escalated to RED** - Severe distress detected in responses\n\n"
            elif flow_result.re_evaluation_result == "green":
                content += "🟢 **Downgraded to GREEN** - No distress indicators in responses\n\n"
        
        return content
    
    @staticmethod
    def _render_green_flow_content(flow_result: ClassificationFlowResult) -> str:
        """Render GREEN flow content (no distress)."""
        content = "### 🟢 GREEN FLAG - No Distress Detected\n\n"
        content += "**Status:** No further steps required\n\n"
        content += "No spiritual distress indicators were detected in this message.\n\n"
        return content
    
    @staticmethod
    def _render_indicators(indicators: List[DistressIndicator]) -> str:
        """Render detected indicators with categories and severity."""
        if not indicators:
            return "### Detected Indicators\n\nNo indicators detected"
        
        content = "### Detected Indicators\n\n"
        
        # Group by severity
        red_indicators = [i for i in indicators if i.severity == "red"]
        yellow_indicators = [i for i in indicators if i.severity == "yellow"]
        
        if red_indicators:
            content += "#### 🔴 RED Indicators (Severe)\n\n"
            for indicator in red_indicators:
                confidence_pct = int(round(indicator.confidence * 100))
                content += f"• **{indicator.subcategory}** ({confidence_pct}% confidence)\n"
                content += f"  - Category: {indicator.category}\n"
                content += f"  - Reference: {indicator.definition_reference}\n\n"
        
        if yellow_indicators:
            content += "#### 🟡 YELLOW Indicators (Potential)\n\n"
            for indicator in yellow_indicators:
                confidence_pct = int(round(indicator.confidence * 100))
                content += f"• **{indicator.subcategory}** ({confidence_pct}% confidence)\n"
                content += f"  - Category: {indicator.category}\n"
                content += f"  - Reference: {indicator.definition_reference}\n\n"
        
        return content
    
    @staticmethod
    def create_tagging_interface() -> Tuple[gr.Component, gr.Component, gr.Component, gr.Component, gr.Component, gr.Component, gr.Component, gr.Component, gr.Component, gr.Component]:
        """
        Create TaggingInterface component.
        
        Provides classification subcategory selector, multi-select for issues,
        and free-text comment fields.
        
        Returns:
            Tuple of individual tagging components for use in event handlers
            
        Requirements: 4.1, 5.1, 5.3, 6.1, 6.3
        """
        # Classification tagging components
        is_correct = gr.Radio(
            choices=[("✓ Correct", True), ("✗ Incorrect", False)],
            label="Is the classification correct?",
            interactive=True,
            visible=False,
        )
        
        subcategory = gr.Dropdown(
            choices=CLASSIFICATION_SUBCATEGORIES,
            label="What type of error? (if incorrect)",
            interactive=True,
            visible=False,
        )
        
        correct_classification = gr.Radio(
            choices=[
                ("🟢 GREEN - No Distress", "green"),
                ("🟡 YELLOW - Potential Distress", "yellow"),
                ("🔴 RED - Severe Distress", "red"),
            ],
            label="What should the correct classification be?",
            interactive=True,
            visible=False,
        )
        
        # Follow-up question issues components
        question_issues = gr.CheckboxGroup(
            choices=QUESTION_ISSUE_TYPES,
            label="Issues with follow-up questions (select all that apply)",
            interactive=True,
            visible=False,
        )
        
        question_comments = gr.Textbox(
            label="Comments on questions",
            placeholder="e.g., 'Too clinical', 'Not spiritually relevant'",
            lines=2,
            interactive=True,
            visible=False,
        )
        
        # Referral message issues components
        referral_issues = gr.CheckboxGroup(
            choices=REFERRAL_ISSUE_TYPES,
            label="Issues with referral message (select all that apply)",
            interactive=True,
            visible=False,
        )
        
        referral_comments = gr.Textbox(
            label="Comments on referral message",
            placeholder="e.g., 'Incomplete summary', 'Tone inappropriate'",
            lines=2,
            interactive=True,
            visible=False,
        )
        
        # Indicator issues components
        indicator_issues = gr.Textbox(
            label="Incorrectly identified indicators",
            placeholder="List indicator IDs or names that were incorrectly identified",
            lines=2,
            interactive=True,
            visible=False,
        )
        
        indicator_comments = gr.Textbox(
            label="Comments on indicators",
            placeholder="e.g., 'Missed anxiety indicators', 'False positive on grief'",
            lines=2,
            interactive=True,
            visible=False,
        )
        
        # General notes component
        notes_section = gr.Textbox(
            label="General Notes",
            placeholder="Any additional feedback or observations",
            lines=3,
            interactive=True,
            visible=False,
        )
        
        return is_correct, subcategory, correct_classification, question_issues, question_comments, referral_issues, referral_comments, indicator_issues, indicator_comments, notes_section
    
    @staticmethod
    def create_indicator_display() -> Tuple[gr.Component, gr.Component]:
        """
        Create IndicatorDisplay component.
        
        Shows indicators with categories and allows tagging incorrect indicators.
        
        Returns:
            Tuple of (indicators_display, indicator_tagging) components
            
        Requirements: 8.1, 8.2, 8.3
        """
        indicators_display = gr.Markdown(
            value="No indicators to display",
            label="Detected Indicators",
        )
        
        indicator_tagging = gr.Group(visible=False)
        with indicator_tagging:
            incorrect_indicators = gr.CheckboxGroup(
                choices=[],
                label="Select indicators that are incorrectly identified",
                interactive=True,
            )
            
            indicator_notes = gr.Textbox(
                label="Why are these indicators incorrect?",
                placeholder="Explain why these indicators don't apply",
                lines=2,
                interactive=True,
            )
        
        return indicators_display, indicator_tagging
    
    @staticmethod
    def create_error_pattern_summary() -> Tuple[gr.Component, gr.Component, gr.Component]:
        """
        Create ErrorPatternSummary component.
        
        Displays error patterns grouped by type with frequent subcategories highlighted.
        
        Returns:
            Tuple of (error_patterns, subcategory_breakdown, recommendations) components
            
        Requirements: 10.1, 10.2, 10.3
        """
        error_patterns = gr.Markdown(
            value="No error patterns yet",
            label="Error Patterns",
        )
        
        subcategory_breakdown = gr.Markdown(
            value="No data",
            label="Subcategory Breakdown",
        )
        
        recommendations = gr.Markdown(
            value="No recommendations yet",
            label="Recommendations for Improvement",
        )
        
        return error_patterns, subcategory_breakdown, recommendations
    
    @staticmethod
    def render_error_patterns(
        classification_errors: Dict[str, int],
        question_errors: Dict[str, int],
        referral_errors: Dict[str, int],
    ) -> Tuple[str, str, str]:
        """
        Render error patterns summary.
        
        Args:
            classification_errors: Dict of classification error subcategories with counts
            question_errors: Dict of question issue types with counts
            referral_errors: Dict of referral issue types with counts
            
        Returns:
            Tuple of (patterns, breakdown, recommendations) markdown strings
        """
        # Error patterns grouped by type
        patterns = "### Error Patterns\n\n"
        
        total_classification_errors = sum(classification_errors.values())
        total_question_errors = sum(question_errors.values())
        total_referral_errors = sum(referral_errors.values())
        
        if total_classification_errors > 0:
            patterns += f"#### Classification Errors: {total_classification_errors} total\n\n"
            for subcategory, count in sorted(classification_errors.items(), key=lambda x: x[1], reverse=True):
                patterns += f"• {subcategory}: {count}\n"
            patterns += "\n"
        
        if total_question_errors > 0:
            patterns += f"#### Follow-Up Question Issues: {total_question_errors} total\n\n"
            for issue_type, count in sorted(question_errors.items(), key=lambda x: x[1], reverse=True):
                patterns += f"• {issue_type}: {count}\n"
            patterns += "\n"
        
        if total_referral_errors > 0:
            patterns += f"#### Referral Message Issues: {total_referral_errors} total\n\n"
            for issue_type, count in sorted(referral_errors.items(), key=lambda x: x[1], reverse=True):
                patterns += f"• {issue_type}: {count}\n"
            patterns += "\n"
        
        # Subcategory breakdown
        breakdown = "### Subcategory Breakdown\n\n"
        
        if classification_errors:
            breakdown += "**Classification Errors:**\n"
            for subcategory, count in sorted(classification_errors.items(), key=lambda x: x[1], reverse=True):
                breakdown += f"- {subcategory}: {count}\n"
            breakdown += "\n"
        
        if question_errors:
            breakdown += "**Question Issues:**\n"
            for issue_type, count in sorted(question_errors.items(), key=lambda x: x[1], reverse=True):
                breakdown += f"- {issue_type}: {count}\n"
            breakdown += "\n"
        
        if referral_errors:
            breakdown += "**Referral Issues:**\n"
            for issue_type, count in sorted(referral_errors.items(), key=lambda x: x[1], reverse=True):
                breakdown += f"- {issue_type}: {count}\n"
            breakdown += "\n"
        
        # Recommendations
        recommendations = "### Recommendations for Improvement\n\n"
        
        # Find most common errors
        all_errors = {}
        for subcategory, count in classification_errors.items():
            all_errors[f"Classification: {subcategory}"] = count
        for issue_type, count in question_errors.items():
            all_errors[f"Questions: {issue_type}"] = count
        for issue_type, count in referral_errors.items():
            all_errors[f"Referral: {issue_type}"] = count
        
        if all_errors:
            sorted_errors = sorted(all_errors.items(), key=lambda x: x[1], reverse=True)
            top_3 = sorted_errors[:3]
            
            recommendations += "**Top areas for improvement:**\n\n"
            for error_type, count in top_3:
                recommendations += f"1. **{error_type}** ({count} occurrences)\n"
                recommendations += f"   - Review prompts and logic for this error type\n"
                recommendations += f"   - Consider additional training data\n\n"
        else:
            recommendations += "No errors detected yet. Great job!\n\n"
        
        return patterns, breakdown, recommendations