"""Pathway Store — persistence via StudioStore.

Handles saving/loading pathways and their revisions to durable storage.
"""

from __future__ import annotations

import logging
import uuid
from typing import TYPE_CHECKING, Any

from core.common import JsonObject
from core.pathways.loader import pathway_from_dict
from core.pathways.pathway import Pathway
from core.pathways.registry import PathwayMeta, PathwayRegistry

if TYPE_CHECKING:
    from persistence import StudioStore

logger = logging.getLogger(__name__)


def serialize_node(node: Any) -> JsonObject:
    """Serialize a node to a dictionary for persistence."""
    result: JsonObject = {
        "id": node.id,
        "type": node.type,
        "name": getattr(node, "name", None),
        "description": getattr(node, "description", None),
        "config": {},
    }

    schemas: dict[str, tuple[list[tuple[str, Any]], list[str]]] = {
        "llm": (
            [("prompt", ""), ("model", "reasoning"), ("temperature", 0.7)],
            ["max_tokens", "system", "response_format", "json_schema", "images"],
        ),
        "tool": ([("tool", ""), ("args", {})], []),
        "transform": ([("expr", "input")], []),
        "router": (
            [("condition", "true"), ("routes", {}), ("default", None)],
            [],
        ),
        "gate": (
            [("condition", "true"), ("true_path", ""), ("false_path", "")],
            [],
        ),
        "code": (
            [("code", ""), ("language", "python")],
            ["profile", "allow_site_packages", "timeout_ms", "memory_mb"],
        ),
        "input": (
            [("input_schema", {}), ("defaults", {})],
            ["description"],
        ),
        "output": (
            [("output_schema", {}), ("mapping", {})],
            ["description"],
        ),
        "handoff": (
            [("candidates", []), ("message", "{{input.message}}")],
            ["fallback", "include_context"],
        ),
    }

    schema = schemas.get(node.type)
    if schema:
        required, optional = schema
        for attr, default in required:
            result["config"][attr] = getattr(node, attr, default)
        for attr in optional:
            value = getattr(node, attr, None)
            if value is not None:
                result["config"][attr] = value

    return result


class PathwayStore:
    """Persistent storage for pathways via StudioStore."""

    def __init__(
        self,
        *,
        store: StudioStore | None = None,
        default_workspace_id: str | None = None,
    ) -> None:
        self._store = store
        self._default_workspace_id = default_workspace_id

    @property
    def has_store(self) -> bool:
        return self._store is not None

    def persist_pathway(
        self,
        pathway: Pathway,
        *,
        source: str,
        version: str | None = None,
        workspace_id: str | None = None,
    ) -> str | None:
        """Persist a pathway to storage. Returns doc_id or None."""
        if not self._store:
            return None

        ws_id = workspace_id or self._default_workspace_id
        if not ws_id:
            logger.warning("Cannot persist pathway: no workspace_id")
            return None

        try:
            ws = self._store.get_workspace(ws_id)
            if not ws:
                self._store.upsert_workspace(
                    {"id": ws_id, "name": f"Pathways ({ws_id})"}
                )
        except Exception as e:
            logger.warning("Failed to ensure workspace: %s", e)

        doc_id = f"doc_{uuid.uuid4().hex[:12]}"

        try:
            self._store.upsert_document(
                {
                    "id": doc_id,
                    "type": "pathway",
                    "name": pathway.name or pathway.id,
                    "workspace_id": ws_id,
                    "parent_id": None,
                    "metadata": {
                        "pathway_id": pathway.id,
                        "source": source,
                        "version": version,
                    },
                    "head_rev": None,
                }
            )

            content = {
                "id": pathway.id,
                "name": pathway.name,
                "description": pathway.description,
                "source": source,
                "version": version,
                "nodes": [serialize_node(node) for node in pathway.nodes.values()],
                "connections": [
                    {"from": c.from_node, "to": c.to_node}
                    for c in pathway.connections
                ],
                "metadata": dict(pathway.metadata),
            }

            rev_id = f"rev_{uuid.uuid4().hex[:12]}"
            self._store.write_revision(doc_id=doc_id, rev_id=rev_id, content=content)
            self._store.upsert_document({"id": doc_id, "head_rev": rev_id})

            return doc_id

        except Exception as e:
            logger.exception("Failed to persist pathway: %s", e)
            return None

    def persist_changes(
        self,
        pathway: Pathway,
        meta: PathwayMeta | None,
    ) -> str | None:
        """Persist pathway changes (update existing or create new doc)."""
        if not self._store:
            return None

        doc_id = meta.doc_id if meta else None

        if not doc_id:
            ws_id = self._default_workspace_id
            if ws_id:
                return self.persist_pathway(
                    pathway,
                    source=meta.source if meta else "user:api",
                    workspace_id=ws_id,
                )
            return None

        try:
            content = {
                "id": pathway.id,
                "name": pathway.name,
                "description": pathway.description,
                "source": meta.source if meta else "user:api",
                "version": meta.version if meta else None,
                "nodes": [serialize_node(node) for node in pathway.nodes.values()],
                "connections": [
                    {
                        "from": c.from_node,
                        "to": c.to_node,
                        "from_output": c.from_output,
                        "to_input": c.to_input,
                    }
                    for c in pathway.connections
                ],
                "metadata": dict(pathway.metadata),
            }

            rev_id = f"rev_{uuid.uuid4().hex[:12]}"
            self._store.write_revision(doc_id=doc_id, rev_id=rev_id, content=content)
            self._store.upsert_document(
                {
                    "id": doc_id,
                    "head_rev": rev_id,
                    "name": pathway.name or pathway.id,
                }
            )
            return doc_id

        except Exception as e:
            logger.exception("Failed to persist pathway changes: %s", e)
            return None

    def load_from_storage(self, pathway_id: str) -> Pathway | None:
        """Load a pathway from storage by doc_id or pathway_id."""
        if not self._store:
            return None

        doc = self._store.get_document(pathway_id)
        if doc:
            return self._load_doc_pathway(pathway_id)

        try:
            all_docs = self._store.list_documents(workspace_id=None)
            for doc in all_docs:
                if doc.type != "pathway":
                    continue
                meta = doc.metadata or {}
                if meta.get("pathway_id") == pathway_id:
                    return self._load_doc_pathway(doc.id)
        except Exception as e:
            logger.debug("Failed to search pathways: %s", e)

        return None

    def _load_doc_pathway(self, doc_id: str) -> Pathway | None:
        """Load a pathway from a document."""
        if not self._store:
            return None

        try:
            doc = self._store.get_document(doc_id)
            if not doc:
                return None

            head_rev = doc.head_rev
            if not head_rev:
                return None

            rev = self._store.get_revision(doc_id=doc_id, rev_id=head_rev)
            if not rev:
                return None

            return pathway_from_dict(rev.content or {})

        except Exception as e:
            logger.exception("Failed to load pathway %s: %s", doc_id, e)
            return None

    def delete_document(self, doc_id: str) -> None:
        """Delete a document from storage."""
        if not self._store:
            return

        try:
            delete_fn = getattr(self._store, "delete_document", None)
            if callable(delete_fn):
                delete_fn(doc_id)
        except Exception as e:
            logger.warning("Failed to delete pathway document: %s", e)

    def create_user_pathway(
        self,
        pathway: Pathway,
        *,
        source: str = "user:unknown",
        workspace_id: str | None = None,
        registry: PathwayRegistry,
    ) -> PathwayMeta:
        """Create a user pathway, persist it, and register it."""
        pid = pathway.id

        doc_id = None
        if self._store:
            ws_id = workspace_id or self._default_workspace_id
            if ws_id:
                doc_id = self.persist_pathway(
                    pathway,
                    source=source,
                    workspace_id=ws_id,
                )

        registry.put(pid, pathway)

        if doc_id:
            registry.put(doc_id, pathway)

        meta = PathwayMeta(
            id=pid,
            name=pathway.name,
            description=pathway.description,
            source=source,
            version=None,
            doc_id=doc_id,
            node_count=len(pathway.nodes),
        )
        registry.set_meta(pid, meta)
        if doc_id:
            registry.set_meta(doc_id, meta)

        logger.debug("Created pathway: %s (source=%s, doc_id=%s)", pid, source, doc_id)
        return meta

    def export(
        self,
        pathway: Pathway,
        meta: PathwayMeta | None,
    ) -> JsonObject:
        """Export a pathway as a portable JSON document."""
        return {
            "format": "server.pathway.v1",
            "pathway": {
                "id": pathway.id,
                "name": pathway.name,
                "description": pathway.description,
                "nodes": [serialize_node(node) for node in pathway.nodes.values()],
                "connections": [
                    {"from": c.from_node, "to": c.to_node}
                    for c in pathway.connections
                ],
                "metadata": dict(pathway.metadata),
            },
            "meta": meta.model_dump() if meta else None,
        }


__all__ = [
    "PathwayStore",
    "serialize_node",
]
