"""
Data Collection Schemas

Defines data structures for interaction logging, expert feedback,
and training data export.
"""

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any


class InteractionOutcome(Enum):
    """Outcome of an interaction."""
    SUCCESS = "success"
    ABSTAINED = "abstained"
    ESCALATED = "escalated"
    ERROR = "error"


class FeedbackRating(Enum):
    """Rating scale for expert feedback."""
    POOR = 1
    BELOW_AVERAGE = 2
    AVERAGE = 3
    GOOD = 4
    EXCELLENT = 5


@dataclass
class RAGContext:
    """Retrieved context from RAG pipeline."""
    doc_id: str
    doc_title: str
    section: str | None
    content: str
    confidence: float
    source: str | None = None


@dataclass
class ToolExecution:
    """Record of a tool execution."""
    tool_name: str
    parameters: dict[str, Any]
    result: Any | None
    success: bool
    error: str | None = None
    execution_time_ms: float = 0.0


@dataclass
class ReasoningStep:
    """Single step in a reasoning chain."""
    step_number: int
    action: str
    result: str
    citations: list[str] = field(default_factory=list)
    confidence: float = 1.0


@dataclass
class VerificationInfo:
    """Verification results for a response."""
    grounded: bool
    conservative: bool
    citations_valid: bool
    verification_checkpoints: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)


@dataclass
class LoggedInteraction:
    """Complete record of a SAGE interaction for training data collection.

    Captures all information needed to:
    1. Analyze system performance
    2. Generate training examples
    3. Support expert review
    """

    # Identifiers
    interaction_id: str
    session_id: str | None = None
    user_id: str | None = None

    # Query
    query: str = ""
    query_type: str = "knowledge"
    complexity: str = "simple"
    query_metadata: dict[str, Any] = field(default_factory=dict)

    # RAG Context
    rag_context: list[RAGContext] = field(default_factory=list)
    rag_confidence: float = 0.0

    # Response
    response: str = ""
    response_confidence: float = 0.0
    outcome: InteractionOutcome = InteractionOutcome.SUCCESS

    # Reasoning
    reasoning_chain: list[ReasoningStep] = field(default_factory=list)
    reasoning_classification: str = "GREEN"

    # Citations
    citations: list[str] = field(default_factory=list)
    citation_doc_ids: list[str] = field(default_factory=list)

    # Verification
    verification: VerificationInfo | None = None

    # Tool Usage
    tool_executions: list[ToolExecution] = field(default_factory=list)

    # Safety
    abstention_triggered: bool = False
    abstention_reason: str | None = None
    abstention_triggers: list[str] = field(default_factory=list)
    escalation_triggered: bool = False
    escalation_request_id: str | None = None

    # Performance
    latency_ms: float = 0.0
    token_usage: dict[str, int] = field(default_factory=dict)

    # Metadata
    timestamp: datetime = field(default_factory=datetime.now)
    model_version: str = ""
    sage_version: str = "1.0.0"

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "interaction_id": self.interaction_id,
            "session_id": self.session_id,
            "user_id": self.user_id,
            "query": self.query,
            "query_type": self.query_type,
            "complexity": self.complexity,
            "query_metadata": self.query_metadata,
            "rag_context": [
                {
                    "doc_id": ctx.doc_id,
                    "doc_title": ctx.doc_title,
                    "section": ctx.section,
                    "content": ctx.content,
                    "confidence": ctx.confidence,
                    "source": ctx.source,
                }
                for ctx in self.rag_context
            ],
            "rag_confidence": self.rag_confidence,
            "response": self.response,
            "response_confidence": self.response_confidence,
            "outcome": self.outcome.value,
            "reasoning_chain": [
                {
                    "step_number": step.step_number,
                    "action": step.action,
                    "result": step.result,
                    "citations": step.citations,
                    "confidence": step.confidence,
                }
                for step in self.reasoning_chain
            ],
            "reasoning_classification": self.reasoning_classification,
            "citations": self.citations,
            "citation_doc_ids": self.citation_doc_ids,
            "verification": {
                "grounded": self.verification.grounded,
                "conservative": self.verification.conservative,
                "citations_valid": self.verification.citations_valid,
                "verification_checkpoints": self.verification.verification_checkpoints,
                "warnings": self.verification.warnings,
            } if self.verification else None,
            "tool_executions": [
                {
                    "tool_name": tool.tool_name,
                    "parameters": tool.parameters,
                    "result": str(tool.result) if tool.result else None,
                    "success": tool.success,
                    "error": tool.error,
                    "execution_time_ms": tool.execution_time_ms,
                }
                for tool in self.tool_executions
            ],
            "abstention_triggered": self.abstention_triggered,
            "abstention_reason": self.abstention_reason,
            "abstention_triggers": self.abstention_triggers,
            "escalation_triggered": self.escalation_triggered,
            "escalation_request_id": self.escalation_request_id,
            "latency_ms": self.latency_ms,
            "token_usage": self.token_usage,
            "timestamp": self.timestamp.isoformat(),
            "model_version": self.model_version,
            "sage_version": self.sage_version,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "LoggedInteraction":
        """Create from dictionary."""
        # Parse RAG context
        rag_context = [
            RAGContext(
                doc_id=ctx["doc_id"],
                doc_title=ctx["doc_title"],
                section=ctx.get("section"),
                content=ctx["content"],
                confidence=ctx["confidence"],
                source=ctx.get("source"),
            )
            for ctx in data.get("rag_context", [])
        ]

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

        # Parse tool executions
        tool_executions = [
            ToolExecution(
                tool_name=tool["tool_name"],
                parameters=tool["parameters"],
                result=tool.get("result"),
                success=tool["success"],
                error=tool.get("error"),
                execution_time_ms=tool.get("execution_time_ms", 0.0),
            )
            for tool in data.get("tool_executions", [])
        ]

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

        return cls(
            interaction_id=data["interaction_id"],
            session_id=data.get("session_id"),
            user_id=data.get("user_id"),
            query=data.get("query", ""),
            query_type=data.get("query_type", "knowledge"),
            complexity=data.get("complexity", "simple"),
            query_metadata=data.get("query_metadata", {}),
            rag_context=rag_context,
            rag_confidence=data.get("rag_confidence", 0.0),
            response=data.get("response", ""),
            response_confidence=data.get("response_confidence", 0.0),
            outcome=InteractionOutcome(data.get("outcome", "success")),
            reasoning_chain=reasoning_chain,
            reasoning_classification=data.get("reasoning_classification", "GREEN"),
            citations=data.get("citations", []),
            citation_doc_ids=data.get("citation_doc_ids", []),
            verification=verification,
            tool_executions=tool_executions,
            abstention_triggered=data.get("abstention_triggered", False),
            abstention_reason=data.get("abstention_reason"),
            abstention_triggers=data.get("abstention_triggers", []),
            escalation_triggered=data.get("escalation_triggered", False),
            escalation_request_id=data.get("escalation_request_id"),
            latency_ms=data.get("latency_ms", 0.0),
            token_usage=data.get("token_usage", {}),
            timestamp=datetime.fromisoformat(data["timestamp"]) if data.get("timestamp") else datetime.now(),
            model_version=data.get("model_version", ""),
            sage_version=data.get("sage_version", "1.0.0"),
        )


@dataclass
class ExpertFeedback:
    """Expert review feedback for an interaction.

    Used to collect human evaluation for training data quality.
    """

    # Identifiers
    feedback_id: str
    interaction_id: str
    reviewer_id: str

    # Ratings (1-5 scale)
    accuracy_score: FeedbackRating = FeedbackRating.AVERAGE
    completeness_score: FeedbackRating = FeedbackRating.AVERAGE
    conservatism_score: FeedbackRating = FeedbackRating.AVERAGE
    citation_quality_score: FeedbackRating = FeedbackRating.AVERAGE

    # Corrections
    correction: str | None = None
    corrected_response: str | None = None
    additional_citations: list[str] = field(default_factory=list)

    # Flags
    is_dual_use_false_positive: bool = False
    is_dual_use_false_negative: bool = False
    requires_escalation_actually: bool = False
    contains_factual_error: bool = False

    # Comments
    reviewer_notes: str = ""

    # Metadata
    timestamp: datetime = field(default_factory=datetime.now)
    review_duration_seconds: float = 0.0

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "feedback_id": self.feedback_id,
            "interaction_id": self.interaction_id,
            "reviewer_id": self.reviewer_id,
            "accuracy_score": self.accuracy_score.value,
            "completeness_score": self.completeness_score.value,
            "conservatism_score": self.conservatism_score.value,
            "citation_quality_score": self.citation_quality_score.value,
            "correction": self.correction,
            "corrected_response": self.corrected_response,
            "additional_citations": self.additional_citations,
            "is_dual_use_false_positive": self.is_dual_use_false_positive,
            "is_dual_use_false_negative": self.is_dual_use_false_negative,
            "requires_escalation_actually": self.requires_escalation_actually,
            "contains_factual_error": self.contains_factual_error,
            "reviewer_notes": self.reviewer_notes,
            "timestamp": self.timestamp.isoformat(),
            "review_duration_seconds": self.review_duration_seconds,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ExpertFeedback":
        """Create from dictionary."""
        return cls(
            feedback_id=data["feedback_id"],
            interaction_id=data["interaction_id"],
            reviewer_id=data["reviewer_id"],
            accuracy_score=FeedbackRating(data.get("accuracy_score", 3)),
            completeness_score=FeedbackRating(data.get("completeness_score", 3)),
            conservatism_score=FeedbackRating(data.get("conservatism_score", 3)),
            citation_quality_score=FeedbackRating(data.get("citation_quality_score", 3)),
            correction=data.get("correction"),
            corrected_response=data.get("corrected_response"),
            additional_citations=data.get("additional_citations", []),
            is_dual_use_false_positive=data.get("is_dual_use_false_positive", False),
            is_dual_use_false_negative=data.get("is_dual_use_false_negative", False),
            requires_escalation_actually=data.get("requires_escalation_actually", False),
            contains_factual_error=data.get("contains_factual_error", False),
            reviewer_notes=data.get("reviewer_notes", ""),
            timestamp=datetime.fromisoformat(data["timestamp"]) if data.get("timestamp") else datetime.now(),
            review_duration_seconds=data.get("review_duration_seconds", 0.0),
        )


@dataclass
class TrainingExample:
    """Formatted training example for SFT or preference learning.

    Can be exported in various formats for different training frameworks.
    """

    # Core data
    prompt: str
    response: str

    # Context (for RAG-aware training)
    context: str | None = None

    # Preference data (for DPO/RLHF)
    preferred_response: str | None = None
    rejected_response: str | None = None

    # Quality signals
    expert_validated: bool = False
    accuracy_score: float = 0.0
    conservatism_score: float = 0.0

    # Metadata
    source_interaction_id: str | None = None
    source_feedback_id: str | None = None
    tags: list[str] = field(default_factory=list)

    def to_sft_format(self) -> dict[str, str]:
        """Convert to SFT format (prompt -> response)."""
        if self.context:
            full_prompt = f"Context:\n{self.context}\n\nQuestion: {self.prompt}"
        else:
            full_prompt = self.prompt

        return {
            "prompt": full_prompt,
            "response": self.response,
        }

    def to_preference_format(self) -> dict[str, str] | None:
        """Convert to preference format for DPO/RLHF."""
        if not self.preferred_response or not self.rejected_response:
            return None

        if self.context:
            full_prompt = f"Context:\n{self.context}\n\nQuestion: {self.prompt}"
        else:
            full_prompt = self.prompt

        return {
            "prompt": full_prompt,
            "chosen": self.preferred_response,
            "rejected": self.rejected_response,
        }

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "prompt": self.prompt,
            "response": self.response,
            "context": self.context,
            "preferred_response": self.preferred_response,
            "rejected_response": self.rejected_response,
            "expert_validated": self.expert_validated,
            "accuracy_score": self.accuracy_score,
            "conservatism_score": self.conservatism_score,
            "source_interaction_id": self.source_interaction_id,
            "source_feedback_id": self.source_feedback_id,
            "tags": self.tags,
        }
