from __future__ import annotations

from datetime import datetime
from typing import Annotated, Any, Literal

from pydantic import BaseModel, Field


# --- Content Blocks (discriminated union) ---


class TextBlock(BaseModel):
    type: Literal["text"] = "text"
    text: str = ""


class ThinkingBlock(BaseModel):
    type: Literal["thinking"] = "thinking"
    thinking: str = ""


class ToolUseBlock(BaseModel):
    type: Literal["tool_use"] = "tool_use"
    name: str = ""
    id: str = ""
    input: dict[str, Any] = Field(default_factory=dict)


class ToolResultBlock(BaseModel):
    type: Literal["tool_result"] = "tool_result"
    tool_use_id: str = ""
    content: str | list[Any] = ""
    is_error: bool = False

    def result_text(self) -> str:
        if isinstance(self.content, str):
            return self.content
        parts = []
        for item in self.content:
            if isinstance(item, dict) and item.get("type") == "text":
                parts.append(item.get("text", ""))
        return "\n".join(parts)


class ImageBlock(BaseModel):
    type: Literal["image"] = "image"
    source: dict[str, str] = Field(default_factory=dict)

    @property
    def media_type(self) -> str:
        return self.source.get("media_type", "")

    @property
    def data(self) -> str:
        return self.source.get("data", "")


ContentBlock = Annotated[
    TextBlock | ThinkingBlock | ToolUseBlock | ToolResultBlock | ImageBlock,
    Field(discriminator="type"),
]


# --- JSONL Record Types ---


class MessagePayload(BaseModel):
    role: str = ""
    content: str | list[ContentBlock] = ""


class UserRecord(BaseModel):
    type: Literal["user"] = "user"
    timestamp: str = ""
    uuid: str = ""
    cwd: str = ""
    gitBranch: str = ""
    message: MessagePayload = Field(default_factory=MessagePayload)


class AssistantRecord(BaseModel):
    type: Literal["assistant"] = "assistant"
    timestamp: str = ""
    uuid: str = ""
    requestId: str = ""
    message: MessagePayload = Field(default_factory=MessagePayload)


class SummaryRecord(BaseModel):
    type: Literal["summary"] = "summary"
    summary: str = ""


# --- App-level models ---


class Message(BaseModel):
    role: str
    timestamp: datetime
    content_blocks: list[TextBlock | ThinkingBlock | ToolUseBlock | ToolResultBlock | ImageBlock] = Field(
        default_factory=list
    )

    @property
    def time_display(self) -> str:
        return self.timestamp.strftime("%-I:%M %p")

    @property
    def relative_time_display(self) -> str:
        now = datetime.now(tz=self.timestamp.tzinfo)
        delta = now - self.timestamp
        seconds = int(delta.total_seconds())
        if seconds < 60:
            return "just now"
        minutes = seconds // 60
        if minutes < 60:
            return f"{minutes}m ago"
        hours = minutes // 60
        if hours < 24:
            return f"{hours}h ago"
        days = hours // 24
        return f"{days}d ago"


class SessionMetadata(BaseModel):
    session_id: str
    project_dir: str
    project_path: str
    file_path: str
    summary: str = ""
    first_message_preview: str = ""
    timestamp: datetime | None = None
    last_timestamp: datetime | None = None
    message_count: int = 0
    git_branch: str = ""
    project_short_name: str = ""
    total_tokens: int = 0
    sparkline_buckets: list[int] = Field(default_factory=list)

    @property
    def token_display(self) -> str:
        if self.total_tokens == 0:
            return ""
        if self.total_tokens >= 1_000_000:
            val = self.total_tokens / 1_000_000
            return f"{val:.1f}M tokens"
        k = self.total_tokens / 1000
        return f"{k:.0f}k tokens"

    @property
    def time_display(self) -> str:
        if not self.timestamp:
            return ""
        return self.timestamp.strftime("%-I:%M %p")

    @property
    def duration_display(self) -> str:
        if not self.timestamp or not self.last_timestamp:
            return ""
        total_seconds = int((self.last_timestamp - self.timestamp).total_seconds())
        if total_seconds < 60:
            return f"{total_seconds}s"
        minutes = total_seconds // 60
        if minutes < 60:
            return f"{minutes}m"
        hours = minutes // 60
        remaining = minutes % 60
        return f"{hours}h {remaining}m"

    @property
    def timestamp_iso(self) -> str:
        if not self.timestamp:
            return ""
        return self.timestamp.isoformat()


class Session(BaseModel):
    metadata: SessionMetadata
    messages: list[Message] = Field(default_factory=list)


class ProjectInfo(BaseModel):
    project_dir: str
    project_path: str
    sessions: list[SessionMetadata] = Field(default_factory=list)
