"""
TemporalQueryHandler: Temporal-aware query processing for Helix.

Handles point-in-time queries and temporal context assembly.
"""

from __future__ import annotations

from datetime import datetime, timezone
from dataclasses import dataclass, field
from typing import Any

from helix.utils.temporal_utils import (
    extract_temporal_params,
    is_temporal_query,
    TemporalParams,
)


@dataclass
class TemporalQueryHandler:
    """
    Handler for temporal-aware queries.
    
    Responsibilities:
    - Detect temporal intent in queries
    - Extract temporal parameters (point-in-time, ranges)
    - Modify graph/vector queries for temporal filtering
    - Assemble temporal context for response generation
    """
    
    # Graphiti instance for temporal graph queries
    graphiti: Any = None
    
    # Default time window for "recent" queries
    default_time_window_days: int = 30
    
    async def detect_temporal_intent(
        self,
        query: str,
    ) -> TemporalParams:
        """
        Detect and extract temporal parameters from a query.
        
        Args:
            query: The query string.
            
        Returns:
            TemporalParams with detected temporal information.
        """
        return extract_temporal_params(query)
    
    async def apply_temporal_filter(
        self,
        results: list[dict],
        params: TemporalParams,
    ) -> list[dict]:
        """
        Apply temporal filtering to search results.
        
        Args:
            results: List of search results.
            params: Temporal parameters for filtering.
            
        Returns:
            Filtered results based on temporal constraints.
        """
        if not params.has_temporal_intent:
            return results
        
        filtered = []
        for result in results:
            # Check timestamp fields
            valid_time = result.get("valid_time") or result.get("t_valid")
            
            if valid_time is None:
                # No temporal data, include by default
                filtered.append(result)
                continue
            
            # Parse timestamp if string
            if isinstance(valid_time, str):
                try:
                    valid_time = datetime.fromisoformat(valid_time.replace("Z", "+00:00"))
                except ValueError:
                    filtered.append(result)
                    continue
            
            # Point-in-time filtering
            if params.valid_at:
                if valid_time <= params.valid_at:
                    filtered.append(result)
                continue
            
            # Range filtering
            if params.valid_from and valid_time < params.valid_from:
                continue
            if params.valid_to and valid_time > params.valid_to:
                continue
            
            filtered.append(result)
        
        return filtered
    
    async def search_temporal_graph(
        self,
        query: str,
        valid_at: datetime | None = None,
        num_results: int = 10,
    ) -> list[dict]:
        """
        Search the temporal knowledge graph.
        
        Args:
            query: Search query.
            valid_at: Point in time for historical query.
            num_results: Maximum results to return.
            
        Returns:
            List of matching facts with temporal metadata.
        """
        if self.graphiti is None:
            return []
        
        # Use Graphiti's search with temporal parameters
        try:
            results = await self.graphiti.search(
                query=query,
                num_results=num_results,
            )
            
            # Convert Graphiti results to standard format
            formatted = []
            for result in results:
                formatted.append({
                    "fact": result.fact if hasattr(result, "fact") else str(result),
                    "score": result.score if hasattr(result, "score") else 0.0,
                    "valid_time": getattr(result, "valid_at", None),
                    "created_time": getattr(result, "created_at", None),
                    "source": getattr(result, "source_description", "unknown"),
                })
            
            return formatted
        except Exception as e:
            print(f"Temporal graph search error: {e}")
            return []
    
    async def assemble_temporal_context(
        self,
        query: str,
        graph_results: list[dict],
        vector_results: list[dict],
        params: TemporalParams,
    ) -> dict[str, Any]:
        """
        Assemble temporal context for the LLM.
        
        Combines graph and vector results with temporal annotations.
        
        Args:
            query: Original query.
            graph_results: Results from graph search.
            vector_results: Results from vector search.
            params: Extracted temporal parameters.
            
        Returns:
            Context dict ready for LLM prompt.
        """
        context = {
            "query": query,
            "has_temporal_intent": params.has_temporal_intent,
            "temporal_keywords": params.temporal_keywords,
            "point_in_time": params.valid_at.isoformat() if params.valid_at else None,
            "facts": [],
            "passages": [],
        }
        
        # Add graph facts with temporal annotations
        for result in graph_results:
            fact_entry = {
                "fact": result.get("fact", ""),
                "confidence": result.get("score", 0.0),
            }
            
            if result.get("valid_time"):
                fact_entry["valid_at"] = result["valid_time"]
                fact_entry["is_current"] = not result.get("is_invalidated", False)
            
            context["facts"].append(fact_entry)
        
        # Add vector passages
        for result in vector_results:
            context["passages"].append({
                "content": result.get("content", ""),
                "score": result.get("score", 0.0),
                "source": result.get("metadata", {}).get("source", "unknown"),
            })
        
        return context
    
    def format_temporal_prompt(
        self,
        context: dict[str, Any],
    ) -> str:
        """
        Format temporal context as a prompt section.
        
        Args:
            context: Assembled temporal context.
            
        Returns:
            Formatted prompt string.
        """
        parts = []
        
        if context.get("has_temporal_intent"):
            parts.append("**Temporal Context:**")
            if context.get("point_in_time"):
                parts.append(f"- Query refers to time: {context['point_in_time']}")
            if context.get("temporal_keywords"):
                parts.append(f"- Temporal keywords: {', '.join(context['temporal_keywords'])}")
            parts.append("")
        
        if context.get("facts"):
            parts.append("**Known Facts from Knowledge Graph:**")
            for fact in context["facts"]:
                valid_info = ""
                if fact.get("valid_at"):
                    valid_info = f" [valid at: {fact['valid_at']}]"
                parts.append(f"- {fact['fact']}{valid_info}")
            parts.append("")
        
        if context.get("passages"):
            parts.append("**Relevant Passages:**")
            for passage in context["passages"]:
                parts.append(f"- {passage['content'][:200]}...")
            parts.append("")
        
        return "\n".join(parts)
