import re
import uuid
from datetime import datetime
from enum import Enum
from typing import Annotated, Any, Literal

from pandas import DataFrame
from pydantic import Field, model_validator

from albert.core.base import BaseAlbertModel
from albert.core.shared.identifiers import LinkId, NotebookId, ProjectId, SynthesisId, TaskId
from albert.core.shared.models.base import BaseResource, EntityLink
from albert.exceptions import AlbertException
from albert.resources.acls import ACL


class ListBlockStyle(str, Enum):
    ORDERED = "ordered"
    UNORDERED = "unordered"


class BlockType(str, Enum):
    PARAGRAPH = "paragraph"
    LIST = "list"
    HEADER = "header"
    CHECKLIST = "checklist"
    IMAGE = "image"
    ATTACHES = "attaches"
    KETCHER = "ketcher"
    TABLE = "table"


class NotebookCopyType(str, Enum):
    TEMPLATE = "template"
    TASK = "Task"
    PROJECT = "Project"
    RESTORE_TEMPLATE = "restoreTemplate"
    GEN_TASK_TEMPLATE = "genTaskTemplate"


class BaseBlock(BaseAlbertModel):
    # id's are autogenerated by EditorJS without any particular pattern. The SDK will use uuid4.
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), alias="albertId")
    version: datetime | None = Field(default=None)


class HeaderContent(BaseAlbertModel):
    level: int = Field(ge=1)
    text: str | None


class HeaderBlock(BaseBlock):
    type: Literal[BlockType.HEADER] = Field(default=BlockType.HEADER, alias="blockType")
    content: HeaderContent


class ParagraphContent(BaseAlbertModel):
    text: str | None


class ParagraphBlock(BaseBlock):
    type: Literal[BlockType.PARAGRAPH] = Field(default=BlockType.PARAGRAPH, alias="blockType")
    content: ParagraphContent


class ChecklistItem(BaseAlbertModel):
    checked: bool
    text: str


class ChecklistContent(BaseAlbertModel):
    items: list[ChecklistItem]


class ChecklistBlock(BaseBlock):
    type: Literal[BlockType.CHECKLIST] = Field(default=BlockType.CHECKLIST, alias="blockType")
    content: ChecklistContent

    def is_checked(self, *, target_text: str) -> bool | None:
        """Get checked state of a checklist item

         Parameters
        ----------
        target_text : str
            The value/text of a checklist entry.

        Returns
        -------
        bool | None
            The checked state of the target entry identified by name.
        """
        # loop items
        for i in self.content.items:
            if i.text == target_text:
                # return check state
                return i.checked

        # return None if no match
        return


class AttachesContent(BaseAlbertModel):
    title: str | None = Field(default=None)
    namespace: str = Field(default="result")
    file_key: str | None = Field(default=None, alias="fileKey")
    format: str | None = Field(default=None, alias="mimeType")
    signed_url: str | None = Field(default=None, alias="signedURL", exclude=True, frozen=True)


class AttachesBlock(BaseBlock):
    type: Literal[BlockType.ATTACHES] = Field(default=BlockType.ATTACHES, alias="blockType")
    content: AttachesContent


class ImageContent(BaseAlbertModel):
    title: str | None = Field(default=None)
    namespace: str = Field(default="result")
    stretched: bool = Field(default=False)
    with_background: bool = Field(default=False, alias="withBackground")
    with_border: bool = Field(default=False, alias="withBorder")
    file_key: str | None = Field(default=None, alias="fileKey")
    format: str | None = Field(default=None, alias="mimeType")
    signed_url: str | None = Field(default=None, alias="signedURL", exclude=True, frozen=True)


class ImageBlock(BaseBlock):
    type: Literal[BlockType.IMAGE] = Field(default=BlockType.IMAGE, alias="blockType")
    content: ImageContent


class KetcherContent(BaseAlbertModel):
    synthesis_id: SynthesisId | None = Field(default=None, alias="synthesisId")
    name: str | None = Field(default=None)
    id: str | None = Field(default=None)
    block_id: str | None = Field(default=None, alias="blockId")
    data: str | None = Field(default=None)
    file_key: str | None = Field(default=None, alias="fileKey")
    s3_key: str | None = Field(default=None, alias="s3Key", exclude=True, frozen=True)
    png: str | None = Field(default=None, exclude=True, frozen=True)
    ketcher_url: str | None = Field(default=None, alias="ketcherUrl", exclude=True, frozen=True)


class KetcherBlock(BaseBlock):
    type: Literal[BlockType.KETCHER] = Field(default=BlockType.KETCHER, alias="blockType")
    content: KetcherContent


class TableContent(BaseAlbertModel):
    content: list[list[str | None]]
    with_headings: bool = Field(default=False, alias="withHeadings")


class TableBlock(BaseBlock):
    type: Literal[BlockType.TABLE] = Field(default=BlockType.TABLE, alias="blockType")
    content: TableContent

    def to_df(self, *, infer_header: bool = True) -> DataFrame:
        """Convert the TableBlock's content to a pd.DataFrame.

        Returns
        -------
        DataFrame
            The block's content as a pd.DataFrame.
        """

        # convert to df
        df = DataFrame(self.content.content)

        if infer_header:
            # clean df -> column name w/o formatting
            df.columns = df.iloc[0, :]
            df.columns = [re.sub(r"<.*?>", "", x) for x in df.columns]
            # discard first
            df = df.iloc[1:, :].reset_index(drop=True)

        # return df
        return df


class NotebookListItem(BaseAlbertModel):
    content: str | None
    items: list["NotebookListItem"] = Field(default_factory=list)


class BulletedListContent(BaseAlbertModel):
    items: list[NotebookListItem]
    style: Literal[ListBlockStyle.UNORDERED] = Field(default=ListBlockStyle.UNORDERED)


class NumberedListContent(BaseAlbertModel):
    items: list[NotebookListItem]
    style: Literal[ListBlockStyle.ORDERED] = Field(default=ListBlockStyle.ORDERED)


ListContent = Annotated[NumberedListContent | BulletedListContent, Field(discriminator="style")]


class ListBlock(BaseBlock):
    type: Literal[BlockType.LIST] = Field(default=BlockType.LIST, alias="blockType")
    content: ListContent


class NotebookLink(BaseAlbertModel):
    id: LinkId | None = Field(default=None)
    child: EntityLink = Field(..., alias="Child")


_NotebookBlockUnion = (
    HeaderBlock
    | ParagraphBlock
    | ChecklistBlock
    | AttachesBlock
    | ImageBlock
    | KetcherBlock
    | TableBlock
    | ListBlock
)
NotebookBlock = Annotated[_NotebookBlockUnion, Field(discriminator="type")]


class Notebook(BaseResource):
    id: NotebookId | None = Field(default=None, alias="albertId")
    name: str = Field(default="Untitled Notebook")
    parent_id: ProjectId | TaskId = Field(..., alias="parentId")
    version: datetime | None = Field(default=None)
    blocks: list[NotebookBlock] = Field(default_factory=list)
    links: list[NotebookLink] | None = Field(default=None)


NotebookContent = (
    HeaderContent
    | ParagraphContent
    | ChecklistContent
    | AttachesContent
    | ImageContent
    | KetcherContent
    | TableContent
    | BulletedListContent
    | NumberedListContent
)

allowed_notebook_contents = {
    BlockType.HEADER: HeaderContent,
    BlockType.PARAGRAPH: ParagraphContent,
    BlockType.CHECKLIST: ChecklistContent,
    BlockType.ATTACHES: AttachesContent,
    BlockType.IMAGE: ImageContent,
    BlockType.KETCHER: KetcherContent,
    BlockType.TABLE: TableContent,
    BlockType.LIST: (BulletedListContent, NumberedListContent),
}


class PutOperation(str, Enum):
    UPDATE = "update"
    DELETE = "delete"


class PutBlockDatum(BaseAlbertModel):
    id: str
    operation: PutOperation
    type: BlockType | None = Field(default=None, alias="blockType")
    content: NotebookContent | None = Field(default=None)
    previous_block_id: str | None = Field(default=None, alias="previousBlockId")

    @model_validator(mode="after")
    def content_matches_type(self) -> "PutBlockDatum":
        if self.content is None:
            return self  # skip check if there's no content

        content_type = allowed_notebook_contents.get(self.type)
        if content_type and not isinstance(self.content, content_type):
            msg = f"The content type and block type do not match. [content_type={type(self.content)}, block_type={self.type}]"
            raise AlbertException(msg)
        return self

    def model_dump(self, **kwargs) -> dict[str, Any]:
        """
        Shallow model_dump to exclude None values (None only removed from top level).
        This ensures required attrs are not removed.
        """
        base = super().model_dump(**kwargs)
        return {k: v for k, v in base.items() if v is not None}


class PutBlockPayload(BaseAlbertModel):
    data: list[PutBlockDatum]

    def model_dump(self, **kwargs) -> dict[str, Any]:
        """model_dump to ensure only top-level None attrs are removed on PutBlockDatum."""
        return {"data": [item.model_dump(**kwargs) for item in self.data]}


class NotebookCopyACL(BaseResource):
    fgclist: list[ACL] = Field(default=None)
    acl_class: str = Field(alias="class")


class NotebookCopyInfo(BaseAlbertModel):
    id: NotebookId
    parent_id: str = Field(alias="parentId")
    notebook_name: str | None = Field(default=None, alias="notebookName")
    name: str | None = Field(default=None)
    acl: NotebookCopyACL | None = Field(default=None)
