from __future__ import annotations

# standard
import json
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

# third party
# custom
from .engine.types import Message, Role, ToolCall
from sunwaee.config import get_module_dir


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _uuid() -> str:
    return str(uuid.uuid4())


def _sessions_dir(workspace: str) -> Path:
    return get_module_dir(workspace, "gen") / "sessions"


### MODEL


@dataclass
class Session:
    id: str = field(default_factory=_uuid)
    title: str = ""
    workspace: str = ""
    created_at: str = field(default_factory=_now)
    updated_at: str = field(default_factory=_now)

    def to_meta(self) -> dict:
        return {
            "type": "meta",
            "id": self.id,
            "title": self.title,
            "workspace": self.workspace,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
        }

    @classmethod
    def from_meta(cls, d: dict) -> "Session":
        return cls(
            id=d["id"],
            title=d.get("title", ""),
            workspace=d.get("workspace", ""),
            created_at=d.get("created_at", _now()),
            updated_at=d.get("updated_at", _now()),
        )


### HELPERS


def _msg_to_dict(m: Message) -> dict:
    d: dict = {"type": "message", "role": m.role.value, "content": m.content}
    if m.reasoning_content:
        d["reasoning_content"] = m.reasoning_content
    if m.reasoning_signature:
        d["reasoning_signature"] = m.reasoning_signature
    if m.tool_call_id:
        d["tool_call_id"] = m.tool_call_id
    if m.tool_calls:
        d["tool_calls"] = [
            {
                "id": tc.id,
                "name": tc.name,
                "arguments": tc.arguments,
                **({"error": tc.error} if tc.error is not None else {}),
                **({"duration": tc.duration} if tc.duration else {}),
                **({"results": tc.results} if tc.results else {}),
            }
            for tc in m.tool_calls
        ]
    return d


def _dict_to_msg(d: dict) -> Message:
    tool_calls = None
    if d.get("tool_calls"):
        tool_calls = [
            ToolCall(
                id=tc["id"],
                name=tc["name"],
                arguments=tc["arguments"],
                error=tc.get("error"),
                duration=tc.get("duration", 0.0),
                results=tc.get("results", []),
            )
            for tc in d["tool_calls"]
        ]
    return Message(
        role=Role(d["role"]),
        content=d.get("content"),
        reasoning_content=d.get("reasoning_content"),
        reasoning_signature=d.get("reasoning_signature"),
        tool_call_id=d.get("tool_call_id"),
        tool_calls=tool_calls,
    )


### OPERATIONS


def _session_path(workspace: str, session_id: str) -> Path:
    return _sessions_dir(workspace) / f"{session_id}.jsonl"


def create_session(workspace: str, title: str = "") -> Session:
    session = Session(workspace=workspace, title=title)
    path = _session_path(workspace, session.id)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        f.write(json.dumps(session.to_meta()) + "\n")
    return session


def load_session(workspace: str, session_id: str) -> tuple[Session, list[Message]]:
    path = _session_path(workspace, session_id)
    if not path.exists():
        raise FileNotFoundError(f"Session not found: {session_id!r}")

    session: Session | None = None
    messages: list[Message] = []

    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:  # pragma: no cover
                continue
            entry = json.loads(line)
            if entry.get("type") == "meta":
                session = Session.from_meta(entry)
            elif entry.get("type") == "message":
                messages.append(_dict_to_msg(entry))

    if session is None:  # pragma: no cover
        raise ValueError(f"Session file missing metadata: {session_id!r}")
    return session, messages


def append_messages(workspace: str, session_id: str, messages: list[Message]) -> None:
    path = _session_path(workspace, session_id)
    with open(path, "a") as f:
        for m in messages:
            f.write(json.dumps(_msg_to_dict(m)) + "\n")


def update_session_meta(workspace: str, session: Session) -> None:
    """Rewrite the first line (metadata) of the session file."""
    path = _session_path(workspace, session.id)
    with open(path, "r") as f:
        lines = f.readlines()
    if lines:
        lines[0] = json.dumps(session.to_meta()) + "\n"
    with open(path, "w") as f:
        f.writelines(lines)


def list_sessions(workspace: str) -> list[Session]:
    sdir = _sessions_dir(workspace)
    if not sdir.exists():
        return []
    sessions = []
    for path in sorted(
        sdir.glob("*.jsonl"), key=lambda p: p.stat().st_mtime, reverse=True
    ):
        try:
            with open(path) as f:
                first_line = f.readline().strip()
            if first_line:
                meta = json.loads(first_line)
                if meta.get("type") == "meta":
                    sessions.append(Session.from_meta(meta))
        except Exception:  # pragma: no cover
            continue
    return sessions


def delete_session(workspace: str, session_id: str) -> None:
    path = _session_path(workspace, session_id)
    if not path.exists():
        raise FileNotFoundError(f"Session not found: {session_id!r}")
    path.unlink()
