from __future__ import annotations

import re
import uuid
import warnings
from datetime import datetime
from enum import Enum
from pathlib import Path
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 (
    CustomTemplateId,
    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, ACLContainer


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")
    file_path: str | Path | None = Field(default=None, exclude=True)
    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")
    file_path: str | Path | None = Field(default=None, exclude=True)
    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")
    id: str | None = Field(default=None)
    block_id: str | None = Field(default=None, alias="blockId")
    data: str = Field(default=None, exclude=True)
    file_key: str | None = Field(default=None, alias="fileKey")
    s3_key: str | None = Field(default=None, alias="s3Key")
    png: str | None = Field(default=None, exclude=True)
    state_type: str = Field(default="project", alias="stateType")
    smiles: str | None = Field(default=None, alias="smiles", exclude=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 | CustomTemplateId = 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):
    """
    Access settings applied to a copied notebook.

    Warning
    -----
    Deprecated and will be removed in 2.0. Use ``ACLContainer`` instead.

    Attributes
    ----------
    fgclist : list[ACL]
        Specific access rules for users or teams.
    acl_class : str
        Default access class (for example, "restricted" or "confidential").
    """

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

    def __init__(self, **data):
        warnings.warn(
            "NotebookCopyACL is deprecated and will be removed in 2.0; use ACLContainer instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        super().__init__(**data)


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: ACLContainer | NotebookCopyACL | None = Field(default=None)
