"""
Zubbl Core Types
"""

from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field


class RLMode(str, Enum):
    """RL computation mode"""

    LOCAL = "local"  # Free - runs on user's machine
    COREWEAVE = "coreweave"  # Cloud PRO - H100 powered
    B200 = "b200"  # Enterprise - NVIDIA B200


class TaskStatus(str, Enum):
    """Task execution status"""

    SUCCESS = "success"
    FAILURE = "failure"
    PARTIAL = "partial"
    TIMEOUT = "timeout"


class ValidationMethod(str, Enum):
    """How the output was validated"""

    EXACT_MATCH = "exact_match"
    SEMANTIC_SIMILARITY = "semantic_similarity"
    SCHEMA_VALIDATION = "schema_validation"
    HUMAN_FEEDBACK = "human_feedback"
    AUTOMATED = "automated"


class Step(BaseModel):
    """A single step in an agent's trajectory"""

    step_id: str = ""
    action: str = ""
    tool_name: Optional[str] = None
    tool_input: Optional[Dict[str, Any]] = None
    tool_output: Optional[Any] = None
    reasoning: Optional[str] = None
    tokens_used: int = 0
    latency_ms: int = 0
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    # Failure context (critical for learning — research shows this is
    # the highest-signal data, per Reflexion/ExpeL/ETO papers)
    error: Optional[str] = None  # Error message if this step failed
    error_type: Optional[str] = None  # Exception class name
    is_recovery_attempt: bool = False  # Was this a Zubbl auto-recovery?
    recovery_fix_type: Optional[str] = None  # What fix was applied (retry, prompt_patch, etc.)
    parent_step_id: Optional[str] = None  # Links recovery steps to the failed step


class Trajectory(BaseModel):
    """
    Complete trajectory of an agent's task execution.
    This is the core data structure Zubbl uses for learning.
    """

    trajectory_id: str
    tenant_id: str
    task_description: str
    steps: List[Step]
    status: TaskStatus
    total_tokens: int = 0
    total_latency_ms: int = 0
    metadata: Dict[str, Any] = Field(default_factory=dict)
    created_at: datetime = Field(default_factory=datetime.utcnow)
    # Validation fields (optional — backwards compatible)
    expected_output: Optional[Any] = None
    actual_output: Optional[Any] = None
    correctness_score: Optional[float] = None  # 0.0-1.0, None = not validated
    agent_confidence: Optional[float] = None  # Agent's self-reported confidence

    @property
    def is_successful(self) -> bool:
        return self.status == TaskStatus.SUCCESS


class Pattern(BaseModel):
    """
    Extracted pattern from successful trajectories.
    Used for policy updates and skill building.
    """

    pattern_id: str
    task_type: str  # Normalized task category
    action_sequence: List[str]  # Abstracted sequence of actions
    success_rate: float = 0.0
    occurrence_count: int = 1
    avg_tokens: float = 0.0
    avg_latency_ms: float = 0.0
    metadata: Dict[str, Any] = Field(default_factory=dict)


class Policy(BaseModel):
    """
    Policy recommendation returned by Zubbl.
    Tells the agent what worked before for similar tasks.
    """

    policy_id: Optional[str] = None
    task_type: str
    recommended_actions: List[str]
    confidence_score: float  # 0.0 - 1.0
    source_patterns: List[str] = Field(default_factory=list)  # Pattern IDs this came from
    rl_mode: RLMode
    metadata: Dict[str, Any] = Field(default_factory=dict)
    # Validation & health fields (optional — backwards compatible)
    is_quarantined: bool = False
    correctness_trend: List[float] = Field(default_factory=list)  # Last N scores
    negative_feedback_streak: int = 0
    last_validated_at: Optional[datetime] = None


class Feedback(BaseModel):
    """
    Human or automated feedback on a trajectory.
    Used to reinforce or penalize patterns.
    """

    feedback_id: str
    trajectory_id: str
    rating: float  # -1.0 to 1.0 (negative = bad, positive = good)
    comment: Optional[str] = None
    source: str = "human"  # "human", "automated", "benchmark"
    created_at: datetime = Field(default_factory=datetime.utcnow)


class ValidationResult(BaseModel):
    """
    Result of validating an agent's output against expected output.
    Core type for confident-but-wrong detection.
    """

    validation_id: str
    trajectory_id: str
    method: ValidationMethod
    expected_output: Optional[Any] = None
    actual_output: Optional[Any] = None
    correctness_score: float = 0.0  # 0.0 to 1.0
    confidence_score: float = 0.0  # Agent's reported confidence
    is_confidence_mismatch: bool = False  # High confidence + low correctness
    details: Dict[str, Any] = Field(default_factory=dict)
    created_at: datetime = Field(default_factory=datetime.utcnow)


class PolicyHealth(BaseModel):
    """Tracks policy health over time for quarantine decisions"""

    policy_id: str
    correctness_trend: List[float] = Field(default_factory=list)  # Last N scores
    negative_feedback_streak: int = 0
    consecutive_failures: int = 0
    is_quarantined: bool = False
    quarantined_at: Optional[datetime] = None
    last_validated_at: Optional[datetime] = None


class UsageStats(BaseModel):
    """Usage statistics for a tenant"""

    tenant_id: str
    period: str  # "day", "week", "month"
    trajectories_ingested: int = 0
    queries_made: int = 0
    rl_cycles_run: int = 0
    credits_used: int = 0
    credits_remaining: int = 0
    task_success_rate: float = 0.0
    avg_improvement: float = 0.0


# =============================================================================
# Onboarding & Profile Types
# =============================================================================


class SignupRequest(BaseModel):
    """Tenant signup request"""

    email: str
    password: str = Field(..., min_length=8, description="Minimum 8 characters")
    company_name: Optional[str] = None
    contact_name: Optional[str] = None
    plan: str = "free"
    industry: Optional[str] = None
    website: Optional[str] = None


class SignupResponse(BaseModel):
    """Tenant signup response — API key shown only once"""

    tenant_id: str
    api_key: str
    plan: str
    message: str


class LoginRequest(BaseModel):
    """Tenant login request"""

    email: str
    password: str


class LoginResponse(BaseModel):
    """Login response — returns API key for session"""

    tenant_id: str
    api_key: str
    plan: str
    email: str
    company_name: Optional[str] = None


class TenantProfile(BaseModel):
    """Full tenant profile"""

    tenant_id: str
    email: Optional[str] = None
    company_name: Optional[str] = None
    contact_name: Optional[str] = None
    industry: Optional[str] = None
    website: Optional[str] = None
    plan: str = "free"
    onboarding_status: str = "pending"
    subscription_status: str = "active"
    features_enabled: List[str] = Field(default_factory=list)
    credits_total: int = 10000
    credits_used: int = 0
    created_at: Optional[datetime] = None


class TenantProfileUpdate(BaseModel):
    """Partial tenant profile update (only non-None fields applied)"""

    company_name: Optional[str] = None
    contact_name: Optional[str] = None
    industry: Optional[str] = None
    website: Optional[str] = None


class AuditLogEntry(BaseModel):
    """Audit log entry for compliance tracking"""

    tenant_id: str
    action: str  # e.g. "tenant.signup", "key.create", "policy.quarantine"
    resource_type: str  # e.g. "tenant", "api_key", "policy", "webhook", "guardrail"
    resource_id: Optional[str] = None
    old_value: Optional[Dict[str, Any]] = None
    new_value: Optional[Dict[str, Any]] = None
    actor: str = "system"
    ip_address: Optional[str] = None
    created_at: datetime = Field(default_factory=datetime.utcnow)
