from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, ConfigDict, Field, field_validator
from typing_extensions import TypedDict

from backend.database.models import PrincipalTypeEnum, PublishingStatusEnum, StatusEnum
from backend.schemas.generated_document import AgentHubDocumentMetadata
from backend.utils.semantic_models_utils import SemanticDatasourceConfig, SemanticDatasourceInfo


def to_camel(string: str) -> str:
    parts = string.split("_")
    return parts[0] + "".join(word.capitalize() for word in parts[1:])


class Artifact(TypedDict, total=False):
    # OTHER TYPES
    name: str
    description: str
    agentHubDocumentMetadata: AgentHubDocumentMetadata
    # Reasoning specific
    id: str
    type: Optional[str]
    hierarchy: Optional[List[Dict[str, Any]]]
    # Common
    # `llmReadableContent` is an optional internal cache used to store an LLM-readable
    # representation for some formats such as PDF and DOCX. It is lazy-populated and
    # should not be assumed to exist on every part.
    parts: List[Dict[str, Any]]
    # preview: Optional[bool] = False # True if this is a preview (not full data)


# exp reasoning artifact
# [{'id': 'rs_04bffab2219fbe3f016979f10efd4881a083bce6096a2b9a0a', 'type': 'REASONING', 'hierarchy': [{'type': 'AGENT', 'agentLoopIteration': 1, 'agentId': 'AgentHubID', 'agentName': 'AgentHub'}, {'type': 'TOOL', 'toolRef': 'AH:agent:MPqDQf5g', 'toolName': 'food expert', 'toolCallId': 'fc_019551c52cba2c56016979f10e0894819f8e4b16132e7d41c4'}, {'type': 'AGENT', 'agentId': 'MPqDQf5g', 'agentName': 'food expert', 'agentLoopIteration': 1}], 'parts': [{'type': 'TEXT', 'index': 2, 'text': '.'}]}]
# exp reasoning artifact [{'id': 'rs_0a24d1c16a7d155d00695506e9668881a0abe7c5a735455214', 'type': 'REASONING', 'parts': [{'type': 'TEXT', 'index': 0, 'text': '**Greeting'}]}]
class ReasoningArtifact(TypedDict):
    id: str
    type: Optional[str]
    hierarchy: Optional[List[Dict[str, Any]]]
    parts: List[Dict[str, Any]]


class ArtifactsMetadata(TypedDict, total=False):
    size_mb: float
    artifacts_id: str
    agentName: str
    query: str
    has_records: bool
    artifacts: Optional[List[Artifact]]  # TODO change this to non optional when data model updated
    preview: bool = False  # True if artifacts is a preview (not full data)


class FeedbackRead(BaseModel):
    rating: Optional[int] = None
    text: Optional[str] = None
    by: Optional[str] = None
    updated_at: Optional[datetime] = None


# Message Schemas
class MessageBase(BaseModel):
    role: str
    content: str
    actions: Optional[Dict[str, Any]] = None
    # artifacts: Optional[ArtifactsMetadata] = None
    selected_agent_ids: Optional[List[str]] = None
    used_agent_ids: Optional[List[str]] = None
    llm_id: Optional[str] = ""
    agents_enabled: Optional[bool] = True
    merged_context: Optional[Dict[str, Any]] = None

    status: Optional[StatusEnum] = StatusEnum.ACTIVE


class MessageCreate(MessageBase):
    id: str
    conversation_id: str
    artifacts: Optional[Dict[str, ArtifactsMetadata]] = None
    tool_validation_requests: Optional[Any] = None
    tool_validation_responses: Optional[Any] = None
    memory_fragment: Optional[Any] = None
    event_log: Optional[List[Any]] = None
    feedback_rating: Optional[int] = None
    feedback_text: Optional[str] = None
    feedback_by: Optional[str] = None
    feedback_updated_at: Optional[datetime] = None
    trace: Optional[Dict[str, Any]] = None
    attachments: Optional[List[Dict[str, Any]]] = None


class MessageUpdate(BaseModel):
    actions: Optional[Dict[str, Any]] = None
    artifacts: Optional[Dict[str, ArtifactsMetadata]] = None
    event_log: Optional[List[Any]] = None
    feedback_rating: Optional[int] = None
    feedback_text: Optional[str] = None
    feedback_by: Optional[str] = None
    feedback_updated_at: Optional[datetime] = None
    trace: Optional[Dict[str, Any]] = None
    tool_validation_responses: Optional[Any] = None
    merged_context: Optional[Dict[str, Any]] = None


class MessageRead(MessageBase):
    id: str
    created_at: datetime
    artifacts_metadata: Optional[Dict[str, ArtifactsMetadata]] = None
    tool_validation_requests: Optional[Any] = None
    tool_validation_responses: Optional[Any] = None
    memory_fragment: Optional[Any] = None
    attachments: Optional[List[Dict[str, Any]]] = None
    feedback: Optional[FeedbackRead] = None  # needed for the frontend
    has_event_log: Optional[bool] = False  # needed for the frontend
    extraction_mode: Optional[str] = None
    quota_exceeded: Optional[bool] = False

    model_config = ConfigDict(from_attributes=True, alias_generator=to_camel, populate_by_name=True)


# Add MessageAgent schema
class MessageAgent(BaseModel):
    message_id: str
    agent_id: str
    selected: int
    used: int

    model_config = ConfigDict(from_attributes=True)


# Agent Schemas
class AgentInfo(BaseModel):
    name: str = ""
    tools: Optional[List] = Field(default_factory=list)
    kb_description: Optional[str] = None
    system_prompt: Optional[str] = None
    llmid: Optional[str] = None
    semantic_datasource: Optional[SemanticDatasourceInfo] = None
    sem_models_llm: Optional[str] = None
    sem_models_interaction_mode: Optional[str] = None


class PublishedVersion(AgentInfo):
    """Published version snapshot - includes all agent fields that are snapshotted."""

    description: Optional[str] = None
    sample_questions: Optional[List] = Field(default_factory=list)
    documents: Optional[List] = Field(default_factory=list)
    datasources: Optional[List[SemanticDatasourceConfig]] = None
    indexing: Optional[Any] = None


class AgentBase(AgentInfo):
    description: Optional[str] = None
    sample_questions: Optional[List] = Field(
        default_factory=list
    )  # TODO: SQLAlchemy/Pydantic is this part of published version too?
    documents: Optional[List] = Field(default_factory=list)
    datasources: Optional[List[SemanticDatasourceConfig]] = None
    indexing: Optional[Any] = None
    published_version: Optional[PublishedVersion] = None
    published_at: Optional[datetime] = None
    publishing_status: Optional[PublishingStatusEnum] = PublishingStatusEnum.IDLE
    publishing_job_id: Optional[str] = None
    publishing_scenario_runs: Optional[List[Dict[str, Any]]] = None


class AgentCreate(AgentBase):
    id: str
    owner: str


class AgentUpdate(AgentBase):
    pass


class AgentRead(AgentCreate):
    created_at: datetime
    last_modified: datetime

    cloud_folders: Optional[Dict[str, Any]] = None
    model_config = ConfigDict(from_attributes=True)


# Conversation Schemas
class ConversationBase(BaseModel):
    user_id: str
    title: Optional[str] = ""
    agent_ids: Optional[List[str]] = Field(default_factory=list)
    llm_id: Optional[str] = ""
    agents_enabled: Optional[bool] = True
    doc_gen_enabled: bool = True
    status: Optional[StatusEnum] = StatusEnum.ACTIVE


class ConversationCreate(ConversationBase):
    conversation_id: str
    messages: Optional[List[MessageBase]] = Field(default_factory=list)


class ConversationUpdate(ConversationBase):
    pass


class ConversationRead(ConversationBase):
    conversation_id: str
    created_at: datetime = datetime.now(timezone.utc)
    last_modified: datetime = datetime.now(timezone.utc)

    model_config = ConfigDict(from_attributes=True)


class FullConversationRead(ConversationRead):
    messages: List[MessageRead] = []


# TODO: Improve as for now this is kept for backward compatibility and used for api calls
class ConversationMetadata(BaseModel):
    conversation_id: str = Field(..., alias="id", validation_alias="conversation_id")
    last_modified: datetime = Field(..., alias="lastUpdated", validation_alias="last_modified")
    agent_ids: Optional[List[str]] = Field(..., alias="agentIds", validation_alias="agent_ids", default_factory=list)
    agents_enabled: bool = Field(True, alias="modeAgents", validation_alias="agents_enabled")
    llm_id: Optional[str] = Field(None, alias="selectedLLM", validation_alias="llm_id")
    doc_gen_enabled: bool = Field(
        True,
        alias="documentGenerationEnabled",
        validation_alias="doc_gen_enabled",
    )
    title: str = ""
    model_config = ConfigDict(from_attributes=True)

    @field_validator("title", mode="before")
    @classmethod
    def title_none_to_empty(cls, v):
        return v or ""


class ConversationMetadataUpdate(BaseModel):
    agent_ids: Optional[List[str]] = Field(..., alias="agentIds", validation_alias="agent_ids", default_factory=list)
    agents_enabled: bool = Field(True, alias="modeAgents", validation_alias="agents_enabled")
    llm_id: Optional[str] = Field(None, alias="selectedLLM", validation_alias="llm_id")
    doc_gen_enabled: bool = Field(
        True,
        alias="documentGenerationEnabled",
        validation_alias="doc_gen_enabled",
    )
    title: str = ""
    model_config = ConfigDict(from_attributes=True)


# AgentShare Schemas
class AgentShareBase(BaseModel):
    principal: str
    principal_type: PrincipalTypeEnum = Field(..., alias="type", validation_alias="principal_type")


class AgentShareCreate(AgentShareBase):
    agent_id: str


class AgentShareUpdate(AgentShareBase):
    pass


class AgentShareRead(AgentShareBase):
    agent_id: str

    model_config = ConfigDict(from_attributes=True)


# Preference Schemas
class PreferenceBase(BaseModel):
    prefs: Optional[Dict[str, Any]] = Field(default_factory=dict)


class PreferenceCreate(PreferenceBase):
    user_id: str


class PreferenceUpdate(PreferenceBase):
    pass


class PreferenceRead(PreferenceBase):
    user_id: str
    created_at: datetime
    last_modified: datetime

    model_config = ConfigDict(from_attributes=True)


# DraftConversation Schemas
class DraftConversationBase(BaseModel):
    convo: Optional[Dict[str, Any]] = Field(default_factory=dict)


class DraftConversationCreate(DraftConversationBase):
    agent_id: str
    user_id: str


class DraftConversationUpdate(DraftConversationBase):
    pass


class DraftConversationRead(DraftConversationBase):
    agent_id: str
    user_id: str
    created_at: datetime
    last_modified: datetime

    model_config = ConfigDict(from_attributes=True)


ConversationRead.model_rebuild()


class DerivedDocument(BaseModel):
    id: Optional[int] = None
    conv_id: str
    user_id: str
    document_name: str
    document_path: str
    document_metadata: Dict[str, Any] = Field(default_factory=lambda: {"snapshots": []})
    created_at: Optional[datetime] = None


class MessageAttachment(BaseModel):
    attachments: List[Dict[str, Any]] = Field(default_factory=list)
    extraction_mode: Optional[str] = None
    quota_exceeded: Optional[bool] = False
