import json
import socket
import time
from collections.abc import Iterator
from datetime import UTC, datetime
from pathlib import Path

from loguru import logger
from pydantic import BaseModel, Field
from sqlalchemy import delete, func, insert, text, update
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential

from session_hub_core.config import DatabaseConfig
from session_hub_core.exceptions import ConnectionError, SyncError
from session_hub_core.models import MessageModel, SessionModel, SyncLogModel
from session_store import SessionExport, SessionInfo, SessionMessage, SessionStore

STATE_DIR = Path.home() / ".session-hub"
HOST_STATE_FILE = STATE_DIR / "host.json"


class HostState(BaseModel):
    hostname: str
    last_synced_at: datetime | None = None


class SyncResult(BaseModel):
    """Result of a sync operation."""

    host: str
    sessions_upserted: int = 0
    sessions_skipped: int = 0
    messages_upserted: int = 0
    messages_deleted: int = 0  # messages removed from source since last sync
    duration_seconds: float = 0.0
    errors: list[str] = Field(default_factory=list)


def load_host_state(state_file: Path | None = None) -> HostState | None:
    path = state_file or HOST_STATE_FILE
    try:
        return HostState.model_validate(json.loads(path.read_text()))
    except (OSError, ValueError):
        return None


def save_host_state(
    hostname: str,
    last_synced_at: datetime | None = None,
    state_file: Path | None = None,
) -> None:
    path = state_file or HOST_STATE_FILE
    state = HostState(hostname=hostname, last_synced_at=last_synced_at or datetime.now(UTC))
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(state.model_dump_json(indent=2))


def validate_hostname(hostname: str) -> str:
    if not hostname:
        raise ValueError("Hostname must not be empty")
    if len(hostname) > 100:
        raise ValueError("Hostname must be at most 100 characters")
    if "\n" in hostname:
        raise ValueError("Hostname must not contain newlines")
    return hostname


def resolve_hostname(
    override: str | None = None,
    state_file: Path | None = None,
    detected: str | None = None,
) -> str:
    if override is not None:
        return validate_hostname(override)
    current = validate_hostname(detected or socket.gethostname())
    state = load_host_state(state_file)
    if state is None:
        return current
    if state.hostname == current:
        return current
    logger.warning(
        "Detected hostname '{}' but previous syncs used '{}'. "
        "Using '{}' to maintain continuity. Use --host to override.",
        current,
        state.hostname,
        state.hostname,
    )
    return state.hostname


def _as_utc(dt: datetime) -> datetime:
    if dt.tzinfo is None:
        return dt.replace(tzinfo=UTC)
    return dt.astimezone(UTC)


def _chunked(items: list[SessionMessage], size: int) -> Iterator[list[SessionMessage]]:
    for i in range(0, len(items), size):
        yield items[i : i + size]


class SyncEngine:
    """Sync local sessions to the central PostgreSQL database."""

    def __init__(self, config: DatabaseConfig, host: str) -> None:
        """Initialize with database config and the local hostname.

        Args:
            config: Database connection configuration.
            host: The hostname to tag synced sessions with.
        """
        self._config = config
        self.host = host
        self._batch_size = config.batch_size
        self._engine: AsyncEngine = create_async_engine(
            config.dsn,
            pool_size=config.pool_size,
            max_overflow=config.max_overflow,
            pool_pre_ping=True,
            pool_recycle=300,
            echo=config.echo,
        )

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=4),
        retry=retry_if_exception_type(ConnectionError),
    )
    async def sync_all(
        self,
        store: SessionStore | None = None,
        dry_run: bool = False,
    ) -> SyncResult:
        """Sync all sessions from all registered tools.

        Uses the default SessionStore (all three adapters) unless a
        custom store is provided. Iterates list_sessions() for each
        tool, loads each session via get_session(), and upserts into
        PostgreSQL.

        Returns:
            SyncResult with counts of sessions and messages upserted.
        """
        store = store or SessionStore()
        started = time.monotonic()
        tools = sorted({s.tool for s in store.list_sessions(limit=10000)})
        result = SyncResult(host=self.host)
        for tool in tools:
            tool_result = await self.sync_tool(tool, store=store, dry_run=dry_run)
            result.sessions_upserted += tool_result.sessions_upserted
            result.sessions_skipped += tool_result.sessions_skipped
            result.messages_upserted += tool_result.messages_upserted
            result.messages_deleted += tool_result.messages_deleted
            result.errors.extend(tool_result.errors)
        result.duration_seconds = time.monotonic() - started
        return result

    async def sync_tool(
        self,
        tool: str,
        dry_run: bool = False,
        store: SessionStore | None = None,
    ) -> SyncResult:
        """Sync sessions from a single tool only.

        Args:
            tool: Tool name ("devin", "opencode", "copilot").
        """
        store = store or SessionStore()
        return await self._sync_tool(tool, store, dry_run)

    async def sync_session(
        self,
        tool: str,
        session_id: str,
    ) -> SyncResult:
        """Sync a single session by tool + session ID.

        Loads the full session with messages and upserts it.
        """
        store = SessionStore()
        try:
            export = store.get_session(session_id, tool=tool)
        except Exception as exc:
            raise SyncError(f"Failed to load session {tool}/{session_id}: {exc}") from exc
        return await self.sync_session_export(export)

    async def sync_session_export(self, export: SessionExport) -> SyncResult:
        result = SyncResult(host=self.host)
        source_node_ids = [m.node_id for m in export.messages]
        started = time.monotonic()
        async with self._engine.begin() as conn:
            session_stmt = (
                pg_insert(SessionModel)
                .values(
                    host=self.host,
                    tool=export.tool,
                    session_id=export.session_id,
                    title=export.title,
                    working_directory=export.working_directory,
                    model_name=export.model,
                    created_at=export.created_at,
                    last_activity_at=export.last_activity_at,
                    message_count=len(export.messages),
                    hidden=False,
                    metadata_=export.metadata,
                )
                .on_conflict_do_update(
                    index_elements=["host", "tool", "session_id"],
                    set_={
                        "title": pg_insert(SessionModel).excluded.title,
                        "working_directory": pg_insert(SessionModel).excluded.working_directory,
                        "model_name": pg_insert(SessionModel).excluded.model_name,
                        "created_at": pg_insert(SessionModel).excluded.created_at,
                        "last_activity_at": pg_insert(SessionModel).excluded.last_activity_at,
                        "message_count": pg_insert(SessionModel).excluded.message_count,
                        "metadata": pg_insert(SessionModel).excluded.metadata,
                        "synced_at": func.now(),
                    },
                )
                .returning(SessionModel.id)
            )
            session_pk = (await conn.execute(session_stmt)).scalar_one()
            result.sessions_upserted = 1
            for batch in _chunked(export.messages, self._batch_size):
                message_stmt = (
                    pg_insert(MessageModel)
                    .values(
                        [
                            {
                                "session_pk": session_pk,
                                "host": self.host,
                                "tool": m.tool,
                                "session_id": m.session_id,
                                "node_id": m.node_id,
                                "parent_node_id": m.parent_node_id,
                                "message_id": m.message_id,
                                "role": m.role,
                                "content": m.content,
                                "created_at": m.created_at,
                                "metadata_": m.metadata,
                            }
                            for m in batch
                        ]
                    )
                    .on_conflict_do_update(
                        index_elements=["host", "tool", "session_id", "node_id"],
                        set_={
                            "parent_node_id": pg_insert(MessageModel).excluded.parent_node_id,
                            "message_id": pg_insert(MessageModel).excluded.message_id,
                            "role": pg_insert(MessageModel).excluded.role,
                            "content": pg_insert(MessageModel).excluded.content,
                            "created_at": pg_insert(MessageModel).excluded.created_at,
                            "metadata": pg_insert(MessageModel).excluded.metadata,
                            "synced_at": func.now(),
                        },
                    )
                )
                await conn.execute(message_stmt)
                result.messages_upserted += len(batch)
            delete_result = await conn.execute(
                delete(MessageModel).where(
                    MessageModel.session_pk == session_pk,
                    MessageModel.node_id.not_in(source_node_ids),
                )
            )
            result.messages_deleted = delete_result.rowcount or 0
        result.duration_seconds = time.monotonic() - started
        return result

    async def _sync_tool(self, tool: str, store: SessionStore, dry_run: bool) -> SyncResult:
        started = time.monotonic()
        result = SyncResult(host=self.host)
        log_id: int | None = None
        try:
            cutoff = await self._get_cutoff(tool, read_only=dry_run)
            if not dry_run:
                log_id = await self._start_sync_log(tool)
            sessions = store.list_sessions(tool=tool, limit=10000)
            now = datetime.now(UTC)
            changed: list[SessionInfo] = []
            for info in sessions:
                last_activity = _as_utc(info.last_activity)
                if last_activity > now:
                    logger.warning(
                        "Session {}/{} has last_activity in the future (clock skew?)",
                        tool,
                        info.session_id,
                    )
                if cutoff is None or last_activity >= cutoff:
                    changed.append(info)
            result.sessions_skipped = len(sessions) - len(changed)
            if dry_run:
                result.sessions_upserted = len(changed)
                result.messages_upserted = sum(i.message_count for i in changed)
                result.duration_seconds = time.monotonic() - started
                return result
            if not changed:
                if log_id is not None:
                    await self._finish_sync_log(log_id, "success")
                save_host_state(self.host)
                result.duration_seconds = time.monotonic() - started
                return result
            for info in changed:
                try:
                    export = store.get_session(info.session_id, tool=tool)
                    session_result = await self.sync_session_export(export)
                    result.sessions_upserted += session_result.sessions_upserted
                    result.messages_upserted += session_result.messages_upserted
                    result.messages_deleted += session_result.messages_deleted
                except Exception as exc:
                    logger.error("Failed to sync session {}/{}: {}", tool, info.session_id, exc)
                    result.errors.append(f"{tool}/{info.session_id}: {exc}")
            if log_id is not None:
                await self._finish_sync_log(
                    log_id,
                    "success",
                    sessions_upserted=result.sessions_upserted,
                    messages_upserted=result.messages_upserted,
                )
            save_host_state(self.host)
            result.duration_seconds = time.monotonic() - started
            return result
        except ConnectionError:
            raise
        except Exception as exc:
            logger.error("Sync failed for tool '{}': {}", tool, exc)
            result.errors.append(f"{tool}: {exc}")
            if log_id is not None:
                try:
                    await self._finish_sync_log(log_id, "failed", error=str(exc))
                except Exception:
                    logger.exception("Failed to mark sync_log entry {} as failed", log_id)
            result.duration_seconds = time.monotonic() - started
            return result

    async def _get_cutoff(self, tool: str, read_only: bool) -> datetime | None:
        try:
            async with self._engine.connect() as conn:
                if read_only:
                    await conn.execute(text("SET TRANSACTION READ ONLY"))
                row = await conn.execute(
                    text(
                        "SELECT MAX(started_at) FROM sync_log "
                        "WHERE host = :host AND tool = :tool AND status = 'success'"
                    ),
                    {"host": self.host, "tool": tool},
                )
                return row.scalar()
        except (OperationalError, OSError) as exc:
            raise ConnectionError(f"Cannot connect to PostgreSQL: {exc}") from exc

    async def _start_sync_log(self, tool: str) -> int:
        async with self._engine.begin() as conn:
            row = await conn.execute(
                insert(SyncLogModel)
                .values(host=self.host, tool=tool, status="running")
                .returning(SyncLogModel.id)
            )
            return row.scalar_one()

    async def _finish_sync_log(
        self,
        log_id: int,
        status: str,
        sessions_upserted: int = 0,
        messages_upserted: int = 0,
        error: str | None = None,
    ) -> None:
        async with self._engine.begin() as conn:
            await conn.execute(
                update(SyncLogModel)
                .where(SyncLogModel.id == log_id)
                .values(
                    status=status,
                    finished_at=func.now(),
                    sessions_upserted=sessions_upserted,
                    messages_upserted=messages_upserted,
                    error=error,
                )
            )

    async def close(self) -> None:
        await self._engine.dispose()
