session_models.py

"""
Session Data Models

Defines dataclasses for session metadata, state, and job tracking.
"""

from dataclasses import dataclass, asdict, field
from datetime import datetime
from typing import Any, Dict, Optional, List
import json


@dataclass
class JobReference:
    """Reference to a submitted cloud job."""
    job_id: str
    service_type: str  # "qiskit_ibm", "ionq", "local_aer", etc.
    status: str = "submitted"  # submitted, running, completed, failed
    created_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
    completed_at: Optional[str] = None
    result_data: Optional[Dict[str, Any]] = None
    
    def to_dict(self) -> Dict:
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict) -> "JobReference":
        return cls(**data)


@dataclass
class SessionMetadata:
    """Metadata about a session."""
    session_id: str
    user_id: str
    alias: str
    app_type: str  # "EM" or "QLBM"
    created_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
    last_modified: str = field(default_factory=lambda: datetime.utcnow().isoformat())
    last_accessed: str = field(default_factory=lambda: datetime.utcnow().isoformat())
    description: str = ""
    
    def to_dict(self) -> Dict:
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict) -> "SessionMetadata":
        return cls(**data)
    
    def update_timestamp(self) -> None:
        """Update last modified and accessed timestamps."""
        now = datetime.utcnow().isoformat()
        self.last_modified = now
        self.last_accessed = now


@dataclass
class SessionState:
    """Complete session state for EM or QLBM app."""
    session_id: str
    app_type: str  # "EM" or "QLBM"
    
    # Generic state container
    state_data: Dict[str, Any] = field(default_factory=dict)
    
    # Job tracking
    submitted_jobs: List[JobReference] = field(default_factory=list)
    
    # Timestamps
    created_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
    last_modified: str = field(default_factory=lambda: datetime.utcnow().isoformat())
    
    def to_dict(self) -> Dict:
        return {
            "session_id": self.session_id,
            "app_type": self.app_type,
            "state_data": self.state_data,
            "submitted_jobs": [job.to_dict() for job in self.submitted_jobs],
            "created_at": self.created_at,
            "last_modified": self.last_modified,
        }
    
    @classmethod
    def from_dict(cls, data: Dict) -> "SessionState":
        jobs = [
            JobReference.from_dict(job) 
            for job in data.get("submitted_jobs", [])
        ]
        return cls(
            session_id=data["session_id"],
            app_type=data["app_type"],
            state_data=data.get("state_data", {}),
            submitted_jobs=jobs,
            created_at=data.get("created_at", datetime.utcnow().isoformat()),
            last_modified=data.get("last_modified", datetime.utcnow().isoformat()),
        )
    
    def update_timestamp(self) -> None:
        """Update last modified timestamp."""
        self.last_modified = datetime.utcnow().isoformat()
    
    def add_job(self, job: JobReference) -> None:
        """Add a submitted job reference."""
        self.submitted_jobs.append(job)
        self.update_timestamp()
    
    def update_job_status(self, job_id: str, status: str, result: Optional[Dict] = None) -> bool:
        """Update status of a job. Returns True if found and updated."""
        for job in self.submitted_jobs:
            if job.job_id == job_id:
                job.status = status
                if status in ["completed", "failed"]:
                    job.completed_at = datetime.utcnow().isoformat()
                if result:
                    job.result_data = result
                self.update_timestamp()
                return True
        return False


@dataclass
class AliasIndexEntry:
    """Entry in the alias index for quick lookups."""
    alias: str
    session_id: str
    created_at: str
    last_modified: str


class AliasIndex:
    """In-memory index of aliases for quick lookups and conflict detection."""
    
    def __init__(self):
        self.entries: Dict[str, List[AliasIndexEntry]] = {}
    
    def add(self, alias: str, entry: AliasIndexEntry) -> None:
        """Add an entry to the index."""
        if alias not in self.entries:
            self.entries[alias] = []
        self.entries[alias].append(entry)
        # Sort by creation time descending (most recent first)
        self.entries[alias].sort(
            key=lambda e: e.created_at, 
            reverse=True
        )
    
    def get_by_alias(self, alias: str) -> List[AliasIndexEntry]:
        """Get all sessions with a given alias (sorted by time, most recent first)."""
        return self.entries.get(alias, [])
    
    def get_most_recent(self, alias: str) -> Optional[AliasIndexEntry]:
        """Get the most recent session with a given alias."""
        entries = self.get_by_alias(alias)
        return entries[0] if entries else None
    
    def remove(self, alias: str, session_id: str) -> bool:
        """Remove an entry from the index. Returns True if found and removed."""
        if alias in self.entries:
            self.entries[alias] = [
                e for e in self.entries[alias]
                if e.session_id != session_id
            ]
            if not self.entries[alias]:
                del self.entries[alias]
            return True
        return False
    
    def to_dict(self) -> Dict:
        """Serialize to dict for storage."""
        return {
            alias: [
                {
                    "alias": entry.alias,
                    "session_id": entry.session_id,
                    "created_at": entry.created_at,
                    "last_modified": entry.last_modified,
                }
                for entry in entries
            ]
            for alias, entries in self.entries.items()
        }
    
    @classmethod
    def from_dict(cls, data: Dict) -> "AliasIndex":
        """Deserialize from dict."""
        index = cls()
        for alias, entries in data.items():
            for entry_data in entries:
                entry = AliasIndexEntry(**entry_data)
                index.add(alias, entry)
        return index