"""
SCOTUS Judgment Environment.

Task-specific environment for predicting Supreme Court case outcomes.
Works with any method agent (IDS-OED, BOED, PSRL, Reflexion, etc.).

Supports THINK, CLOSED_SEARCH, OPEN_WEB_SEARCH, and PROVIDE_FINAL_RESPONSE actions.
"""

import os
import logging
import json
import glob
from typing import Dict, Any, Optional, List, Set
from datetime import date
import datetime
from .base import Environment, Observation
from .utils import EM, get_timestamp, parse_date_to_timestamp
from ..prompts.environments.scotus_judgment import (
    get_environment_description,
    get_task_instance_description,
    get_initial_observation,
    get_search_capabilities,
    get_response_requirements,
)
from ..agents.action import ActionType, Action
from ..agents.action.closed_search.dense import search as search_dense_index, load_embedding_model, search_with_model
from ..agents.action.closed_search.bm25 import search as search_bm25_index
from ..agents.action.open_search.search import search as search_web
from agents_framework.agents.action.closed_search.metadata_filters import MetadataFilterPolicy  # type: ignore
from ..models.llm import ModelAPI


class ScotusMetadataIntentResolver:
    """
    LLM-driven intent resolver that keeps agents/environment prompts agnostic.
    Uses the LLM prompt to decide appropriate filters - no code-level fallbacks.
    """

    def __init__(
        self,
        model_api: Optional[ModelAPI],
        model_id: str = "gemini-2.5-pro",
        provider: str = "gemini",
        max_tokens: int = 10000,
    ):
        self.model_api = model_api
        self.model_id = model_id
        self.provider = provider
        self.max_tokens = max_tokens

    def infer_filters(
        self,
        query: str,
        search_type: str,
        get_field_values,
    ) -> Dict[str, str]:
        # Handle both agent-facing and internal search type names
        if search_type not in ("task_specific_documents", "docket_file_documents"):
            return {}

        available_titles = get_field_values("proceeding_title", search_type) or set()
        available_links = get_field_values("link_name", search_type) or set()
        filters: Dict[str, str] = {}

        if self.model_api and (available_titles or available_links):
            prompt = self._build_prompt(query, available_titles, available_links)
            try:
                response = self.model_api(
                    model_id=self.model_id,
                    prompt=[{"role": "user", "content": prompt}],
                    provider=self.provider,
                    max_tokens=self.max_tokens,
                    temperature=0.0,
                )
                filters = self._parse_response(response, available_titles, available_links)
                logger.debug(f"Metadata intent inference succeeded: {filters}")
            except Exception as exc:
                import traceback
                logger.warning(f"Metadata intent LLM classification failed: {exc}")
                logger.debug(f"Metadata intent inference traceback:\n{traceback.format_exc()}")
                # Return empty dict on failure - search will proceed without filters
                return {}

        return filters

    def _build_prompt(
        self, query: str, titles: Set[str], links: Set[str]
    ) -> str:
        titles_str = "\n".join(f"- {value}" for value in sorted(titles)) or "- NONE"
        links_str = "\n".join(f"- {value}" for value in sorted(links)) or "- NONE"
        return f"""You classify whether a closed-search query wants a specific docket document type.

Query: "{query}"

Available proceeding_title values (document type + party/proceeding labels):
{titles_str}

Available link_name values (administrative/filing metadata labels):
{links_str}

Rules:
- Return NONE for proceeding_title only when the query is very broad and should search all documents.
- For link_name, use this mapping:
  * "Petition" = the actual petition content for any petition
  * "Main Document" = the actual brief content for any brief (petitioner briefs, respondent briefs, amicus briefs)
  * "Appendix" = supporting materials attached to filings
  * "Proof of Service" / "Certificate of Word Count" = administrative metadata only (no substantive legal content)
- When the query asks for a "petition" or "writ of certiorari", set link_name to "Petition".
- When the query asks for any brief (petitioner, respondent, amicus), set link_name to "Main Document".
- Only return "Appendix","Proof of Service" or "Certificate of Word Count" if the query is looking for supporting materials

Respond exactly in two lines:
proceeding_title: <exact value from the list above or NONE>
link_name: <exact value from the list above or NONE>"""

    def _parse_response(
        self, response: str, titles: Set[str], links: Set[str]
    ) -> Dict[str, str]:
        filters: Dict[str, str] = {}
        if not response:
            return filters

        for line in response.strip().splitlines():
            line = line.strip()
            if line.lower().startswith("proceeding_title:"):
                value = line.split(":", 1)[1].strip()
                matched = self._match_value(value, titles)
                if matched:
                    filters["proceeding_title"] = matched
            elif line.lower().startswith("link_name:"):
                value = line.split(":", 1)[1].strip()
                matched = self._match_value(value, links)
                if matched:
                    filters["link_name"] = matched

        return filters

    def _match_value(self, value: str, candidates: Set[str]) -> Optional[str]:
        if not candidates or not value or value.upper() == "NONE":
            return None
        # Ensure value is a string
        if not isinstance(value, str):
            logger.warning(f"_match_value received non-string value: {value} (type: {type(value).__name__})")
            return None
        if value in candidates:
            return value

        value_lower = value.lower()
        for candidate in candidates:
            if value_lower in candidate.lower():
                return candidate
        return None


logger = logging.getLogger(__name__)

class SCOTUSJudgment(Environment):
    """
    SCOTUS Judgment Prediction environment that supports:
    - PROVIDE_FINAL_RESPONSE: Make legal prediction
    - THINK: Internal reasoning (no environment interaction)
    - CLOSED_SEARCH: Search indexed legal documents
    - OPEN_WEB_SEARCH: Search the web for legal information
    """
    
    def __init__(self, 
                 case_info: Dict[str, Any],
                 max_steps: int = 10, 
                 search_config: Dict[str, Any] = None,
                 embedding_model_name: str = "infgrad/dewey_en_beta",
                 task_specific_documents_dir: str = None,
                 metadocuments_dir: str = None,
                 device: str = None,
                 metadata_filter_max_tokens: int = 10000):
        """
        Initialize the SCOTUS Judgment Prediction environment.
        
        Args:
            case_info: Dictionary containing case information
            max_steps: Maximum number of steps allowed
            search_config: Configuration for search functionality
            embedding_model_name: Name of the embedding model to use for search
            task_specific_documents_dir: Path to directory containing task-specific documents
            metadocuments_dir: Path to directory containing metadocuments
            device: Device to use for embedding model ('cpu' or 'cuda'). Defaults to 'cpu' to avoid Triton/flash-attention issues.
        """
        # Initialize with required base class parameters
        # Handle multiple legal questions by joining them
        legal_questions = case_info.get('legal_questions', ['Unknown question'])
        question = '; '.join(legal_questions) if len(legal_questions) > 1 else legal_questions[0]
        
        # Keep legal_answers as a list for custom is_correct logic
        legal_answers = case_info.get('legal_answers', [])
        key = legal_answers  # Store as list
        
        action_space = [
            ActionType.PROVIDE_FINAL_RESPONSE,
            ActionType.THINK,
            ActionType.CLOSED_SEARCH,
            ActionType.OPEN_WEB_SEARCH,
        ]
        
        super().__init__(question=question, key=key, action_space=action_space, max_steps=max_steps)
        
        self.case_info = case_info
        
        # Search configuration
        self.search_config = search_config or {}
        self.search_top_k = self.search_config.get('top_k', 5)
        self.use_readability_reranking = self.search_config.get('use_readability_reranking', True)
        self.readability_bias = self.search_config.get('readability_bias', 0.4)
        self.search_index_type = self.search_config.get('search_index_type', 'dense')
        self.embedding_model_name = self.search_config.get('embedding_model', embedding_model_name)
        self.device = device
        self._metadata_model_api = self._init_metadata_model_api()
        self._metadata_filter_policy = self._build_metadata_filter_policy()
        
        # Index paths
        self.indexes = self.search_config.get('indexes', {})
        self.task_specific_indexes = self.indexes.get('task_specific_documents', {})
        self.meta_indexes = self.indexes.get('metadocuments', {})
        
        # Action space is already set by parent class
        
        # Track agent performance
        self.search_history = []
        
        # Preload embedding model for dense search if configured
        self.embedding_model = None
        self._preload_embedding_model()
        
        # Cache for available field values to avoid repeated file scanning
        self._field_values_cache = {}
        
        # Set document directories
        self.task_specific_documents_dir = task_specific_documents_dir
        self.metadocuments_dir = metadocuments_dir
        
        # Parse and cache resolution date filter once during initialization
        self._resolution_date_filter = self._parse_resolution_date_filter()
        
        # Store metadata filter max_tokens
        self.metadata_filter_max_tokens = metadata_filter_max_tokens
        
        logger.info(f"Initialized SCOTUSJudgment environment")
        logger.info(f"Case: {self.case_info.get('case_title', '')}")
        logger.info(f"Max steps: {self.max_steps}")
        logger.info(f"Task-specific documents directory: {self.task_specific_documents_dir}")
        logger.info(f"Metadocuments directory: {self.metadocuments_dir}")
        if self._resolution_date_filter:
            cutoff_date = datetime.datetime.fromtimestamp(self._resolution_date_filter['before'])
            logger.info(f"Date filter: documents before {cutoff_date.date()} (resolution_date={self.case_info.get('resolution_date', 'N/A')})")
        else:
            logger.warning("⚠️  No date filter applied - resolution_date could not be parsed or is missing!")
    
    def _preload_embedding_model(self):
        """Preload the embedding model if dense search is configured."""
        if self.search_index_type == "dense":
            try:
                logger.info("Preloading embedding model for dense search...")
                logger.info("About to call load_embedding_model...")
                logger.info(f"Model name: {self.embedding_model_name}")
                logger.info(f"Device: {self.device}")
                self.embedding_model = load_embedding_model(
                    model_name=self.embedding_model_name,
                    device=self.device
                )
                logger.info("load_embedding_model completed successfully")
                logger.info("✓ Embedding model preloaded successfully")
                logger.info(f"Embedding model type: {type(self.embedding_model)}")
            except Exception as e:
                logger.error(f"Failed to preload embedding model: {e}")
                import traceback
                logger.error(f"Full traceback: {traceback.format_exc()}")
                logger.warning("Will load model on-demand for each search")
                self.embedding_model = None
        else:
            logger.info("Dense search not configured, skipping model preload")

    def _init_metadata_model_api(self) -> Optional[ModelAPI]:
        """Initialize ModelAPI for metadata intent classification, if enabled."""
        if self.search_config.get("disable_metadata_intent_llm", False):
            return None
        try:
            return ModelAPI()
        except Exception as exc:
            logger.warning(f"Failed to initialize ModelAPI for metadata filters: {exc}")
            return None

    def _build_metadata_filter_policy(self) -> Optional[MetadataFilterPolicy]:
        """Configure metadata filter inference for docket file searches."""
        if self.search_config.get("disable_auto_metadata_filters", False):
            return None

        model_id = self.search_config.get("metadata_filter_model_id", "gemini-2.5-pro")
        provider = self.search_config.get("metadata_filter_provider", "gemini")

        resolver = ScotusMetadataIntentResolver(
            model_api=self._metadata_model_api,
            model_id=model_id,
            provider=provider,
            max_tokens=getattr(self, 'metadata_filter_max_tokens', 10000),
        )

        return MetadataFilterPolicy(
            resolver=resolver,
            get_field_values_fn=self._get_available_field_values,
            applicable_search_types=["docket_file_documents", "task_specific_documents"],  # Support both agent-facing and internal names
        )

    def _get_metadata_filter_policy_for(self, search_type: str) -> Optional[MetadataFilterPolicy]:
        """Return the metadata filter policy for a search type if enabled."""
        if self._metadata_filter_policy and self._metadata_filter_policy.applies_to(search_type):
            return self._metadata_filter_policy
        return None
    
    def _parse_resolution_date_filter(self) -> dict:
        """Parse resolution date filter once during initialization to prevent data leakage."""
        # Extract resolution date from case info
        resolution_date = self.case_info.get('resolution_date')
        if not resolution_date:
            logger.warning("⚠️  resolution_date is missing from case_info - no date filter will be applied!")
            return {}
        
        cutoff_timestamp = parse_date_to_timestamp(resolution_date)
        
        if cutoff_timestamp:
            return {
                'field': 'published_epoch_seconds',
                'before': cutoff_timestamp
            }
        else:
            logger.warning(f"⚠️  Could not parse resolution_date '{resolution_date}' - no date filter will be applied!")
            return {}
    
    def _get_resolution_date_filter(self) -> dict:
        """Get cached resolution date filter (parsed during initialization)."""
        return self._resolution_date_filter
    
    def _get_resolution_cutoff_date(self) -> Optional['date']:
        """Get cutoff date as date object for web search."""
        date_filter = self._get_resolution_date_filter()
        if date_filter and 'before' in date_filter:
            from datetime import datetime
            cutoff_timestamp = date_filter['before']
            return datetime.fromtimestamp(cutoff_timestamp).date()
        return None
    
    def _get_available_field_values(self, field_name: str, search_type: str = "task_specific_documents") -> Set[str]:
        """Get all unique values for a given field from task-specific documents or metadocuments."""
        cache_key = f"{field_name}_{search_type}"
        if cache_key in self._field_values_cache:
            return self._field_values_cache[cache_key]
        
        # Get directory based on search type (handle both agent-facing and internal names)
        if search_type in ("task_specific_documents", "docket_file_documents"):
            documents_dir = self.task_specific_documents_dir
            dir_name = "task_specific_documents_dir"
        elif search_type in ("metadocuments", "legal_research_documents"):
            documents_dir = self.metadocuments_dir
            dir_name = "metadocuments_dir"
        else:
            documents_dir = None
            dir_name = "unknown"
        
        if documents_dir is None:
            logger.warning(f"{dir_name} is None in config for search_type '{search_type}'")
            return set()
        
        if not os.path.exists(documents_dir):
            return set()
        
        # Read unique field values from JSON files
        field_values = set()
        json_files = glob.glob(os.path.join(documents_dir, "**", "*.json"), recursive=True)
        
        for json_file in json_files:
            try:
                with open(json_file, 'r') as f:
                    data = json.load(f)
                    if field_name in data and data[field_name] is not None:
                        field_values.add(str(data[field_name]).strip())
            except (json.JSONDecodeError, IOError) as e:
                logger.debug(f"Error reading {json_file}: {e}")
                continue
        
        self._field_values_cache[cache_key] = field_values
        logger.info(f"Cached {len(field_values)} unique values for field '{field_name}' in {search_type}")
        return field_values
    
    def step(self, action: Action) -> Observation:
        """Execute an action and return the resulting observation."""
        # Handle parsing failures by skipping the round
        if action is None:
            logger.warning("Agent returned None (parsing failed), skipping this round")
            return Observation(
                result="Round skipped due to parsing failure",
                metadata={"skipped": True, "reason": "parsing_failure"}
            )
        
        # Increment step counter (same as parent class)
        self.current_step += 1
        
        if action.action_type == ActionType.PROVIDE_FINAL_RESPONSE:
            # Handle legal prediction
            return self._handle_final_response_action(action)
            
        elif action.action_type == ActionType.THINK:
            # Handle internal reasoning - no environment interaction
            return self._handle_think_action(action)
            
        elif action.action_type == ActionType.CLOSED_SEARCH:
            # Handle indexed document search
            return self._handle_closed_search_action(action)
            
        elif action.action_type == ActionType.OPEN_WEB_SEARCH:
            # Handle web search
            return self._handle_web_search_action(action)
            
        else:
            raise NotImplementedError(f"Action type {action.action_type} not implemented in SCOTUSJudgment environment.")
    
    def _handle_final_response_action(self, action: Action) -> Observation:
        """Handle PROVIDE_FINAL_RESPONSE action (making legal prediction)."""
        response = getattr(action, "response", "")
        
        # Set the answer for the base class's is_correct() method
        self.answer = response
        
        # Terminate the episode after final response
        self.terminated = True
        
        # Use the custom is_correct method for ordered questions/answers
        accuracy = self.is_correct()
        is_fully_correct = accuracy == 1.0
        
        observation = Observation(
            result=f"Legal prediction submitted: {response}",
            metadata={
                "action_type": "PROVIDE_FINAL_RESPONSE",
                "response": response,
                "is_correct": is_fully_correct,
                "accuracy": accuracy
            }
        )
        return observation
    
    def _handle_think_action(self, action: Action) -> Observation:
        """Handle THINK action (internal reasoning)."""
        thought = getattr(action, "thought", "")
        
        observation = Observation(
            result=f"Internal reasoning: {thought}",
            metadata={
                "action_type": "THINK",
                "thought": thought
            }
        )
        return observation
    
    def _handle_closed_search_action(self, action: Action) -> Observation:
        """Handle CLOSED_SEARCH action (indexed document search)."""
        query = getattr(action, "query", "")
        search_type = getattr(action, "search_type", "task_specific_documents")
        
        # Extract optional parameters, using config defaults if not provided
        k = getattr(action, "k", self.search_top_k)
        
        # Note: field_filters are now handled automatically by MetadataFilterPolicy
        # which infers appropriate filters from the natural language query.
        # date_filter and include_null_dates are also handled automatically by environment
        # to prevent data leakage and ensure consistent execution details
        
        try:
            # Map agent-facing search_type to internal index name
            # Agent uses "docket_file_documents" -> internally "task_specific_documents"
            # Agent uses "legal_research_documents" -> internally "metadocuments"
            if search_type == "docket_file_documents":
                internal_search_type = "task_specific_documents"
            elif search_type == "legal_research_documents":
                internal_search_type = "metadocuments"
            else:
                internal_search_type = search_type  # Fallback for any other values
            
            # Determine which index to search
            if internal_search_type == "task_specific_documents":
                index_path = self.task_specific_indexes.get(self.search_index_type)
            elif internal_search_type == "metadocuments":
                index_path = self.meta_indexes.get(self.search_index_type)
            else:
                raise ValueError(f"Unknown search type: {search_type}. Use 'docket_file_documents' or 'legal_research_documents'.")
            
            if not index_path or not os.path.exists(index_path):
                raise FileNotFoundError(f"Index path not found: {index_path}")
            
            # Get automatic date filter to prevent data leakage
            date_filter = self._get_resolution_date_filter()
            metadata_filter_policy = self._get_metadata_filter_policy_for(search_type)
            
            # NEW APPROACH: Infer metadata filters FIRST, then apply them directly in search query
            # This avoids the prefetch problem - filters are applied natively by Qdrant
            inferred_field_filters = None
            if metadata_filter_policy and metadata_filter_policy.applies_to(search_type):
                try:
                    get_values = metadata_filter_policy.get_field_values_fn or (lambda _field, _stype: set())
                    candidates = metadata_filter_policy.resolver.infer_filters(query, search_type, get_values)
                    if candidates:
                        # Validate and clean candidates - ensure all values are valid types for Qdrant
                        # Qdrant MatchValue accepts: str, int, float, bool, or list of these
                        cleaned_filters = {}
                        for field, value in candidates.items():
                            if value is None:
                                logger.warning(f"Skipping filter {field}=None (None values not supported by Qdrant)")
                                continue
                            # Convert to supported type if needed
                            if isinstance(value, (str, int, float, bool, list)):
                                cleaned_filters[field] = value
                            else:
                                # Convert other types to string
                                logger.warning(f"Converting filter {field}={value} (type: {type(value).__name__}) to string for Qdrant")
                                cleaned_filters[field] = str(value)
                        inferred_field_filters = cleaned_filters if cleaned_filters else None
                    else:
                        inferred_field_filters = None
                    # Always log inferred filters (even if None/empty) BEFORE search
                    logger.debug(f"Inferred field filters: {inferred_field_filters if inferred_field_filters else None}")
                except Exception as exc:
                    logger.warning(f"Metadata intent inference failed: {exc}, proceeding without filters")
                    inferred_field_filters = None
                    logger.debug(f"Inferred field filters: {inferred_field_filters}")
            else:
                # No metadata filter policy - log that no filters will be applied
                logger.debug(f"Inferred field filters: {inferred_field_filters}")
            
            # Determine null date handling based on search type:
            # - task_specific_documents: strict (no null dates allowed)
            # - metadocuments: allow null dates for google_scholar (unlikely to have leakage)
            #                  but not for web results (could be news articles about decision)
            if internal_search_type == "metadocuments":
                allow_null_for_sources = ["google_scholar"]
            else:
                allow_null_for_sources = None
            
            # Perform search based on index type
            if self.search_index_type == "dense":
                # Determine backend and collection name based on path structure
                if "qdrant" in index_path:
                    backend = "qdrant"
                    collection_name = "task_specific_documents" if "task_specific" in index_path else "metadocuments"
                else:
                    backend = "faiss"
                    collection_name = None
                
                # Use preloaded model if available, otherwise fallback to loading on-demand
                if self.embedding_model is not None:
                    results = search_with_model(
                        model=self.embedding_model,
                        index_path=index_path,
                        query_text=query,
                        model_name=self.embedding_model_name,
                        k=k,
                        use_readability_reranking=self.use_readability_reranking,
                        readability_bias=self.readability_bias,
                        backend=backend,
                        collection_name=collection_name,
                        date_filter=date_filter,  # Automatic date cutoff
                        include_null_dates=False,  # Base: no null dates allowed
                        allow_null_for_sources=allow_null_for_sources,  # Override for specific sources
                        field_filters=inferred_field_filters,  # Apply inferred filters directly in Qdrant query
                        metadata_filter_policy=None,  # No post-processing filtering needed
                    )
                else:
                    # Fallback to loading model on-demand
                    results = search_dense_index(
                        index_path=index_path,
                        query_text=query,
                        model_name=self.embedding_model_name,
                        k=k,
                        cache_folder=os.environ.get('SENTENCE_TRANSFORMERS_HOME'),
                        use_readability_reranking=self.use_readability_reranking,
                        readability_bias=self.readability_bias,
                        backend=backend,
                        collection_name=collection_name,
                        date_filter=date_filter,  # Automatic date cutoff
                        include_null_dates=False,  # Base: no null dates allowed
                        allow_null_for_sources=allow_null_for_sources,  # Override for specific sources
                        field_filters=inferred_field_filters,  # Apply inferred filters directly in Qdrant query
                        metadata_filter_policy=None,  # No post-processing filtering needed
                    )
            elif self.search_index_type == "bm25":
                results = search_bm25_index(
                    index_path=index_path,
                    query=query,
                    k=k,
                    date_filter=date_filter,  # Automatic date cutoff
                    include_null_dates=False,  # Base: no null dates allowed
                    allow_null_for_sources=allow_null_for_sources,  # Override for specific sources
                    field_filters=inferred_field_filters,  # Apply inferred filters directly in Qdrant query
                    metadata_filter_policy=None,  # No post-processing filtering needed
                )
            else:
                raise ValueError(f"Unknown search index type: {self.search_index_type}")
            
            # Store the inferred field filters that were applied
            applied_filters = inferred_field_filters if inferred_field_filters else {}
            
            # Debug logging for empty results
            if len(results) == 0:
                logger.warning(f"⚠️  No results found for query: '{query}' (search_type: {search_type})")
                logger.warning(f"   Date filter: {date_filter}")
                logger.warning(f"   Applied field filters: {applied_filters}")
                logger.warning(f"   Index path: {index_path}")
                logger.warning(f"   include_null_dates: False")
            
            # Store date_filter and field_filters in observation metadata for logging by runner
            
            # Create structured search results with truncated contents (for observation.result - agent sees this)
            structured_results = []
            # Create structured search results with full contents (for observation.metadata - evaluation needs this)
            metadata_results = []
            
            for i, search_result in enumerate(results):
                # Get full contents (no truncation for metadata)
                full_contents = search_result.metadata.get('contents', '')
                
                # Get truncated contents for agent observation
                contents = full_contents
                if internal_search_type == "task_specific_documents" and len(contents) > 26500:  # ~6,600 tokens
                    contents = contents[:26500] + "\n\n[Document truncated - showing first 26,500 characters]"
                elif internal_search_type == "metadocuments" and len(contents) > 42000:  # ~10,500 tokens  
                    contents = contents[:42000] + "\n\n[Document truncated - showing first 42,000 characters]"
                
                # Helper function to create result dict
                def create_result_dict(contents_to_use):
                    # Build dict manually since closed_search SearchResult doesn't have to_dict()
                    result_dict = {
                        "result_id": search_result.result_id,
                        "title": search_result.title,
                        "url": search_result.url,
                        "snippet": search_result.snippet,
                        "source": search_result.source,
                        "published_date_raw": getattr(search_result, 'published_date_raw', None),
                        "published_date": search_result.published_date.isoformat() if hasattr(search_result, 'published_date') and search_result.published_date else None,
                        "metadata": search_result.metadata.copy() if search_result.metadata else {}
                    }
                    # Update metadata to include contents
                    result_dict['metadata']['contents'] = contents_to_use
                    return result_dict
                
                # Create truncated version for agent
                structured_results.append(create_result_dict(contents))
                # Create full version for metadata
                metadata_results.append(create_result_dict(full_contents))
            
            # Create structured observation result (matching observation_metadata format)
            observation_result = {
                "action_type": "CLOSED_SEARCH",
                "query": query,
                "search_type": search_type,
                "num_results": len(results),
                "search_results": structured_results
            }
            
            # Store search in history
            search_record = {
                "step": self.current_step,
                "query": query,
                "search_type": search_type,
                "num_results": len(results),
                "timestamp": get_timestamp()
            }
            self.search_history.append(search_record)
            
            observation = Observation(
                result=observation_result,
                metadata={
                    "action_type": "CLOSED_SEARCH",
                    "query": query,
                    "search_type": search_type,
                    "num_results": len(results),
                    "search_results": metadata_results,  # Full contents for evaluation (JSON serializable dicts)
                    "date_filter": date_filter if date_filter else None,
                    "field_filters": applied_filters if applied_filters else None
                }
            )
            
        except Exception as e:
            error_msg = str(e)
            import traceback
            error_traceback = traceback.format_exc()
            logger.error(f"Closed search failed: {error_msg}")
            logger.error(f"Full traceback:\n{error_traceback}")
            observation = Observation(
                result=f"Search failed: {error_msg}",
                metadata={
                    "action_type": "CLOSED_SEARCH",
                    "query": query,
                    "search_type": search_type,
                    "error": error_msg,
                    "error_traceback": error_traceback,
                    "search_results": []
                }
            )
        
        return observation
    
    def _handle_web_search_action(self, action: Action) -> Observation:
        """Handle OPEN_WEB_SEARCH action (web search)."""
        query = getattr(action, "query", "")
        search_type = getattr(action, "search_type", "web")  # Default to "web" if not specified
        k = getattr(action, "k", self.search_top_k)
        
        try:
            # Get automatic date cutoff to prevent data leakage
            cutoff_date = self._get_resolution_cutoff_date()
            
            # Perform web search with automatic date cutoff
            search_results = search_web(
                query=query, 
                search_type=search_type,  # Respect action's search_type parameter
                num_results=k,
                cutoff_date=cutoff_date,  # Automatic date cutoff to prevent data leakage
                exclude_undated=True,  # Conservative: exclude undated results
            )
            
            # Create structured search results
            structured_results = []
            for i, result in enumerate(search_results):
                # Get snippet as contents (most relevant part from Google)
                contents = result.snippet if hasattr(result, 'snippet') else ""
                
                # Use position from search results (lower is better)
                # Check for web_position, google_scholar_position, or position in metadata
                position = None
                if hasattr(result, 'metadata'):
                    position = (result.metadata.get('web_position') or 
                              result.metadata.get('google_scholar_position') or 
                              result.metadata.get('position') or 
                              i+1)
                else:
                    position = i+1
                
                # Create structured result dict
                result_dict = result.to_dict()
                # Add position to metadata if not already there
                if 'metadata' not in result_dict:
                    result_dict['metadata'] = {}
                result_dict['metadata']['position'] = position
                structured_results.append(result_dict)
            
            # Create structured observation result (matching observation_metadata format)
            observation_result = {
                "action_type": "OPEN_WEB_SEARCH",
                "query": query,
                "search_type": search_type,
                "num_results": len(search_results),
                "search_results": structured_results
            }
            
            # Store search in history
            search_record = {
                "step": self.current_step,
                "query": query,
                "search_type": search_type,  # Use actual search_type from action
                "num_results": len(search_results),
                "timestamp": get_timestamp()
            }
            self.search_history.append(search_record)
            
            observation = Observation(
                result=observation_result,
                metadata={
                    "action_type": "OPEN_WEB_SEARCH",
                    "query": query,
                    "search_type": search_type,
                    "num_results": len(search_results),
                    "search_results": structured_results  # Use structured dicts for JSON serialization
                }
            )
            
        except Exception as e:
            error_msg = str(e)
            logger.error(f"Web search failed: {error_msg}")
            observation = Observation(
                result=f"Web search failed: {error_msg}",
                metadata={
                    "action_type": "OPEN_WEB_SEARCH",
                    "query": query,
                    "error": error_msg,
                    "web_search_results": []
                }
            )
        
        return observation
    
    def get_environment_description(self) -> str:
        """Get environment description for agent prompts."""
        sanitized_case_info = self.get_sanitized_case_info()
        return get_environment_description(sanitized_case_info)
    
    def get_case_info(self) -> Dict[str, Any]:
        """Get case information."""
        return self.case_info.copy()
    
    def get_sanitized_case_info(self) -> Dict[str, Any]:
        """Get case information without legal_answers for agent consumption."""
        sanitized = self.case_info.copy()
        sanitized.pop('legal_answers', None)
        return sanitized

    def get_response_requirements(self) -> str:
        """Get task-specific response requirements for predictions."""
        return get_response_requirements()
    
    def get_search_capabilities(self) -> str:
        """Get description of search capabilities for the agent."""
        indexes_config = {}
        if self.task_specific_indexes.get(self.search_index_type):
            indexes_config["task_specific_documents"] = True
        if self.meta_indexes.get(self.search_index_type):
            indexes_config["metadocuments"] = True
        
        return get_search_capabilities(indexes_config)
    
    def get_task_instance_description(self) -> str:
        """Get formatted task instance description for agent prompts."""
        sanitized_case_info = self.get_sanitized_case_info()
        return get_task_instance_description(sanitized_case_info)
    
    
    def get_search_history(self) -> List[Dict[str, Any]]:
        """Get history of searches performed."""
        return self.search_history.copy()
    
    def is_correct(self) -> float:
        """Check accuracy of the current answer using ordered questions/answers logic.
        
        Returns:
            float: Accuracy as fraction of correct answers (0.0 to 1.0)
        """
        if not self.key or not self.answer:
            return 0.0
        
        # Parse the response to extract predicted answers
        # The response should contain answers delimited by semicolons
        # Format: "Answer1; Answer2; Answer3"
        
        # Split response by semicolons and clean up whitespace
        predicted_answers_raw = [answer.strip() for answer in self.answer.split(';')]
        
        # Normalize predicted answers (handle variations like yes, Yes, YES, Y)
        def normalize_answer(answer):
            answer_upper = answer.upper()
            if answer_upper in ['YES', 'Y']:
                return 'Yes'
            elif answer_upper in ['NO', 'N']:
                return 'No'
            else:
                return answer  # Keep original if not a recognized variation
        
        predicted_answers = [normalize_answer(answer) for answer in predicted_answers_raw]
        
        # Count correct answers
        correct_count = 0
        total_questions = len(self.key)
        
        if len(predicted_answers) == len(self.key):
            # Check each predicted answer against the corresponding correct answer
            for predicted, correct in zip(predicted_answers, self.key):
                if predicted == correct:
                    correct_count += 1
        else:
            # If counts don't match, count only the overlapping answers as correct
            min_length = min(len(predicted_answers), len(self.key))
            for i in range(min_length):
                if predicted_answers[i] == self.key[i]:
                    correct_count += 1
            # No credit for missing or extra answers
        
        # Return accuracy as fraction
        return correct_count / total_questions if total_questions > 0 else 0.0
    
    def is_done(self) -> bool:
        """Check if the environment is done (max steps reached or final prediction made)."""
        if self.is_truncated():
            return True
        
        # Check if a final prediction was made (base class sets terminated=True)
        if self.is_terminated():
            return True
        
        return False
    
    def reset(self):
        """Reset the environment to initial state."""
        super().reset()  # Call parent reset (resets self.answer, self.current_step, self.terminated)
        self.search_history = []
        
        # Create the initial observation using prompt module
        sanitized_case_info = self.get_sanitized_case_info()
        result = get_initial_observation(sanitized_case_info, self.max_steps)
        
        self.initial_observation = Observation(
            result=result,
            metadata={
                "case_title": self.case_info.get('case_title', ''),
                "case_docket": self.case_info.get('case_docket', ''),
                "legal_questions": self.case_info.get('legal_questions', []),
                "max_steps": self.max_steps,
                "available_actions": [action_type.value for action_type in self.action_space]
            }
        )
        
        logger.info("Environment reset")

    def get_initial_observation(self) -> Observation:
        """Get the initial observation with case information."""
        return self.initial_observation
    
    def get_stats(self) -> Dict[str, Any]:
        """Get environment statistics."""
        return {
            "current_step": self.current_step,
            "max_steps": self.max_steps,
            "has_answer": bool(self.answer),
            "num_searches": len(self.search_history),
            "is_done": self.is_done()
        }
