"""
Interaction Logger

Logs all SAGE interactions for training data collection and analysis.
"""

import json
import logging
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any

from sage.data_collection.schemas import (
    InteractionOutcome,
    LoggedInteraction,
    RAGContext,
    ReasoningStep,
    ToolExecution,
    VerificationInfo,
)

logger = logging.getLogger(__name__)


class InteractionLogger:
    """Logs SAGE interactions for training data collection.

    Provides thread-safe logging of interactions to JSON files,
    with support for batch export and analysis.

    Example:
        logger = InteractionLogger(storage_path="data/interactions/")
        await logger.initialize()

        # Log an interaction
        interaction = await logger.log_interaction(
            query="What is the subcritical mass limit for U-235?",
            response="Per ANSI/ANS-8.1, the limit is 20.1 kg.",
            rag_context=[...],
            tool_results=[...],
        )

        # Export for analysis
        await logger.export_to_jsonl("training_data.jsonl")
    """

    def __init__(
        self,
        storage_path: str | Path = "data/interactions/",
        enabled: bool = True,
        batch_size: int = 100,
    ):
        """Initialize interaction logger.

        Args:
            storage_path: Directory for storing interaction logs.
            enabled: Whether logging is enabled.
            batch_size: Number of interactions to batch before flushing.
        """
        self.storage_path = Path(storage_path)
        self.enabled = enabled
        self.batch_size = batch_size
        self._buffer: list[LoggedInteraction] = []
        self._initialized = False

    async def initialize(self) -> None:
        """Initialize the logger and create storage directory."""
        if self._initialized:
            return

        self.storage_path.mkdir(parents=True, exist_ok=True)
        self._initialized = True
        logger.info(f"Interaction logger initialized at {self.storage_path}")

    async def close(self) -> None:
        """Close logger and flush any remaining interactions."""
        await self._flush()
        self._initialized = False

    async def log_interaction(
        self,
        query: str,
        response: str,
        session_id: str | None = None,
        user_id: str | None = None,
        query_type: str = "knowledge",
        complexity: str = "simple",
        query_metadata: dict[str, Any] | None = None,
        rag_context: list[dict[str, Any]] | None = None,
        rag_confidence: float = 0.0,
        response_confidence: float = 0.0,
        outcome: InteractionOutcome = InteractionOutcome.SUCCESS,
        reasoning_chain: list[dict[str, Any]] | None = None,
        reasoning_classification: str = "GREEN",
        citations: list[str] | None = None,
        citation_doc_ids: list[str] | None = None,
        verification: dict[str, Any] | None = None,
        tool_executions: list[dict[str, Any]] | None = None,
        abstention_triggered: bool = False,
        abstention_reason: str | None = None,
        abstention_triggers: list[str] | None = None,
        escalation_triggered: bool = False,
        escalation_request_id: str | None = None,
        latency_ms: float = 0.0,
        token_usage: dict[str, int] | None = None,
        model_version: str = "",
    ) -> LoggedInteraction:
        """Log a single interaction.

        Args:
            query: User's query.
            response: SAGE's response.
            session_id: Session identifier.
            user_id: User identifier.
            query_type: Type of query (knowledge, calculation, analysis).
            complexity: Query complexity level.
            query_metadata: Additional query metadata.
            rag_context: Retrieved context from RAG.
            rag_confidence: RAG retrieval confidence.
            response_confidence: Response confidence score.
            outcome: Interaction outcome.
            reasoning_chain: Steps in reasoning chain.
            reasoning_classification: GREEN/YELLOW/RED classification.
            citations: List of citation strings.
            citation_doc_ids: List of cited document IDs.
            verification: Verification results.
            tool_executions: List of tool executions.
            abstention_triggered: Whether abstention was triggered.
            abstention_reason: Reason for abstention.
            abstention_triggers: Abstention trigger types.
            escalation_triggered: Whether escalation was triggered.
            escalation_request_id: Escalation request ID.
            latency_ms: Response latency in milliseconds.
            token_usage: Token usage statistics.
            model_version: Model version string.

        Returns:
            LoggedInteraction record.
        """
        if not self.enabled:
            return LoggedInteraction(
                interaction_id=str(uuid.uuid4()),
                query=query,
                response=response,
            )

        if not self._initialized:
            await self.initialize()

        # Generate interaction ID
        interaction_id = str(uuid.uuid4())

        # Parse RAG context
        parsed_rag_context = []
        if rag_context:
            for ctx in rag_context:
                parsed_rag_context.append(
                    RAGContext(
                        doc_id=ctx.get("doc_id", ""),
                        doc_title=ctx.get("doc_title", ""),
                        section=ctx.get("section"),
                        content=ctx.get("content", ""),
                        confidence=ctx.get("confidence", 0.0),
                        source=ctx.get("source"),
                    )
                )

        # Parse reasoning chain
        parsed_reasoning = []
        if reasoning_chain:
            for step in reasoning_chain:
                parsed_reasoning.append(
                    ReasoningStep(
                        step_number=step.get("step_number", 0),
                        action=step.get("action", ""),
                        result=step.get("result", ""),
                        citations=step.get("citations", []),
                        confidence=step.get("confidence", 1.0),
                    )
                )

        # Parse tool executions
        parsed_tools = []
        if tool_executions:
            for tool in tool_executions:
                parsed_tools.append(
                    ToolExecution(
                        tool_name=tool.get("tool_name", ""),
                        parameters=tool.get("parameters", {}),
                        result=tool.get("result"),
                        success=tool.get("success", True),
                        error=tool.get("error"),
                        execution_time_ms=tool.get("execution_time_ms", 0.0),
                    )
                )

        # Parse verification
        parsed_verification = None
        if verification:
            parsed_verification = VerificationInfo(
                grounded=verification.get("grounded", True),
                conservative=verification.get("conservative", True),
                citations_valid=verification.get("citations_valid", True),
                verification_checkpoints=verification.get("verification_checkpoints", []),
                warnings=verification.get("warnings", []),
            )

        # Create interaction record
        interaction = LoggedInteraction(
            interaction_id=interaction_id,
            session_id=session_id,
            user_id=user_id,
            query=query,
            query_type=query_type,
            complexity=complexity,
            query_metadata=query_metadata or {},
            rag_context=parsed_rag_context,
            rag_confidence=rag_confidence,
            response=response,
            response_confidence=response_confidence,
            outcome=outcome,
            reasoning_chain=parsed_reasoning,
            reasoning_classification=reasoning_classification,
            citations=citations or [],
            citation_doc_ids=citation_doc_ids or [],
            verification=parsed_verification,
            tool_executions=parsed_tools,
            abstention_triggered=abstention_triggered,
            abstention_reason=abstention_reason,
            abstention_triggers=abstention_triggers or [],
            escalation_triggered=escalation_triggered,
            escalation_request_id=escalation_request_id,
            latency_ms=latency_ms,
            token_usage=token_usage or {},
            timestamp=datetime.now(),
            model_version=model_version,
        )

        # Add to buffer
        self._buffer.append(interaction)

        # Flush if buffer is full
        if len(self._buffer) >= self.batch_size:
            await self._flush()

        logger.debug(f"Logged interaction {interaction_id}")
        return interaction

    async def log_from_sage_response(
        self,
        query: str,
        response: Any,
        session_id: str | None = None,
        user_id: str | None = None,
    ) -> LoggedInteraction:
        """Log interaction from a SAGEResponse object.

        Args:
            query: Original query.
            response: SAGEResponse object.
            session_id: Session identifier.
            user_id: User identifier.

        Returns:
            LoggedInteraction record.
        """
        # Extract RAG context if available
        rag_context = []
        if response.rag_response and hasattr(response.rag_response, "context"):
            for ctx in response.rag_response.context:
                rag_context.append({
                    "doc_id": getattr(ctx, "doc_id", ""),
                    "doc_title": getattr(ctx, "doc_title", ""),
                    "section": getattr(ctx, "section", None),
                    "content": getattr(ctx, "content", ""),
                    "confidence": getattr(ctx, "confidence", 0.0),
                })

        # Extract tool executions
        tool_executions = []
        for tool_result in response.tool_results:
            tool_executions.append({
                "tool_name": tool_result.tool_name,
                "parameters": getattr(tool_result, "parameters", {}),
                "result": tool_result.result,
                "success": tool_result.success,
                "error": tool_result.error,
                "execution_time_ms": tool_result.execution_time_ms,
            })

        # Determine outcome
        if response.error:
            outcome = InteractionOutcome.ERROR
        elif response.should_abstain:
            outcome = InteractionOutcome.ABSTAINED
        elif response.requires_escalation:
            outcome = InteractionOutcome.ESCALATED
        else:
            outcome = InteractionOutcome.SUCCESS

        # Extract abstention info
        abstention_triggered = response.should_abstain
        abstention_reason = None
        abstention_triggers = []
        if response.abstention and abstention_triggered:
            abstention_reason = response.abstention.reason
            abstention_triggers = [t.value for t in response.abstention.triggers]

        # Extract verification results
        verification = None
        if response.verification_results:
            verification = {
                "grounded": response.grounded,
                "conservative": True,  # Default to conservative
                "citations_valid": True,
                "verification_checkpoints": [
                    v.checkpoint for v in response.verification_results
                ],
                "warnings": [
                    w for v in response.verification_results for w in v.warnings
                ],
            }

        return await self.log_interaction(
            query=query,
            response=response.answer,
            session_id=session_id,
            user_id=user_id,
            query_type=response.query_type.value if response.query_type else "knowledge",
            complexity=response.complexity.value if response.complexity else "simple",
            rag_context=rag_context,
            rag_confidence=response.confidence,
            response_confidence=response.confidence,
            outcome=outcome,
            verification=verification,
            tool_executions=tool_executions,
            abstention_triggered=abstention_triggered,
            abstention_reason=abstention_reason,
            abstention_triggers=abstention_triggers,
            escalation_triggered=response.requires_escalation,
            escalation_request_id=response.escalation_request_id,
            latency_ms=response.latency_ms,
            model_version=response.model,
        )

    async def _flush(self) -> None:
        """Flush buffered interactions to storage."""
        if not self._buffer:
            return

        # Write to daily log file
        date_str = datetime.now().strftime("%Y-%m-%d")
        log_file = self.storage_path / f"interactions_{date_str}.jsonl"

        with open(log_file, "a") as f:
            for interaction in self._buffer:
                f.write(json.dumps(interaction.to_dict()) + "\n")

        logger.info(f"Flushed {len(self._buffer)} interactions to {log_file}")
        self._buffer.clear()

    async def get_interactions(
        self,
        start_date: datetime | None = None,
        end_date: datetime | None = None,
        query_type: str | None = None,
        outcome: InteractionOutcome | None = None,
        limit: int = 1000,
    ) -> list[LoggedInteraction]:
        """Retrieve logged interactions.

        Args:
            start_date: Start of date range.
            end_date: End of date range.
            query_type: Filter by query type.
            outcome: Filter by outcome.
            limit: Maximum number of interactions to return.

        Returns:
            List of LoggedInteraction objects.
        """
        interactions = []

        # List all log files
        log_files = sorted(self.storage_path.glob("interactions_*.jsonl"))

        for log_file in log_files:
            with open(log_file) as f:
                for line in f:
                    if len(interactions) >= limit:
                        break

                    data = json.loads(line)
                    interaction = LoggedInteraction.from_dict(data)

                    # Apply filters
                    if start_date and interaction.timestamp < start_date:
                        continue
                    if end_date and interaction.timestamp > end_date:
                        continue
                    if query_type and interaction.query_type != query_type:
                        continue
                    if outcome and interaction.outcome != outcome:
                        continue

                    interactions.append(interaction)

        return interactions

    async def get_pending_reviews(
        self,
        limit: int = 50,
    ) -> list[LoggedInteraction]:
        """Get interactions pending expert review.

        Prioritizes interactions with:
        - Low confidence scores
        - Abstention triggered
        - Escalation triggered
        - Tool usage

        Args:
            limit: Maximum number to return.

        Returns:
            List of interactions to review.
        """
        all_interactions = await self.get_interactions(limit=limit * 2)

        # Sort by priority (lower confidence first, abstentions, escalations)
        def priority(i: LoggedInteraction) -> float:
            score = i.response_confidence

            # Boost priority for special cases
            if i.abstention_triggered:
                score -= 0.3
            if i.escalation_triggered:
                score -= 0.2
            if i.tool_executions:
                score -= 0.1

            return score

        all_interactions.sort(key=priority)

        return all_interactions[:limit]

    async def export_to_jsonl(
        self,
        output_path: str | Path,
        start_date: datetime | None = None,
        end_date: datetime | None = None,
    ) -> int:
        """Export interactions to JSONL format.

        Args:
            output_path: Path to output file.
            start_date: Start of date range.
            end_date: End of date range.

        Returns:
            Number of interactions exported.
        """
        await self._flush()  # Ensure buffer is flushed

        interactions = await self.get_interactions(
            start_date=start_date,
            end_date=end_date,
            limit=100000,  # Large limit for export
        )

        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        with open(output_path, "w") as f:
            for interaction in interactions:
                f.write(json.dumps(interaction.to_dict()) + "\n")

        logger.info(f"Exported {len(interactions)} interactions to {output_path}")
        return len(interactions)

    def get_statistics(self) -> dict[str, Any]:
        """Get logging statistics.

        Returns:
            Dictionary of statistics.
        """
        stats = {
            "storage_path": str(self.storage_path),
            "enabled": self.enabled,
            "buffer_size": len(self._buffer),
            "log_files": [],
            "total_interactions": 0,
        }

        for log_file in self.storage_path.glob("interactions_*.jsonl"):
            with open(log_file) as f:
                count = sum(1 for _ in f)
            stats["log_files"].append({
                "file": log_file.name,
                "count": count,
            })
            stats["total_interactions"] += count

        return stats
