import base64
import json
import time
from datetime import datetime
from enum import Enum
from typing import Literal, List, Any, Dict, Optional

from pydantic import BaseModel, Field, ConfigDict
from pydantic.alias_generators import to_camel

from persona_models.common import create_uuid


class SourceChunk(BaseModel):
    start: int
    end: int


class Source(BaseModel):
    type: Literal["web", "document"]
    url: Optional[str] = None
    document_id: Optional[str] = None
    page_id: Optional[str] = None
    page_number: Optional[int] = None
    chunks: list[SourceChunk] = Field(default_factory=list)
    most_relevant: bool = False

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


class FunctionResponse(BaseModel):
    name: str
    result: Dict[str, Any]
    function_call_id: str | None = None

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


class FunctionCall(BaseModel):
    name: str
    args: Dict[str, Any]
    id: str | None = None

    def __format__(self, format_spec):
        return self.name

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


class File(BaseModel):
    name: str
    url: str | None = None
    base64: str | None = None
    content_type: str

    def is_base64(self) -> bool:
        return self.base64 is not None

    def is_url(self) -> bool:
        return self.url is not None and self.url.startswith("http")

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


class Image(BaseModel):
    content: str | None = None
    content_type: str

    def to_bytes(self) -> bytes:
        if self.content is None:
            raise ValueError("Content is none")
        return base64.b64decode(self.content)

    @classmethod
    def from_bytes(cls, image_bytes: bytes, content_type: str) -> "Image":
        return cls(
            content=base64.b64encode(image_bytes).decode(), content_type=content_type
        )

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


MessageType = Literal[
    "text", "image", "file", "json", "function_call", "function_response", "thought"
]


class Message(BaseModel):
    id: str = Field(default_factory=create_uuid)
    session_id: str | None = None
    text: str | None = None
    image: Image | None = None
    file: File | None = None
    role: Literal["user", "assistant", "function", "tool"]
    timestamp: float = Field(default_factory=time.time)
    created_at: datetime = Field(default_factory=datetime.now)
    function_calls: List[FunctionCall] | None = None
    function_response: FunctionResponse | None = None
    finish_reason: Literal["stop", "function_call"] | None = None
    sources: list | None = None
    sender: str | None = None
    thought: bool = False

    @property
    def type(self) -> MessageType:
        if self.is_function_call():
            return "function_call"
        elif self.is_function_response():
            return "function_response"
        elif self.is_image():
            return "image"
        elif self.is_file():
            return "file"
        elif self.is_json():
            return "json"
        elif self.is_thought():
            return "thought"
        elif self.is_text():
            return "text"
        else:
            raise ValueError("Message has no valid content type")

    def is_text(self):
        return (
            not self.is_function_call()
            and not self.is_function_response()
            and not self.is_image()
            and not self.is_json()
            and not self.is_file()
            and not self.is_thought()
            and self.text is not None
        )

    def is_image(self):
        return (
            not self.is_function_call()
            and not self.is_function_response()
            and self.image is not None
        )

    def is_file(self):
        return (
            not self.is_function_call()
            and not self.is_function_response()
            and self.file is not None
        )

    def is_json(self):
        return (
            not self.is_function_call()
            and not self.is_function_response()
            and self.text is not None
            and self.text.startswith("{")
            and self.text.endswith("}")
        )

    def is_function_call(self):
        return self.function_calls is not None and len(self.function_calls) > 0

    def is_function_response(self):
        return self.function_response is not None

    def is_thought(self):
        return self.thought

    def merge_text(self, other: "Message") -> bool:
        if (
            (self.is_text() or self.is_thought())
            and (self.type == other.type)
            and self.finish_reason is None
        ):
            self.text = (
                other.text
                if self.text is None
                else self.text + other.text if other.text else self.text
            )
            self.finish_reason = other.finish_reason
            self.session_id = other.session_id
            self.role = other.role
            self.timestamp = other.timestamp
            self.created_at = other.created_at
            self.file = other.file
            self.image = other.image
            self.sources = other.sources
            self.sender = other.sender

        return self.finish_reason is not None

    @staticmethod
    def render_text(content: "Message | List[Message]") -> str:
        if isinstance(content, list):
            return "\n".join(
                [
                    f"{message.role}> {message.text}"
                    for message in content
                    if message.is_text()
                ]
            )
        elif isinstance(content, Message) and content.is_text():
            return content.text

        raise ValueError("Content is not a text message")

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


class EmptyMessage(Message):
    def __init__(self):
        super().__init__(text=None, role="user", finish_reason=None)

    def is_text(self):
        return False

    def is_function_call(self):
        return False

    def is_function_response(self):
        return False


SessionStatus = Literal["active", "finished"]
SessionType = Literal["realtime", "async"]
SessionVisibility = Literal["active", "archived", "deleted"]


class Session(BaseModel):
    id: str = Field(default_factory=create_uuid)
    code: str = Field(default_factory=create_uuid)
    name: str | None = None
    project_id: str
    agent_id: str
    current_agent_id: Optional[str] = None
    user_id: str = "anonymous"
    visibility: SessionVisibility | None = "active"
    status: SessionStatus = "active"
    type: SessionType = "async"
    created_at: datetime = Field(default_factory=datetime.now)
    finished_at: datetime | None = None
    parent_session_id: Optional[str] = None
    protocol: str | None = None
    extras: Dict[str, Any] = Field(default_factory=dict)

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


class MessageAccumulator:
    _current: Message | None = None

    def accumulate(self, other: Message) -> Message:
        if not self._current:
            self._current = other.model_copy()
            return self._current

        if self._current.role != other.role:
            if not self._current.finish_reason:
                self._current.finish_reason = "stop"
            return self._current

        if self._current.merge_text(other):
            return self._current

        return other


def merge_messages(messages: List[Message]) -> List[Message]:
    merged = []
    last_message: Message | None = None
    for message in messages:
        if not last_message and message.finish_reason is not None:
            merged.append(message)
            continue

        if not last_message:
            last_message = message.model_copy()
            continue

        if last_message.role != message.role or last_message.thought != message.thought:
            if not last_message.finish_reason:
                last_message.finish_reason = "stop"

            merged.append(last_message)

            if message.finish_reason:
                merged.append(message)
                last_message = None
                continue
            else:
                last_message = message
                continue

        if last_message.merge_text(message):
            merged.append(last_message)
            last_message = None

    if last_message and not any(m.id == last_message.id for m in merged):
        merged.append(last_message)

    return [message for message in merged if not message.is_text() or message.text]


class AgentContext(BaseModel):
    data: dict[str, Any] = Field(default_factory=dict)
    files: Optional[List[File]] = None

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


class CollaborationMode(str, Enum):
    PERSISTENT_TRANSFER = "persistent_transfer"
    DELEGATE_AND_RETURN = "delegate_and_return"
    CONSULT_THEN_RESPOND = "consult_then_respond"


class TransferAction(BaseModel):
    request: Optional[str] = None
    thought: Optional[str] = None
    agent: str
    function_call_message: Message
    function_response_message: Message
    collaboration_mode: CollaborationMode
    instructions: Optional[str] = None

    def set_response(self, messages_from_collaborator: List[Message]):
        response = Message.render_text(messages_from_collaborator)
        self.function_response_message.function_response.result = {"response": response}
        return self.function_response_message

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


class AgentResponse(BaseModel):
    messages: List[Message]
    transfer: Optional[TransferAction] = None

    @property
    def text(self) -> str | None:
        return "\n".join(m.text for m in self.messages if m.is_text())

    def is_json(self) -> bool:
        if not self.messages:
            return False
        return self.messages[-1].is_json()

    def to_json(self) -> Optional[dict]:
        json_responses = []
        for m in self.messages:
            if not m.is_json():
                continue
            try:
                json_responses.append(json.loads(m.text))
            except json.JSONDecodeError:
                pass

        return json_responses[0] if len(json_responses) == 1 else json_responses

    @property
    def first_function_call(self) -> Optional[FunctionCall]:
        for m in self.messages:
            if m.is_function_call():
                return m.function_calls[-1]
        return None

    @property
    def function_calls(self) -> List[FunctionCall]:
        result = []
        for m in self.messages:
            if m.is_function_call():
                for fc in m.function_calls:
                    result.append(fc)
        return result

    def is_transfer(self):
        return self.transfer is not None

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


class GenerationResponse(BaseModel):
    json_: dict | None = Field(default=None, alias="json")
    response: AgentResponse
    session_id: str

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


class Reasoning(BaseModel):
    thought: str | None = None
    image_url: str | None = None

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