# 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.

"""
Government Service Application Assistant Environment Implementation.

A simple implementation that simulates government service document assistance.
"""

from uuid import uuid4
from typing import Dict, Any, Optional

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

from models import GovAction, GovObservation


class GovEnvironment(Environment):
    """
    Government Service Application Assistant Environment.

    This environment simulates an AI assistant helping users with government service applications.
    """

    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        """Initialize the government service environment."""
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._service_type: Optional[str] = None
        self._current_stage = "service_selection"
        self._required_documents: list[Dict[str, Any]] = []
        self._submitted_documents: list[Dict[str, Any]] = []
        self._validation_results: Optional[Dict[str, Any]] = None
        self._correction_suggestions: list[Dict[str, Any]] = []
        self._application_id: Optional[str] = None
        self._done = False

    def reset(self) -> GovObservation:
        """
        Reset the environment to initial state.

        Returns:
            GovObservation with welcome message
        """
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._service_type = None
        self._current_stage = "service_selection"
        self._required_documents = []
        self._submitted_documents = []
        self._validation_results = None
        self._correction_suggestions = []
        self._application_id = None
        self._done = False

        return GovObservation(
            message="Government Service Assistant ready. Please select a service to begin.",
            current_stage=self._current_stage,
            reward=0.0,
            done=False,
        )

    def step(self, action: GovAction) -> GovObservation:  # type: ignore[override]
        """
        Execute a step in the environment.

        Args:
            action: GovAction containing the action to perform

        Returns:
            GovObservation with updated state and feedback
        """
        self._state.step_count += 1

        # Handle the action
        observation = self._handle_action(action)

        # Check if episode should end
        observation.done = self._done or self._state.step_count >= 10

        return observation

    def _handle_action(self, action: GovAction) -> GovObservation:
        """Handle the specific action type."""
        # For backward compatibility, if no action_type, treat as echo
        if not action.action_type:
            return GovObservation(
                echoed_message=action.message,
                message_length=len(action.message),
                message=f"Echo: {action.message}",
                current_stage=self._current_stage,
                reward=min(len(action.message) * 0.1, 1.0),
                done=False,
            )

        if action.action_type == "select_service":
            return self._handle_select_service(action)
        elif action.action_type == "list_required_documents":
            return self._handle_list_required_documents(action)
        elif action.action_type == "validate_documents":
            return self._handle_validate_documents(action)
        elif action.action_type == "suggest_corrections":
            return self._handle_suggest_corrections(action)
        elif action.action_type == "submit_application":
            return self._handle_submit_application(action)
        else:
            return GovObservation(
                message=f"Unknown action type: {action.action_type}",
                current_stage=self._current_stage,
                reward=-0.1,
                done=self._done,
            )

    def _handle_select_service(self, action: GovAction) -> GovObservation:
        """Handle service selection action."""
        if not action.service_type:
            return GovObservation(
                message="Please specify a service type to select.",
                current_stage=self._current_stage,
                reward=-0.05,
                done=False,
            )

        service_type = action.service_type
        available_services = ["passport_new", "passport_renewal", "pan_new", "aadhaar_update", "driving_license_new"]

        if service_type not in available_services:
            return GovObservation(
                message=f"Unknown service type: {service_type}. Available: {available_services}",
                current_stage=self._current_stage,
                reward=-0.1,
                done=False,
            )

        # Update state
        self._service_type = service_type
        self._required_documents = self._get_required_documents(service_type)
        self._current_stage = "document_identification"

        return GovObservation(
            message=f"Selected service: {service_type}. {len(self._required_documents)} documents required.",
            current_stage=self._current_stage,
            service_type=self._service_type,
            required_documents=self._required_documents,
            reward=0.1,
            done=False,
        )

    def _handle_list_required_documents(self, action: GovAction) -> GovObservation:
        """Handle listing required documents."""
        if not self._service_type:
            return GovObservation(
                message="Please select a service first.",
                current_stage=self._current_stage,
                reward=-0.05,
                done=False,
            )

        return GovObservation(
            message=f"Required documents for {self._service_type}:",
            current_stage=self._current_stage,
            service_type=self._service_type,
            required_documents=self._required_documents,
            reward=0.05,
            done=False,
        )

    def _handle_validate_documents(self, action: GovAction) -> GovObservation:
        """Handle document validation."""
        if not action.documents:
            return GovObservation(
                message="Please provide documents to validate.",
                current_stage=self._current_stage,
                reward=-0.05,
                done=False,
            )

        # Update submitted documents
        self._submitted_documents = [doc.model_dump() for doc in action.documents]

        # Perform validation
        validation_results = self._validate_documents(action.documents)
        self._validation_results = validation_results

        # Update stage
        self._current_stage = "validation"

        # Generate message
        if validation_results["is_complete"] and validation_results["is_valid"]:
            message = "All documents are valid and complete. You can now submit the application."
            self._current_stage = "submission"
        elif validation_results["is_complete"] and not validation_results["is_valid"]:
            message = f"Documents are complete but some are invalid. Found {len(validation_results['invalid_documents'])} invalid documents."
            self._current_stage = "correction"
        else:
            message = f"Documents are incomplete. Missing {len(validation_results['missing_documents'])} required documents."
            self._current_stage = "correction"

        return GovObservation(
            message=message,
            current_stage=self._current_stage,
            service_type=self._service_type,
            submitted_documents=self._submitted_documents,
            validation_results=validation_results,
            reward=0.0,
            done=False,
        )

    def _handle_suggest_corrections(self, action: GovAction) -> GovObservation:
        """Handle suggesting corrections."""
        if not self._validation_results:
            return GovObservation(
                message="Please validate documents first.",
                current_stage=self._current_stage,
                reward=-0.05,
                done=False,
            )

        corrections = self._generate_corrections()
        self._correction_suggestions = corrections

        if corrections:
            message = f"Generated {len(corrections)} correction suggestions."
        else:
            message = "No corrections needed - all documents appear valid."

        return GovObservation(
            message=message,
            current_stage=self._current_stage,
            service_type=self._service_type,
            submitted_documents=self._submitted_documents,
            validation_results=self._validation_results,
            correction_suggestions=corrections,
            reward=0.0,
            done=False,
        )

    def _handle_submit_application(self, action: GovAction) -> GovObservation:
        """Handle application submission."""
        if not self._validation_results:
            return GovObservation(
                message="Please validate documents before submitting.",
                current_stage=self._current_stage,
                reward=-0.1,
                done=False,
            )

        validation = self._validation_results
        if not validation["is_complete"] or not validation["is_valid"]:
            return GovObservation(
                message="Cannot submit - documents are incomplete or invalid.",
                current_stage=self._current_stage,
                reward=-0.1,
                done=False,
            )

        # Generate application ID and complete
        application_id = f"APP-{str(uuid4())[:8].upper()}"
        self._application_id = application_id
        self._current_stage = "completed"
        self._done = True

        return GovObservation(
            message=f"Application submitted successfully! ID: {application_id}",
            current_stage=self._current_stage,
            service_type=self._service_type,
            application_id=application_id,
            is_complete=True,
            reward=0.5,
            done=True,
        )

    def _get_required_documents(self, service_type: str) -> list[Dict[str, Any]]:
        """Get required documents for a service type."""
        if service_type == "passport_new":
            return [
                {"type": "Proof of Address", "description": "Aadhaar/Voter ID/Utility bill", "mandatory": True},
                {"type": "Proof of Date of Birth", "description": "Birth certificate/PAN/Aadhaar", "mandatory": True},
                {"type": "Photograph", "description": "Recent passport-size photo", "mandatory": True}
            ]
        elif service_type == "passport_renewal":
            return [
                {"type": "Old Passport", "description": "Original passport to renew", "mandatory": True},
                {"type": "Proof of Address", "description": "Aadhaar/Voter ID/Utility bill", "mandatory": True},
                {"type": "Photograph", "description": "Recent passport-size photo", "mandatory": True}
            ]
        elif service_type == "pan_new":
            return [
                {"type": "Proof of Identity", "description": "Aadhaar/Voter ID/Driving license", "mandatory": True},
                {"type": "Proof of Address", "description": "Aadhaar/Voter ID/Utility bill", "mandatory": True},
                {"type": "Proof of Date of Birth", "description": "Birth certificate/PAN/Aadhaar", "mandatory": True}
            ]
        elif service_type == "aadhaar_update":
            return [
                {"type": "Proof of Address", "description": "Utility bill/Bank statement/Rental agreement", "mandatory": True},
                {"type": "Proof of Identity", "description": "PAN/Passport/Voter ID", "mandatory": True}
            ]
        elif service_type == "driving_license_new":
            return [
                {"type": "Proof of Age", "description": "Birth certificate/School certificate", "mandatory": True},
                {"type": "Proof of Address", "description": "Aadhaar/Voter ID/Utility bill", "mandatory": True},
                {"type": "Application Form", "description": "Form 2 or Form 4", "mandatory": True}
            ]
        return []

    def _validate_documents(self, documents: list) -> Dict[str, Any]:
        """Validate submitted documents."""
        result = {
            "is_complete": False,
            "is_valid": False,
            "missing_documents": [],
            "invalid_documents": [],
            "valid_documents": []
        }

        required_types = {doc["type"] for doc in self._required_documents if doc["mandatory"]}
        submitted_types = {doc.type for doc in documents}

        # Check missing
        result["missing_documents"] = list(required_types - submitted_types)

        # Check invalid
        for doc in documents:
            is_valid = True
            reason = ""

            details_lower = doc.details.lower()

            if "expired" in details_lower or "2023" in details_lower:
                is_valid = False
                reason = "Document appears to be expired"
            elif doc.type == "Proof of Address" and "spouse" in details_lower:
                is_valid = False
                reason = "Must be in applicant's name"
            elif doc.type == "Proof of Date of Birth" and "10th marksheet" in details_lower:
                is_valid = False
                reason = "Not sufficient as standalone DOB proof"

            if is_valid:
                result["valid_documents"].append(doc.type)
            else:
                result["invalid_documents"].append({"type": doc.type, "reason": reason})

        result["is_complete"] = len(result["missing_documents"]) == 0
        result["is_valid"] = len(result["invalid_documents"]) == 0

        return result

    def _generate_corrections(self) -> list[Dict[str, Any]]:
        """Generate correction suggestions."""
        corrections = []

        if not self._validation_results:
            return corrections

        # Handle missing documents
        for missing in self._validation_results["missing_documents"]:
            if missing == "Proof of Address":
                corrections.append({
                    "issue_type": "missing",
                    "suggested_action": "Submit a valid address proof",
                    "alternative_documents": ["Aadhaar Card", "Voter ID", "Utility Bill"]
                })
            elif missing == "Proof of Date of Birth":
                corrections.append({
                    "issue_type": "missing",
                    "suggested_action": "Submit a valid DOB proof",
                    "alternative_documents": ["Birth Certificate", "PAN Card", "Aadhaar Card"]
                })

        # Handle invalid documents
        for invalid in self._validation_results["invalid_documents"]:
            doc_type = invalid["type"]
            reason = invalid["reason"]

            if reason == "Document appears to be expired":
                corrections.append({
                    "issue_type": "expired",
                    "current_document": doc_type,
                    "suggested_action": f"Replace expired {doc_type} with current version",
                    "alternative_documents": ["Updated version of same document"]
                })
            elif "applicant's name" in reason:
                corrections.append({
                    "issue_type": "name_mismatch",
                    "current_document": doc_type,
                    "suggested_action": f"Submit {doc_type} in applicant's name",
                    "alternative_documents": ["Alternative address proof in applicant's name"]
                })

        return corrections

    @property
    def state(self) -> State:
        """Get the current environment state."""
        return self._state