import json
import logging
from pydantic import BaseModel, Field, validator, field_validator
from typing import Literal, Optional, Any, Dict, Type, TypeVar, Union, List

logger = logging.getLogger(__name__)
from ..registry.registry import register_model

T = TypeVar('T', bound='ReflectionCritique')

# --- Enumerations for Clarity ---

# Common component types for better IDE support and documentation
COMMON_COMPONENTS = [
    "PlannerAgent",
    "DiagnosisAgent",
    "ToolCall_CMDB",
    "ToolCall_Metrics",
    "TemporalWorkflow",
    "LangGraphStateNode",
    "PromptTemplate",
    'ToolRegistryService',
    'ExternalLLMService',
    'ActivityError',
    'HTTPCommunication'
]

# FailingComponent is now a string with validation
def validate_failing_component(value: str) -> str:
    """Validate that the failing component is a non-empty string.
    
    Args:
        value: The component name to validate
        
    Returns:
        The validated component name
        
    Raises:
        ValueError: If the value is not a string or is empty
    """
    if not isinstance(value, str) or not value.strip():
        raise ValueError("Failing component must be a non-empty string")
    return value

# Type alias for documentation and IDE support
FailingComponent = str

# Define the severity of the identified issue
IssueSeverity = Literal[
    "Critical_Tool_Failure",
    "Major_Logic_Error",
    "Minor_Prompt_Drift",
    "Schema_Mismatch",
    'Critical_Service_Unreachable',
    'Critical_Server_Error_5xx',
    'Minor_Tool_Schema',
    'Major_Read_Timeout'
]

# --- Update 3: Fix Scope ---
FixScope = Literal[
    'Prompt',
    'Code',
    'Configuration',
    'Tool_Schema',
    # NEW TECHNICAL SCOPE:
    'Infrastructure_Kubernetes',
    'Timeout_Setting'
]
@register_model
class ReflectionSuccessOptimization(BaseModel):
    """Structured output for optimizing successful agent executions."""
    trace_id: str = Field(..., description="The unique Langfuse ID of the trace being optimized")
    optimization_opportunity: str = Field(..., description="Brief description of the optimization opportunity")
    impact_level: Literal["Low", "Medium", "High"] = Field(..., description="Expected impact of implementing the optimization")
    optimization_type: Literal["Efficiency", "Completeness", "Accuracy", "UserExperience", "CostSaving"]
    suggested_improvement: str = Field(..., description="Detailed suggestion for improvement")
    expected_benefit: str = Field(..., description="Expected benefits of implementing the suggestion")
    confidence_score: int = Field(ge=0, le=100, default=50, description="Confidence in this optimization (0-100)")
    optimization_priority: int = Field(ge=1, le=5, default=3, description="Suggested implementation priority (1-5)")
    related_components: List[str] = Field(default_factory=list, description="Components that would be affected by this optimization")

# --- Main Structured Output ---
@register_model
class ReflectionCritique(BaseModel):
    """
    A structured output generated by the Reflection Agent after analyzing a failed or
    sub-optimal Langfuse trace. This object is the key deliverable for the Governance layer (JARVIS v1).
    """
    trace_id: str = Field(..., description="The unique Langfuse ID of the trace that was analyzed.")
    run_status: Literal["Failure", "Suboptimal"] = Field(...,
                                                         description="The final status of the run being critiqued.")

    # --- Core Findings ---
    failing_component: str = Field(
        ...,
        description="The specific component, agent, or tool that ultimately caused the failure."
    )
    
    root_cause_summary: str = Field(
        ...,
        description="A concise, one-sentence summary of the single root cause."
    )
    
    issue_severity: IssueSeverity = Field(
        ...,
        description="The assessed severity of the root cause."
    )

    # --- Suggested Fix & Confidence ---
    suggested_fix_action: str = Field(
        ...,
        description="A detailed, actionable recommendation to fix the root cause."
    )
    
    fix_scope: FixScope = Field(
        ...,
        description="The area where the fix needs to be applied."
    )
    
    fix_confidence_score: int = Field(
        default=50,
        ge=1,
        le=100,
        description="Confidence score (1-100) for the suggested fix. Defaults to 50."
    )

    # --- Audit Score ---
    critique_score: int = Field(
        default=3,
        ge=0,
        le=5,
        description="Quality score (0-5) for the failed run. Defaults to 3."
    )

    # --- Detailed Analysis ---
    detailed_analysis: str = Field(
        ...,
        description="Detailed explanation of the root cause and fix rationale."
    )
    is_capability_gap: bool = Field(
        default=False,
        description="True if the failure was caused by a missing tool or environmental state (e.g., ServiceNow hibernation)."
    )

    ace_bullet_type: Literal["PITFALL", "STRATEGY", "MISSING_CAPABILITY"] = Field(
        default="PITFALL",
        description="The category of the lesson for the Knowledge Graph."
    )

    missing_tool_blueprint: Optional[Dict[str, Any]] = Field(
        default=None,
        description="A JSON schema defining the requirements for a new tool if is_capability_gap is True."
    )

    # --- Pillar Mapping (Functionality Preservation) ---
    # We map your existing pillars directly into the ACE Loop:
    # root_cause_summary -> ACE Pitfall Content
    # issue_severity -> Utility Score weighting
    # suggested_fix_action -> ACE Strategy Content
    # detailed_analysis -> Blueprint Context
    
    # --- Validation ---
    @field_validator('failing_component')
    def validate_failing_component(cls, v: str) -> str:
        """
        Ensure failing_component is a valid non-empty string.
        
        This validator is intentionally permissive to allow any non-empty string
        while providing validation for common issues like whitespace-only values.
        """
        try:
            # Use the standalone validator function for consistency
            return validate_failing_component(v)
        except ValueError as e:
            # Provide a more helpful error message
            raise ValueError(
                f"Invalid failing component: {str(e)}. "
                f"Common components include: {', '.join(COMMON_COMPONENTS)}"
            ) from e
    
    @field_validator('root_cause_summary')
    def validate_root_cause_summary(cls, v: str) -> str:
        """Ensure root_cause_summary is a non-empty string."""
        if not isinstance(v, str) or not v.strip():
            raise ValueError("Root cause summary must be a non-empty string")
        return v.strip()
    
    @field_validator('suggested_fix_action')
    def validate_suggested_fix(cls, v: str) -> str:
        """Ensure suggested_fix_action is a non-empty string."""
        if not isinstance(v, str) or not v.strip():
            raise ValueError("Suggested fix must be a non-empty string")
        return v.strip()
    
    @field_validator('detailed_analysis')
    def validate_detailed_analysis(cls, v: str) -> str:
        """Ensure detailed_analysis is a non-empty string."""
        if not isinstance(v, str) or not v.strip():
            raise ValueError("Detailed analysis must be a non-empty string")
        return v.strip()
    
    # --- Serialization ---
    def model_dump_json(self, **kwargs) -> str:
        """Serialize to JSON with proper error handling."""
        try:
            return super().model_dump_json(**kwargs)
        except Exception as e:
            logger.error(f"Failed to serialize ReflectionCritique: {str(e)}")
            # Return a minimal valid JSON object with error information
            return json.dumps({
                "error": "Failed to serialize reflection critique",
                "trace_id": self.trace_id,
                "failing_component": getattr(self, 'failing_component', 'unknown'),
                "issue_severity": getattr(self, 'issue_severity', 'unknown')
            })
    
    @classmethod
    def parse_raw(cls: Type[T], json_data: Union[str, bytes], **kwargs) -> T:
        """Parse JSON data with improved error handling."""
        try:
            if isinstance(json_data, bytes):
                json_data = json_data.decode('utf-8')
            return super().model_validate_json(json_data, **kwargs)
        except Exception as e:
            logger.error(f"Failed to parse ReflectionCritique: {str(e)}\nData: {json_data[:500]}")
            raise ValueError(f"Invalid reflection critique data: {str(e)}") from e
