# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Data models for the Government Service Application Assistant Environment.
"""

from typing import List, Optional, Dict, Any, Literal
from pydantic import BaseModel, Field
from openenv.core.env_server.types import Action, Observation


class DocumentInfo(BaseModel):
    """Information about a required or submitted document."""
    type: str = Field(..., description="Document type (e.g., 'Aadhaar Card')")
    description: Optional[str] = Field(None, description="Additional details about the document")
    mandatory: bool = Field(default=True, description="Whether the document is mandatory")
    validity_period: Optional[str] = Field(None, description="Required validity (e.g., 'Must be current')")
    name_match_required: bool = Field(default=False, description="Whether name must exactly match application")
    format_requirements: List[str] = Field(default_factory=list, description="Specific format requirements")


class SubmittedDocument(BaseModel):
    """A document submitted by the user for validation."""
    type: str = Field(..., description="Document type")
    details: str = Field(..., description="Details about the submitted document")
    validation_status: Literal["valid", "invalid", "missing_info"] = Field(
        default="valid", description="Validation status"
    )
    validation_reason: Optional[str] = Field(
        None, description="Reason if invalid or missing info"
    )


class ValidationResult(BaseModel):
    """Results of document validation."""
    is_complete: bool = Field(default=False, description="Whether all required documents are present")
    is_valid: bool = Field(default=False, description="Whether all submitted documents are valid")
    missing_documents: List[str] = Field(default_factory=list, description="List of missing required document types")
    invalid_documents: List[Dict[str, str]] = Field(
        default_factory=list, description="List of invalid documents with reasons"
    )
    valid_documents: List[str] = Field(default_factory=list, description="List of valid document types")


class CorrectionSuggestion(BaseModel):
    """A suggestion for correcting document issues."""
    issue_type: Literal["missing", "invalid_format", "name_mismatch", "expired", "insufficient"] = Field(
        ..., description="Type of issue"
    )
    current_document: Optional[str] = Field(
        None, description="The document that needs to be changed (if applicable)"
    )
    suggested_action: str = Field(..., description="Action to take to fix the issue")
    alternative_documents: List[str] = Field(
        default_factory=list, description="Alternative documents that could be used"
    )


class GovServiceState(BaseModel):
    """The internal state of the government service environment."""
    service_type: Optional[str] = Field(None, description="Type of government service")
    user_profile: Dict[str, Any] = Field(default_factory=dict, description="User profile information")
    current_stage: Literal[
        "service_selection", "document_identification", "validation", "correction", "submission", "completed"
    ] = Field(default="service_selection", description="Current stage of the application process")
    required_documents: List[DocumentInfo] = Field(
        default_factory=list, description="List of required documents for the selected service"
    )
    submitted_documents: List[SubmittedDocument] = Field(
        default_factory=list, description="Documents submitted by the user"
    )
    validation_results: Optional[ValidationResult] = Field(
        None, description="Results from document validation"
    )
    correction_suggestions: List[CorrectionSuggestion] = Field(
        default_factory=list, description="Suggestions for correcting document issues"
    )
    application_id: Optional[str] = Field(None, description="Generated application ID if submitted")
    step_count: int = Field(default=0, description="Number of steps taken in current episode")
    max_steps: int = Field(default=10, description="Maximum steps allowed per episode")
    done: bool = Field(default=False, description="Whether the episode is complete")


class GovAction(Action):
    """Action for the Government Service Application Assistant Environment."""

    # Backward compatibility - keep original message field
    message: str = Field(default="", description="Message (for backward compatibility)")

    # New structured action fields
    action_type: Optional[Literal[
        "select_service",
        "list_required_documents",
        "validate_documents",
        "suggest_corrections",
        "submit_application"
    ]] = Field(None, description="Type of action to perform")

    # Service selection
    service_type: Optional[str] = Field(
        None, description="Service type for select_service action"
    )

    # Document validation
    documents: Optional[List[SubmittedDocument]] = Field(
        None, description="Documents to validate or suggest corrections for"
    )


class GovObservation(Observation):
    """Observation from the Government Service Application Assistant Environment."""

    # Backward compatibility - keep original fields
    echoed_message: str = Field(default="", description="Echoed message (for backward compatibility)")
    message_length: int = Field(default=0, description="Message length (for backward compatibility)")

    # New feedback message
    message: str = Field(default="", description="Feedback message from the environment")

    # Current state information
    current_stage: str = Field(default="service_selection", description="Current stage of processing")
    service_type: Optional[str] = Field(None, description="Currently selected service type")

    # Document information
    required_documents: List[Dict[str, Any]] = Field(
        default_factory=list, description="Required documents for the service"
    )
    submitted_documents: List[Dict[str, Any]] = Field(
        default_factory=list, description="Currently submitted documents"
    )

    # Validation results
    validation_results: Optional[Dict[str, Any]] = Field(
        None, description="Results of document validation"
    )

    # Correction suggestions
    correction_suggestions: List[Dict[str, Any]] = Field(
        default_factory=list, description="Suggestions for correcting issues"
    )

    # Application status
    application_id: Optional[str] = Field(None, description="Application ID if submitted")
    is_complete: bool = Field(default=False, description="Whether application is complete and valid")