"""OpenCode session source — reads sessions from the OpenCode SQLite database.

Implements the :class:`session_store.protocol.SessionSource` protocol by
reading from the OpenCode SQLite database at ``~/.local/share/opencode/opencode.db``.
"""

from __future__ import annotations

import json
import os
import sqlite3
from datetime import UTC, datetime
from typing import Any

from loguru import logger

from opencode_session.config import OpenCodeConfig
from opencode_session.queries import (
    COUNT_MESSAGES,
    DETECT_BY_CWD,
    DETECT_MOST_RECENT,
    GET_MESSAGES,
    GET_PARTS,
    GET_SESSION,
    LIST_SESSIONS,
)
from session_store import SessionExport, SessionInfo, SessionMessage
from session_store.errors import SessionLoadError, SessionNotFoundError


def parse_datetime(value: str | None) -> datetime:
    """Parse an ISO 8601 timestamp string into a timezone-aware datetime.

    Handles a trailing ``Z`` (UTC) suffix. Falls back to the current UTC time
    when the value is missing or cannot be parsed.

    Args:
        value: ISO timestamp string or ``None``.

    Returns:
        A timezone-aware :class:`~datetime.datetime`.
    """
    if not value:
        return datetime.now(UTC)
    text = value.strip()
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        parsed = datetime.fromisoformat(text)
    except ValueError:
        logger.warning("Failed to parse datetime {!r}, using now() as fallback", value)
        return datetime.now(UTC)
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=UTC)
    return parsed


class OpenCodeSource:
    """Session source adapter for OpenCode's SQLite-backed session store."""

    def __init__(self, config: OpenCodeConfig | None = None) -> None:
        self.config = config or OpenCodeConfig()

    def tool_name(self) -> str:
        """Return the name of the tool this adapter reads from."""
        return "opencode"

    # ------------------------------------------------------------------ #
    # Internal helpers
    # ------------------------------------------------------------------ #

    def _connect(self) -> sqlite3.Connection:
        """Open a read-only connection to the OpenCode SQLite database.

        Raises:
            SessionLoadError: If the database file is missing or cannot be opened.
        """
        db_path = self.config.db_path
        if not db_path.exists():
            raise SessionLoadError(f"OpenCode database not found at {db_path}")
        try:
            # uri + mode=ro keeps us from accidentally writing to the DB.
            uri = f"file:{db_path}?mode=ro"
            conn = sqlite3.connect(uri, uri=True)
        except sqlite3.Error as exc:
            raise SessionLoadError(f"Failed to open OpenCode database at {db_path}") from exc
        return conn

    @staticmethod
    def _load_json(raw: str, context: str) -> dict[str, Any] | None:
        """Parse a JSON blob, logging a warning and returning ``None`` on failure.

        Args:
            raw: The raw JSON string.
            context: A human-readable description used in the warning log.
        """
        try:
            parsed = json.loads(raw)
        except (json.JSONDecodeError, TypeError) as exc:
            logger.warning("Failed to parse JSON for {}: {}", context, exc)
            return None
        if not isinstance(parsed, dict):
            logger.warning("Expected JSON object for {} but got {}", context, type(parsed))
            return None
        return parsed

    # ------------------------------------------------------------------ #
    # Protocol methods
    # ------------------------------------------------------------------ #

    def list_sessions(self, limit: int = 20) -> list[SessionInfo]:
        """List sessions ordered by most recent activity.

        If the database does not exist, returns an empty list and logs a warning.
        """
        try:
            conn = self._connect()
        except SessionLoadError:
            logger.warning("OpenCode database not found; returning empty session list")
            return []

        try:
            cur = conn.execute(LIST_SESSIONS, (limit,))
            rows = cur.fetchall()
        finally:
            conn.close()

        sessions: list[SessionInfo] = []
        for row in rows:
            session_id, title, directory, time_updated, _time_created = row
            message_count = self._count_messages(session_id)
            sessions.append(
                SessionInfo(
                    tool=self.tool_name(),
                    session_id=session_id,
                    title=title or "",
                    message_count=message_count,
                    last_activity=parse_datetime(time_updated),
                    working_directory=directory or "",
                )
            )
        return sessions

    def _count_messages(self, session_id: str) -> int:
        """Count messages for a session, returning 0 on any error."""
        try:
            conn = self._connect()
        except SessionLoadError:
            return 0
        try:
            cur = conn.execute(COUNT_MESSAGES, (session_id,))
            row = cur.fetchone()
        finally:
            conn.close()
        return int(row[0]) if row else 0

    def load_session(self, session_id: str) -> SessionExport:
        """Load a full session with all messages built from text parts.

        Raises:
            SessionNotFoundError: If ``session_id`` does not exist.
            SessionLoadError: If the database cannot be opened.
        """
        conn = self._connect()
        try:
            # --- session metadata -------------------------------------------------
            cur = conn.execute(GET_SESSION, (session_id,))
            session_row = cur.fetchone()
            if session_row is None:
                raise SessionNotFoundError(f"OpenCode session {session_id!r} not found")
            (
                _sid,
                title,
                directory,
                time_updated,
                time_created,
                model_raw,
                slug,
                cost,
                tokens_input,
                tokens_output,
                tokens_reasoning,
                session_metadata_raw,
            ) = session_row

            # --- messages ---------------------------------------------------------
            cur = conn.execute(GET_MESSAGES, (session_id,))
            message_rows = cur.fetchall()
        finally:
            conn.close()

        messages: list[SessionMessage] = []
        for msg_id, data_raw, time_raw in message_rows:
            data = self._load_json(data_raw, f"message {msg_id}")
            if data is None:
                continue
            role = str(data.get("role", ""))
            parent_id = data.get("parentID")
            part_ids: list[str] = [p for p in data.get("parts", []) if isinstance(p, str)]
            content, parts_metadata = self._build_content(part_ids)
            # Build message-level metadata from tool-specific fields.
            msg_metadata: dict[str, Any] = dict(parts_metadata)
            for key in ("agent", "model", "cost", "tokens", "path"):
                if key in data:
                    msg_metadata[key] = data[key]
            messages.append(
                SessionMessage(
                    tool=self.tool_name(),
                    session_id=session_id,
                    node_id=msg_id,
                    parent_node_id=parent_id if isinstance(parent_id, str) else None,
                    message_id=msg_id,
                    role=role,
                    content=content,
                    created_at=parse_datetime(time_raw),
                    metadata=msg_metadata,
                )
            )

        # Extract model.id from the JSON model column.
        model_id: str | None = None
        if model_raw:
            model_data = self._load_json(model_raw, f"session {session_id} model")
            if model_data is not None and isinstance(model_data.get("id"), str):
                model_id = model_data["id"]

        # Build session-level metadata from tool-specific columns.
        session_metadata: dict[str, Any] = {}
        if slug is not None:
            session_metadata["slug"] = slug
        if cost is not None:
            session_metadata["cost"] = cost
        if tokens_input is not None:
            session_metadata["tokens_input"] = tokens_input
        if tokens_output is not None:
            session_metadata["tokens_output"] = tokens_output
        if tokens_reasoning is not None:
            session_metadata["tokens_reasoning"] = tokens_reasoning
        if session_metadata_raw:
            parsed = self._load_json(session_metadata_raw, f"session {session_id} metadata")
            if parsed is not None:
                session_metadata.update(parsed)

        return SessionExport(
            tool=self.tool_name(),
            session_id=session_id,
            title=title or "",
            working_directory=directory or "",
            created_at=parse_datetime(time_created) if time_created else None,
            last_activity_at=parse_datetime(time_updated) if time_updated else None,
            model=model_id,
            messages=messages,
            metadata=session_metadata,
        )

    def _build_content(self, part_ids: list[str]) -> tuple[str, dict[str, Any]]:
        """Fetch parts by UUID and concatenate the text parts' content.

        Non-text parts (reasoning, tool) are collected into a metadata dict
        rather than included in the unified ``content`` field. Missing parts
        are logged and skipped. Returns a tuple of (content, metadata) where
        metadata contains ``reasoning`` (list of reasoning texts) and
        ``tool_calls`` (list of tool part dicts).
        """
        if not part_ids:
            return "", {}

        conn = self._connect()
        try:
            placeholders = ", ".join("?" for _ in part_ids)
            query = GET_PARTS.format(placeholders=placeholders)
            cur = conn.execute(query, tuple(part_ids))
            rows = cur.fetchall()
        finally:
            conn.close()

        # Map part id -> data dict so we can preserve the order of part_ids.
        found: dict[str, dict[str, Any]] = {}
        for part_id, data_raw in rows:
            data = self._load_json(data_raw, f"part {part_id}")
            if data is None:
                continue
            found[part_id] = data

        texts: list[str] = []
        reasoning: list[str] = []
        tool_calls: list[dict[str, Any]] = []
        for pid in part_ids:
            data = found.get(pid)
            if data is None:
                logger.warning("Part {} referenced by message not found in DB", pid)
                continue
            part_type = data.get("type")
            if part_type == "text":
                text = data.get("text", "")
                if isinstance(text, str):
                    texts.append(text)
            elif part_type == "reasoning":
                text = data.get("text", "")
                if isinstance(text, str):
                    reasoning.append(text)
            elif part_type == "tool":
                tool_calls.append(data)
        metadata: dict[str, Any] = {}
        if reasoning:
            metadata["reasoning"] = reasoning
        if tool_calls:
            metadata["tool_calls"] = tool_calls
        return "".join(texts), metadata

    def detect_current_session(self) -> str | None:
        """Detect the most likely current session.

        First tries to match a session whose ``directory`` equals the current
        working directory, falling back to the most recently updated session.
        Returns ``None`` if the database is missing or empty. Never raises.
        """
        try:
            conn = self._connect()
        except SessionLoadError:
            return None

        try:
            cwd = os.getcwd()
            cur = conn.execute(DETECT_BY_CWD, (cwd,))
            row = cur.fetchone()
            if row is not None:
                return str(row[0])

            cur = conn.execute(DETECT_MOST_RECENT)
            row = cur.fetchone()
            if row is not None:
                return str(row[0])
            return None
        except sqlite3.Error as exc:
            logger.warning("Error while detecting current OpenCode session: {}", exc)
            return None
        finally:
            conn.close()
