import dataclasses
import logging
import shutil
import sys
import textwrap
from dataclasses import dataclass, field
from datetime import datetime
from typing import Literal, TypedDict
from xml.sax.saxutils import escape as xml_escape
from xml.sax.saxutils import quoteattr

import tomlkit
from dateutil.parser import isoparse
from rich.markup import escape as escape_markup
from rich.syntax import Syntax
from tomlkit._utils import escape_string
from typing_extensions import Self

from gptme.llm.models import get_default_model

from .codeblock import Codeblock
from .constants import ROLE_COLOR
from .util import console
from .util.prompt import rich_to_str
from .util.tokens import len_tokens
from .util.uri import URI, FilePath, parse_file_reference

logger = logging.getLogger(__name__)


class UsageData(TypedDict, total=False):
    """Token usage data from LLM API responses.

    Nested under ``usage`` in :class:`MessageMetadata` to mirror the structure
    returned by LLM provider APIs (Anthropic, OpenAI, etc.).
    """

    input_tokens: int
    output_tokens: int
    cache_read_tokens: int
    cache_creation_tokens: int


class MessageMetadata(TypedDict, total=False):
    """
    Metadata stored with each message.

    All fields are optional for compact storage - only non-None values are serialized.

    Token/cost fields are populated for assistant messages when telemetry is enabled.

    Token counts are nested under ``usage`` to match LLM API response structure::

        {
            "model": "claude-sonnet",
            "cost": 0.005,
            "usage": {
                "input_tokens": 100,
                "output_tokens": 50,
                "cache_read_tokens": 80,
                "cache_creation_tokens": 10,
            }
        }
    """

    model: str
    cost: float  # Cost in USD
    usage: UsageData


_TOKEN_KEYS = (
    "input_tokens",
    "output_tokens",
    "cache_read_tokens",
    "cache_creation_tokens",
)


def _migrate_metadata(meta: dict) -> MessageMetadata:
    """Migrate flat token format to nested ``usage`` format.

    Old format had token fields at the top level alongside ``model`` and ``cost``.
    New format nests them under ``usage``.  This function detects the old format
    and converts it, allowing transparent deserialization of old logs.
    """
    # Already migrated (or new format)
    if "usage" in meta or not any(k in meta for k in _TOKEN_KEYS):
        return MessageMetadata(**meta)

    usage: UsageData = {}
    remaining: dict = {}
    for k, v in meta.items():
        if k in _TOKEN_KEYS:
            usage[k] = v  # type: ignore[literal-required]
        else:
            remaining[k] = v
    if usage:
        remaining["usage"] = usage
    return MessageMetadata(**remaining)


def _format_toml_value(value: object) -> str:
    """Format a value for TOML inline table, properly escaping strings."""
    if isinstance(value, str):
        return f'"{escape_string(value)}"'
    if isinstance(value, float):
        return f"{value:.6f}"
    return str(value)


def _format_metadata_toml(metadata: MessageMetadata) -> str:
    """Format metadata as TOML inline table, handling nested ``usage``."""
    meta_items = []
    for k, v in metadata.items():
        if k == "usage" and isinstance(v, dict):
            # Format usage as a nested inline table
            usage_items = [f'"{uk}" = {_format_toml_value(uv)}' for uk, uv in v.items()]
            meta_items.append(f'"usage" = {{ {", ".join(usage_items)} }}')
        else:
            meta_items.append(f'"{k}" = {_format_toml_value(v)}')
    return f"metadata = {{ {', '.join(meta_items)} }}"


@dataclass(frozen=True, eq=False)
class Message:
    """
    A message in the assistant conversation.

    Attributes:
        role: The role of the message sender (system, user, or assistant).
        content: The content of the message.
        timestamp: The timestamp of the message.
        files: Files attached to the message, could e.g. be images for vision.
        pinned: Whether this message should be pinned to the top of the chat, and never context-trimmed.
        hide: Whether this message should be hidden from the chat output (but still be sent to the assistant).
        quiet: Whether this message should be printed on execution (will still print on resume, unlike hide).
               This is not persisted to the log file.
        metadata: Optional metadata including token usage and cost information.
    """

    role: Literal["system", "user", "assistant"]
    content: str
    timestamp: datetime = field(default_factory=datetime.now)
    files: list[FilePath] = field(default_factory=list)
    file_hashes: dict[str, str] = field(default_factory=dict)  # {filepath: hash}
    call_id: str | None = None

    pinned: bool = False
    hide: bool = False
    quiet: bool = False

    # Metadata for token usage and cost tracking
    metadata: MessageMetadata | None = None

    def __post_init__(self):
        assert isinstance(self.timestamp, datetime)

    def __repr__(self):
        content = textwrap.shorten(self.content, 20, placeholder="...")
        return f"<Message role={self.role} content={content}>"

    def __eq__(self, other):
        if not isinstance(other, Message):
            return False
        return self.role == other.role and self.content == other.content

    def concat(self, other: "Message", separator: str = "\n\n") -> "Message":
        """Concatenate two messages of the same role.

        Merges content with separator, and combines files and file_hashes.
        Preserves timestamp from the first message.

        Args:
            other: Message to concatenate with this one
            separator: String to join content with (default: "\\n\\n")

        Returns:
            New Message with merged content and files

        Raises:
            ValueError: If messages have different roles
        """
        if self.role != other.role:
            raise ValueError(
                f"Cannot concatenate messages with different roles: {self.role} vs {other.role}"
            )

        # Merge file_hashes (other's hashes take precedence for same paths)
        merged_hashes = {**self.file_hashes, **other.file_hashes}

        return self.replace(
            content=f"{self.content}{separator}{other.content}",
            files=self.files + other.files,
            file_hashes=merged_hashes,
            # Keep pinned/hide if either message has it set
            pinned=self.pinned or other.pinned,
            hide=self.hide or other.hide,
        )

    def __hash__(self):
        return hash((self.role, self.content))

    def len_tokens(self, model: str) -> int:
        return len_tokens(self, model=model)

    def replace(self, **kwargs) -> Self:
        """Replace attributes of the message."""
        return dataclasses.replace(self, **kwargs)

    def to_dict(self, keys=None) -> dict:
        """Return a dict representation of the message, serializable to JSON."""

        d: dict = {
            "role": self.role,
            "content": self.content,
            "timestamp": self.timestamp.isoformat(),
        }
        if self.files:
            # Resolve Paths to absolute paths, keep URIs as-is
            d["files"] = [
                str(f) if isinstance(f, URI) else str(f.resolve()) for f in self.files
            ]
        if self.file_hashes:
            d["file_hashes"] = self.file_hashes
        if self.pinned:
            d["pinned"] = True
        if self.hide:
            d["hide"] = True
        if self.call_id:
            d["call_id"] = self.call_id
        # Only serialize metadata if it has content (compact storage)
        if self.metadata:
            meta_dict = dict(self.metadata)
            # Ensure nested usage dict is also a plain dict (not tomlkit types)
            usage = meta_dict.get("usage")
            if isinstance(usage, dict):
                meta_dict["usage"] = dict(usage)
            d["metadata"] = meta_dict
        if keys:
            return {k: d[k] for k in keys if k in d}
        return d

    def to_xml(self) -> str:
        """Converts a message to an XML string with proper escaping."""
        # Use quoteattr for role to handle quotes and special chars safely
        # Use xml_escape for content to handle <, >, & characters
        return f"<message role={quoteattr(self.role)}>\n{xml_escape(self.content)}\n</message>"

    def format(
        self,
        oneline: bool = False,
        highlight: bool = False,
        max_length: int | None = None,
    ) -> str:
        """Format the message for display.

        Args:
            oneline: Whether to format the message as a single line
            highlight: Whether to highlight code blocks
            max_length: Maximum length of the message. If None, no truncation is applied.
                       If set, will truncate at first newline or max_length, whichever comes first.
        """
        if max_length is not None:
            first_newline = self.content.find("\n")
            max_length = (
                min(max_length, first_newline) if first_newline != -1 else max_length
            )
            content = self.content[:max_length]
            if len(content) < len(self.content):
                content += "..."
            temp_msg = self.replace(content=content)
            return format_msgs([temp_msg], oneline=True, highlight=highlight)[0]
        return format_msgs([self], oneline=oneline, highlight=highlight)[0]

    def print(self, oneline: bool = False, highlight: bool = True) -> None:
        print_msg(self, oneline=oneline, highlight=highlight)

    def to_toml(self) -> str:
        """Converts a message to a TOML string, for easy editing by hand in editor to then be parsed back."""
        flags = []
        if self.pinned:
            flags.append("pinned")
        if self.hide:
            flags.append("hide")
        flags_toml = "\n".join(f"{flag} = true" for flag in flags)
        # Use proper TOML array syntax with escaped strings (not Python repr)
        if self.files:
            escaped_files = ", ".join(f'"{escape_string(str(f))}"' for f in self.files)
            files_toml = f"files = [{escaped_files}]"
        else:
            files_toml = ""
        # Serialize file_hashes as TOML inline table with proper escaping
        if self.file_hashes:
            items = ", ".join(
                f'"{escape_string(k)}" = "{escape_string(v)}"'
                for k, v in self.file_hashes.items()
            )
            file_hashes_toml = f"file_hashes = {{ {items} }}"
        else:
            file_hashes_toml = ""
        # Serialize metadata as TOML inline table if present
        if self.metadata:
            metadata_toml = _format_metadata_toml(self.metadata)
        else:
            metadata_toml = ""
        # Serialize call_id only if present (avoid serializing "None" as string)
        call_id_toml = f'call_id = "{self.call_id}"' if self.call_id else ""
        extra = (
            flags_toml
            + "\n"
            + files_toml
            + "\n"
            + file_hashes_toml
            + "\n"
            + metadata_toml
            + "\n"
            + call_id_toml
        ).strip()

        # doublequotes need to be escaped
        # content = self.content.replace('"', '\\"')
        content = escape_string(self.content)
        content = content.replace("\\n", "\n")
        # Don't strip - preserve whitespace for data integrity

        return f'''[message]
role = "{self.role}"
content = """
{content}
"""
timestamp = "{self.timestamp.isoformat()}"
{extra}
'''

    @classmethod
    def from_toml(cls, toml: str) -> Self:
        """
        Converts a TOML string to a message.

        The string can be a single [[message]].
        """

        t = tomlkit.parse(toml)
        assert "message" in t and isinstance(t["message"], dict)
        msg: dict = t["message"]

        # Parse metadata if present (migrate old flat format if needed)
        metadata: MessageMetadata | None = None
        if msg.get("metadata"):
            metadata = _migrate_metadata(dict(msg["metadata"]))

        return cls(
            msg["role"],
            _fix_toml_content(msg["content"]),
            pinned=msg.get("pinned", False),
            hide=msg.get("hide", False),
            files=[parse_file_reference(f) for f in msg.get("files", [])],
            file_hashes=msg.get("file_hashes", {}),
            timestamp=isoparse(msg["timestamp"]),
            call_id=msg.get("call_id"),
            metadata=metadata,
        )

    def get_codeblocks(self) -> list[Codeblock]:
        """
        Get all codeblocks from the message content.
        """
        content_str = self.content

        # prepend newline to make sure we get the first codeblock
        if not content_str.startswith("\n"):
            content_str = "\n" + content_str

        # check if message contains a code block
        backtick_count = content_str.count("\n```")
        if backtick_count < 2:
            return []

        return Codeblock.iter_from_markdown(content_str)

    def cost(self, model: str | None = None, output=False) -> float:
        """Get the input cost of the message in USD."""
        from .llm.models import get_model  # noreorder

        m = get_model(model) if model else get_default_model()
        assert m, "No model specified or loaded"
        tok = len_tokens(self, f"{m.provider}/{m.model}")
        price = (m.price_output if output else m.price_input) / 1_000_000
        return tok * price


def format_msgs(
    msgs: list[Message],
    oneline: bool = False,
    highlight: bool = False,
    indent: int = 0,
) -> list[str]:
    """Formats messages for printing to the console."""
    # Import here to avoid circular import
    from .config import get_config

    outputs = []
    for msg in msgs:
        # Use configured username for user messages, otherwise capitalize the role
        if msg.role == "user":
            userprefix = get_config().user.user.name
        else:
            userprefix = msg.role.capitalize()
        if highlight:
            color = ROLE_COLOR[msg.role]
            userprefix = f"[bold {color}]{userprefix}[/bold {color}]"

        # get terminal width
        max_len = shutil.get_terminal_size().columns - len(userprefix)
        output = ""
        if oneline:
            content = msg.content.replace("\n", "\\n")
            if highlight:
                content = escape_markup(content)
            output += textwrap.shorten(content, width=max_len, placeholder="...")
            if len(output) < 20:
                output = content[:max_len] + "..."
        else:
            multiline = len(msg.content.split("\n")) > 1
            output += "\n" + indent * " " if multiline else ""
            for i, block in enumerate(msg.content.split("```")):
                if i % 2 == 0:
                    # Escape Rich markup in non-code-block content
                    if highlight:
                        block = escape_markup(block)
                    output += textwrap.indent(block, prefix=indent * " ")
                    continue
                if highlight:
                    lang = block.split("\n", 1)[0]
                    content = block.split("\n", 1)[-1]
                    fmt = "underline blue"
                    block = f"[{fmt}]{lang}\n[/{fmt}]" + rich_to_str(
                        Syntax(
                            content.rstrip().replace("[", r"\["),
                            lang,
                        )
                    )
                output += f"```{block.rstrip()}\n```"

        status_emoji = ""
        if msg.role == "system":
            first_line = msg.content.split("\n", 1)[0].lower()
            first_three_words = first_line.split()[:3]
            isSuccess = first_line.startswith(("saved", "appended")) or any(
                word in ["success", "successfully"] for word in first_three_words
            )
            isError = first_line.startswith(("error", "failed"))
            if isSuccess:
                status_emoji = "✅ "
            elif isError:
                status_emoji = "❌ "

        outputs.append(f"{userprefix}: {status_emoji}{output.rstrip()}")
    return outputs


def print_msg(
    msg: Message | list[Message],
    oneline: bool = False,
    highlight: bool = True,
    show_hidden: bool = False,
) -> None:
    """Prints the log to the console."""
    # if not tty, force highlight=False (for tests and such)
    if not sys.stdout.isatty():
        highlight = False

    msgs = msg if isinstance(msg, list) else [msg]
    msgstrs = format_msgs(msgs, highlight=highlight, oneline=oneline)
    skipped_hidden = 0
    for m, s in zip(msgs, msgstrs):
        if m.hide and not show_hidden:
            skipped_hidden += 1
            continue
        try:
            console.print(s)
        except Exception:
            # rich can throw errors, if so then print the raw message
            logger.exception("Error printing message")
            print(s)
    if skipped_hidden:
        console.print(
            f"[grey30]Skipped {skipped_hidden} hidden system messages, show with --show-hidden[/]"
        )


def msgs_to_toml(msgs: list[Message]) -> str:
    """Converts a list of messages to a TOML string, for easy editing by hand in editor to then be parsed back."""
    t = ""
    for msg in msgs:
        t += msg.to_toml().replace("[message]", "[[messages]]") + "\n\n"

    return t


def _fix_toml_content(content: str) -> str:
    """
    Remove exactly one trailing newline that TOML multiline format adds.

    TOML multiline strings (using triple quotes) add a newline before the
    closing delimiter. This function removes that artifact while preserving
    all other whitespace.
    """
    content = content.removesuffix("\n")
    return content


def toml_to_msgs(toml: str) -> list[Message]:
    """
    Converts a TOML string to a list of messages.

    The string can be a whole file with multiple [[messages]].
    """
    t = tomlkit.parse(toml)
    assert "messages" in t and isinstance(t["messages"], list)
    msgs: list[dict] = t["messages"]

    return [
        Message(
            msg["role"],
            _fix_toml_content(msg["content"]),
            pinned=msg.get("pinned", False),
            hide=msg.get("hide", False),
            timestamp=isoparse(msg["timestamp"]),
            files=[parse_file_reference(f) for f in msg.get("files", [])],
            file_hashes=dict(msg.get("file_hashes", {})),
            call_id=msg.get("call_id"),
            metadata=_migrate_metadata(dict(msg["metadata"]))
            if msg.get("metadata")
            else None,
        )
        for msg in msgs
    ]


def msgs2dicts(msgs: list[Message]) -> list[dict]:
    """Convert a list of Message objects to a list of dicts ready to pass to an LLM."""
    return [msg.to_dict(keys=["role", "content", "files", "call_id"]) for msg in msgs]
