state.py

"""Per-user state management with plain classes and CRUD helpers."""

from datetime import datetime
import uuid


class Source:
    def __init__(self, id, filename, file_type, size_mb, source_url, chunk_count, status, error_message, created_at, file_path=None):
        self.id = id
        self.filename = filename
        self.file_type = file_type  # "pdf", "pptx", "txt", "url", "youtube"
        self.size_mb = size_mb
        self.source_url = source_url
        self.chunk_count = chunk_count
        self.status = status  # "ready", "processing", "failed"
        self.error_message = error_message
        self.created_at = created_at
        self.file_path = file_path

    def to_dict(self):
        return {
            "id": self.id,
            "filename": self.filename,
            "file_type": self.file_type,
            "size_mb": self.size_mb,
            "source_url": self.source_url,
            "chunk_count": self.chunk_count,
            "status": self.status,
            "error_message": self.error_message,
            "created_at": self.created_at,
        }

    @classmethod
    def from_dict(cls, d):
        return cls(
            id=d["id"],
            filename=d["filename"],
            file_type=d["file_type"],
            size_mb=d["size_mb"],
            source_url=d["source_url"],
            chunk_count=d["chunk_count"],
            status=d["status"],
            error_message=d["error_message"],
            created_at=d["created_at"],
        )


class Message:
    def __init__(self, id, role, content, citations, created_at):
        self.id = id
        self.role = role  # "user" or "assistant"
        self.content = content
        self.citations = citations  # [{source, page, text}]
        self.created_at = created_at

    def to_dict(self):
        return {
            "id": self.id,
            "role": self.role,
            "content": self.content,
            "citations": self.citations,
            "created_at": self.created_at,
        }

    @classmethod
    def from_dict(cls, d):
        return cls(
            id=d["id"],
            role=d["role"],
            content=d["content"],
            citations=d.get("citations", []),
            created_at=d["created_at"],
        )


class Artifact:
    def __init__(self, id, type, title, content, audio_path, created_at):
        self.id = id
        self.type = type  # "conversation_summary", "document_summary", "podcast", "quiz"
        self.title = title
        self.content = content
        self.audio_path = audio_path
        self.created_at = created_at

    def to_dict(self):
        return {
            "id": self.id,
            "type": self.type,
            "title": self.title,
            "content": self.content,
            "audio_path": self.audio_path,
            "created_at": self.created_at,
        }

    @classmethod
    def from_dict(cls, d):
        return cls(
            id=d["id"],
            type=d["type"],
            title=d["title"],
            content=d.get("content", ""),
            audio_path=d.get("audio_path"),
            created_at=d["created_at"],
        )


class Notebook:
    def __init__(self, id, title, created_at, sources=None, messages=None, artifacts=None):
        self.id = id
        self.title = title
        self.created_at = created_at
        self.sources = sources if sources is not None else []
        self.messages = messages if messages is not None else []
        self.artifacts = artifacts if artifacts is not None else []

    def to_dict(self):
        return {
            "id": self.id,
            "title": self.title,
            "created_at": self.created_at,
            "sources": [s.to_dict() for s in self.sources],
            "messages": [m.to_dict() for m in self.messages],
            "artifacts": [a.to_dict() for a in self.artifacts],
        }

    @classmethod
    def from_dict(cls, d):
        return cls(
            id=d["id"],
            title=d["title"],
            created_at=d["created_at"],
            sources=[Source.from_dict(s) for s in d.get("sources", [])],
            messages=[Message.from_dict(m) for m in d.get("messages", [])],
            artifacts=[Artifact.from_dict(a) for a in d.get("artifacts", [])],
        )


class UserData:
    def __init__(self, user_id, user_name, notebooks=None, active_notebook_id=None):
        self.user_id = user_id
        self.user_name = user_name
        self.notebooks = notebooks if notebooks is not None else {}
        self.active_notebook_id = active_notebook_id

    def to_dict(self):
        return {
            "user_id": self.user_id,
            "user_name": self.user_name,
            "notebooks": {nb_id: nb.to_dict() for nb_id, nb in self.notebooks.items()},
            "active_notebook_id": self.active_notebook_id,
        }

    @classmethod
    def from_dict(cls, d):
        notebooks = {
            nb_id: Notebook.from_dict(nb_data)
            for nb_id, nb_data in d.get("notebooks", {}).items()
        }
        return cls(
            user_id=d["user_id"],
            user_name=d["user_name"],
            notebooks=notebooks,
            active_notebook_id=d.get("active_notebook_id"),
        )


def create_default_user_data(user_id, user_name):
    nb_id = str(uuid.uuid4())
    default_nb = Notebook(
        id=nb_id,
        title="My First Notebook",
        created_at=datetime.now().isoformat(),
    )
    return UserData(
        user_id=user_id,
        user_name=user_name,
        notebooks={nb_id: default_nb},
        active_notebook_id=nb_id,
    )


def get_active_notebook(state):
    if state and state.active_notebook_id and state.active_notebook_id in state.notebooks:
        return state.notebooks[state.active_notebook_id]
    return None


def create_notebook(state, title):
    nb_id = str(uuid.uuid4())
    state.notebooks[nb_id] = Notebook(
        id=nb_id,
        title=title,
        created_at=datetime.now().isoformat(),
    )
    state.active_notebook_id = nb_id
    return state


def delete_notebook(state, nb_id):
    if nb_id in state.notebooks:
        # Clean up Pinecone vectors for this notebook
        try:
            from persistence.vector_store import VectorStore
            VectorStore().delete_namespace(nb_id)
        except Exception:
            pass  # Best-effort cleanup
        del state.notebooks[nb_id]
        remaining = list(state.notebooks.keys())
        state.active_notebook_id = remaining[0] if remaining else None
    return state


def rename_notebook(state, nb_id, new_title):
    if nb_id in state.notebooks:
        state.notebooks[nb_id].title = new_title
    return state


def get_notebook_choices(state):
    """Return list of (display_label, notebook_id) for gr.Radio."""
    choices = []
    for nb_id, nb in state.notebooks.items():
        src_count = len(nb.sources)
        msg_count = len(nb.messages)
        label = nb.title
        if src_count > 0 or msg_count > 0:
            label += f"  ({src_count}s, {msg_count}m)"
        choices.append((label, nb_id))
    return choices


def get_all_artifacts(notebook, artifact_type):
    return [a for a in reversed(notebook.artifacts) if a.type == artifact_type]


def get_latest_artifact(notebook, artifact_type):
    for artifact in reversed(notebook.artifacts):
        if artifact.type == artifact_type:
            return artifact
    return None