from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, List, Dict, Any, Union, Literal, TypeAlias, Protocol, runtime_checkable
import datetime as _dt

from pydantic import BaseModel, field_validator, field_serializer

try:
    # py>=3.9
    from zoneinfo import ZoneInfo
except Exception:  # pragma: no cover
    ZoneInfo = None  # type: ignore


@dataclass
class FunctionInfo:
    name: str
    arguments: Any = None

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "FunctionInfo":
        """Create ``FunctionInfo`` from a JSON-like dict."""
        return cls(
            name=data.get("name", ""),
            arguments=data.get("arguments"),
        )

    def to_dict(self) -> Dict[str, Any]:
        """Serialize to a plain dict."""
        result = {"name": self.name}
        if self.arguments is not None:
            result["arguments"] = self.arguments
        return result


@dataclass
class ToolCallInfo:
    id: str
    type: str
    function: FunctionInfo

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "ToolCallInfo":
        """Create ``ToolCallInfo`` from a JSON-like dict."""
        function_data = data.get("function", {})
        function = FunctionInfo.from_json(function_data)
        return cls(
            id=data.get("id", ""),
            type=data.get("type", ""),
            function=function,
        )

    def to_dict(self) -> Dict[str, Any]:
        """Serialize to a plain dict."""
        return {
            "id": self.id,
            "type": self.type,
            "function": self.function.to_dict(),
        }


def _format_utc_offset(dt: _dt.datetime) -> str:
    """Format UTC offset as ``+HH:MM`` or ``-HH:MM``."""
    off = dt.utcoffset()
    if off is None:
        return "+00:00"
    total = int(off.total_seconds())
    sign = "+" if total >= 0 else "-"
    total = abs(total)
    hh = total // 3600
    mm = (total % 3600) // 60
    return f"{sign}{hh:02d}:{mm:02d}"


def _get_local_aware_now() -> _dt.datetime:
    """Return current local timezone-aware datetime."""
    return _dt.datetime.now().astimezone()


def _coerce_to_local_tz(dt: _dt.datetime) -> _dt.datetime:
    """Normalize datetime to local timezone-aware representation."""
    if dt.tzinfo is None:
        # For working-memory semantics we treat naive datetimes as local time.
        local = _get_local_aware_now().tzinfo
        return dt.replace(tzinfo=local)
    return dt.astimezone(_get_local_aware_now().tzinfo)


class UserContent(BaseModel):
    """User message payload with identity and timestamp metadata."""

    user_name: str
    date: _dt.datetime
    content: str
    email: Optional[str] = None
    tool_contents: Optional[List["UserContent"]] = None
    __chronml_template__ = (
        "From: {.user_name}\n"
        "Date: {.date}\n\n"
        "{.content:block}\n"
    )

    @field_validator("date", mode="before")
    @classmethod
    def parse_date(cls, v: Any) -> _dt.datetime:
        """Parse supported date forms and normalize to local timezone-aware datetime."""
        # 1) JSON object: {"ts": ..., "tz": ...}
        if isinstance(v, dict) and "ts" in v:
            ts = v["ts"]
            tz = v.get("tz")

            if not isinstance(ts, (int, float)):
                raise TypeError(f"date.ts must be int/float seconds, got {type(ts)}")

            # Prefer IANA timezone name if available (preserves DST rules).
            if tz and isinstance(tz, str) and ZoneInfo is not None:
                try:
                    dt = _dt.datetime.fromtimestamp(ts, tz=ZoneInfo(tz))
                    return dt
                except Exception:
                    pass

            # Fallback to fixed offset (+HH:MM / -HH:MM).
            if tz and isinstance(tz, str) and (tz.startswith("+") or tz.startswith("-")):
                try:
                    sign = 1 if tz[0] == "+" else -1
                    hh, mm = tz[1:].split(":")
                    off = _dt.timedelta(hours=int(hh) * sign, minutes=int(mm) * sign)
                    dt = _dt.datetime.fromtimestamp(ts, tz=_dt.timezone(off))
                    return dt
                except Exception:
                    pass

            # Final fallback: interpret as UTC.
            return _dt.datetime.fromtimestamp(ts, tz=_dt.timezone.utc)

        # 2) Raw timestamp in seconds: interpret as UTC and convert to local tz.
        if isinstance(v, (int, float)):
            dt = _dt.datetime.fromtimestamp(v, tz=_dt.timezone.utc)
            return dt.astimezone(_get_local_aware_now().tzinfo)

        # 3) Datetime value.
        if isinstance(v, _dt.datetime):
            return _coerce_to_local_tz(v)

        # Let pydantic parse string formats (e.g. ISO8601).
        return v

    @field_validator("date", mode="after")
    @classmethod
    def normalize_date(cls, v: _dt.datetime) -> _dt.datetime:
        """Ensure in-memory date is always local timezone-aware."""
        return _coerce_to_local_tz(v)

    @field_serializer("date")
    def dump_date(self, v: _dt.datetime) -> Dict[str, Union[float, str]]:
        """Serialize date as ``{ts, tz}`` with IANA tz or numeric offset."""
        if v.tzinfo is None:
            v = _coerce_to_local_tz(v)

        ts = v.timestamp()
        tz_name = None
        if ZoneInfo is not None:
            tz_name = getattr(v.tzinfo, "key", None)

        return {
            "ts": ts,
            "tz": tz_name if isinstance(tz_name, str) and tz_name else _format_utc_offset(v),
        }

    def to_dict(self) -> Dict[str, Any]:
        """Serialize with pydantic ``model_dump``."""
        return self.model_dump()

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "UserContent":
        """Create ``UserContent`` from dict."""
        return cls(
            user_name=data.get("user_name", ""),
            email=data.get("email"),
            date=data.get("date", _dt.datetime.now()),
            content=data.get("content", ""),
            tool_contents=data.get("tool_contents"),
        )

    def to_message_text(self) -> str:
        """Render payload as plain text for model-input serialization."""
        date_str = self.date.strftime("%a, %d %b %Y %H:%M:%S %z")
        tool_prefix = ""
        if self.tool_contents:
            parts = []
            for idx, item in enumerate(self.tool_contents, 1):
                parts.append(f"[Tool Notice {idx}]\\n{item.to_message_text()}")
            tool_prefix = "\\n\\n".join(parts) + "\\n\\n"
        if self.email:
            return f"From: {self.user_name} <{self.email}>\\nDate: {date_str}\\n\\n{tool_prefix}{self.content}"
        return f"From: {self.user_name}\\nDate: {date_str}\\n\\n{tool_prefix}{self.content}"


class ToolMessageContent(BaseModel):
    """Standard tool output payload represented as a base model."""

    status: Optional[str] = None
    data: Optional[Any] = None
    error: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    info: Optional[Dict[str, Any]] = None
    __chronml_template__ = "{.status}\n{.data:code-block-json}\n{.error}\n{.metadata:code-block-json}\n{.info:code-block-json}\n"


@runtime_checkable
class MessageTextRenderable(Protocol):
    def to_message_text(self) -> str:
        """Render object as plain text for LLM conversation payloads."""


Role: TypeAlias = Literal["user", "assistant", "system", "tool"]
ROLE_USER: Role = "user"
ROLE_ASSISTANT: Role = "assistant"
ROLE_SYSTEM: Role = "system"
ROLE_TOOL: Role = "tool"


class Message(BaseModel):
    """Conversation message with optional tool and reasoning metadata."""

    role: Role
    content: Union[str, UserContent, ToolMessageContent, BaseModel]
    name: Optional[str] = None
    tool_call_id: Optional[str] = None
    tool_calls: Optional[List[ToolCallInfo]] = None
    usage: Optional[Dict[str, Any]] = None
    reasoning_content: Optional[str] = None
    model_name: Optional[str] = None
    extra: Optional[Dict[str, Any]] = None

    @property
    def __chronml_template__(self) -> str:
        if self.role == ROLE_SYSTEM:
            return ""
        if self.role == ROLE_TOOL:
            return "# [{.role}]\n\n{.content}\n"
        head = "# [{.role}]\n\n"
        if self.role == ROLE_ASSISTANT:
            return head + "{.reasoning_content:block}\n{.content:block}\n{.tool_calls:code-block-json}\n{.usage:code-block-json}\n"
        return head + "{.content}\n"


class CompactRecord(BaseModel):
    """Compaction marker in history.

    Indicates that all previous content has been compacted into ``value``.
    When building messages for the LLM, clear earlier messages and insert
    a single assistant message with this ``value``.
    """

    type: str = "compact"
    value: str = ""
    user_intent: str = ""
    __chronml_template__ = "**Compact Summary:** {.value}\n\n**User Intent:** {.user_intent}\n"
