backend/database/models.py

"""
Database models for the agent monitoring system.
"""

from datetime import datetime, timezone
import json
from sqlalchemy import Column, Integer, String, Float, Boolean, DateTime, Text, ForeignKey, Table, UniqueConstraint, Index
from sqlalchemy.orm import relationship
from sqlalchemy.types import JSON, TypeDecorator
from sqlalchemy.ext.declarative import declarative_base
import uuid

Base = declarative_base()

class SafeJSON(TypeDecorator):
    """Custom JSON type that handles circular references using default=str"""
    impl = Text
    
    def process_bind_param(self, value, dialect):
        if value is not None:
            return json.dumps(value, default=str)
        return value
    
    def process_result_value(self, value, dialect):
        if value is not None:
            return json.loads(value)
        return value


class Trace(Base):
    """Model for storing agent traces (conversations, interactions, etc.)."""
    __tablename__ = "traces"

    id = Column(Integer, primary_key=True, index=True)
    trace_id = Column(String(36), unique=True, index=True, default=lambda: str(uuid.uuid4()))
    filename = Column(String(255), nullable=True, index=True)
    title = Column(String(255), nullable=True)
    description = Column(Text, nullable=True)
    content = Column(Text, nullable=True)  # Full trace content
    content_hash = Column(String(64), nullable=True, index=True)  # Hash of content for deduplication
    upload_timestamp = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    update_timestamp = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))
    uploader = Column(String(255), nullable=True)
    trace_type = Column(String(50), nullable=True)  # e.g., 'conversation', 'code_execution', etc.
    trace_source = Column(String(50), nullable=True)  # e.g., 'user_upload', 'api', 'generated'
    character_count = Column(Integer, default=0)
    turn_count = Column(Integer, default=0)
    status = Column(String(50), default="uploaded")  # uploaded, processed, analyzed, etc.
    processing_method = Column(String(50), nullable=True)  # e.g., 'sliding_window', 'single_pass', etc.
    tags = Column(JSON, nullable=True)  # Store tags as JSON array
    trace_metadata = Column(JSON, nullable=True)  # Additional metadata as JSON

    # Relationships
    knowledge_graphs = relationship("KnowledgeGraph", back_populates="trace", 
                                   foreign_keys="KnowledgeGraph.trace_id", 
                                   cascade="all, delete-orphan")

    __table_args__ = (
        UniqueConstraint('trace_id', name='uix_trace_id'),
        Index('idx_trace_content_hash', 'content_hash'),
        Index('idx_trace_title', 'title'),
        Index('idx_trace_status', 'status'),
    )

    def to_dict(self):
        """Convert to dictionary representation."""
        return {
            "id": self.id,
            "trace_id": self.trace_id,
            "filename": self.filename,
            "title": self.title,
            "description": self.description,
            "upload_timestamp": self.upload_timestamp.isoformat() if self.upload_timestamp else None,
            "update_timestamp": self.update_timestamp.isoformat() if self.update_timestamp else None,
            "uploader": self.uploader,
            "trace_type": self.trace_type,
            "trace_source": self.trace_source,
            "character_count": self.character_count,
            "turn_count": self.turn_count,
            "status": self.status,
            "processing_method": self.processing_method,
            "tags": self.tags,
            "metadata": self.trace_metadata,
            "knowledge_graph_count": len(self.knowledge_graphs) if self.knowledge_graphs else 0
        }

    @classmethod
    def from_content(cls, content, filename=None, title=None, description=None, trace_type=None, 
                     trace_source="user_upload", uploader=None, tags=None, trace_metadata=None):
        """Create a Trace instance from content."""
        import hashlib
        
        trace = cls()
        trace.trace_id = str(uuid.uuid4())
        trace.filename = filename
        trace.title = title or f"Trace {trace.trace_id[:8]}"
        trace.description = description
        trace.content = content
        
        # Calculate content hash for deduplication
        if content:
            content_hash = hashlib.sha256(content.encode('utf-8')).hexdigest()
            trace.content_hash = content_hash
            
            # Set character count
            trace.character_count = len(content)
            
            # Estimate turn count (approximate)
            turn_markers = [
                "user:", "assistant:", "system:", "human:", "ai:", 
                "User:", "Assistant:", "System:", "Human:", "AI:"
            ]
            turn_count = 0
            for marker in turn_markers:
                turn_count += content.count(marker)
            trace.turn_count = max(1, turn_count)  # At least 1 turn
        
        trace.trace_type = trace_type
        trace.trace_source = trace_source
        trace.uploader = uploader
        trace.tags = tags or []
        trace.trace_metadata = trace_metadata or {}
        trace.status = "uploaded"
        
        return trace


class KnowledgeGraph(Base):
    """Model for storing knowledge graphs."""
    __tablename__ = "knowledge_graphs"

    id = Column(Integer, primary_key=True, index=True)
    filename = Column(String(255), unique=True, index=True)
    creation_timestamp = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    update_timestamp = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))
    entity_count = Column(Integer, default=0)
    relation_count = Column(Integer, default=0)
    _graph_data = Column("graph_data", Text, nullable=True)  # Underlying TEXT field
    status = Column(String(50), default="created", nullable=False)  # Status of processing: created, enriched, perturbed, causal
    
    # Add fields for trace and window tracking
    trace_id = Column(String(36), ForeignKey("traces.trace_id"), nullable=True, index=True, 
                    comment="ID to group knowledge graphs from the same trace")
    window_index = Column(Integer, nullable=True, 
                        comment="Sequential index of window within a trace")
    window_total = Column(Integer, nullable=True, 
                        comment="Total number of windows in the trace")
    window_start_char = Column(Integer, nullable=True, 
                             comment="Starting character position in the original trace")
    window_end_char = Column(Integer, nullable=True, 
                           comment="Ending character position in the original trace")
    processing_run_id = Column(String(36), nullable=True, index=True,
                             comment="ID to distinguish multiple processing runs of the same trace")
    
    # Relationships
    entities = relationship("Entity", back_populates="graph", cascade="all, delete-orphan")
    relations = relationship("Relation", back_populates="graph", cascade="all, delete-orphan")
    trace = relationship("Trace", back_populates="knowledge_graphs", foreign_keys=[trace_id])
    prompt_reconstructions = relationship(
        "PromptReconstruction", back_populates="knowledge_graph", cascade="all, delete-orphan"
    )
    perturbation_tests = relationship("PerturbationTest", back_populates="knowledge_graph", 
                                     cascade="all, delete-orphan")
    causal_analyses = relationship("CausalAnalysis", back_populates="knowledge_graph",
                                 cascade="all, delete-orphan")

    __table_args__ = (
        UniqueConstraint('filename', name='uix_knowledge_graph_filename'),
    )
    
    @property
    def graph_data(self):
        """Get the graph_data as a parsed JSON object"""
        if self._graph_data is None:
            return None
            
        if isinstance(self._graph_data, dict):
            # Already a dictionary, return as is
            return self._graph_data
            
        # Try to parse as JSON
        try:
            return json.loads(self._graph_data)
        except:
            # If parsing fails, return None
            return None
            
    @graph_data.setter
    def graph_data(self, value):
        """Set graph_data, converting to a JSON string if it's a dictionary"""
        if value is None:
            self._graph_data = None
        elif isinstance(value, dict):
            self._graph_data = json.dumps(value)
        else:
            # Assume it's already a string
            self._graph_data = value

    @property
    def graph_content(self):
        """Get the graph content from graph_data field"""
        # Return graph_data
        return self.graph_data or {}
    
    @graph_content.setter
    def graph_content(self, data):
        """Set graph content from a dictionary."""
        self.graph_data = data
        # Update counts
        if isinstance(data, dict):
            if 'entities' in data and isinstance(data['entities'], list):
                self.entity_count = len(data['entities'])
            if 'relations' in data and isinstance(data['relations'], list):
                self.relation_count = len(data['relations'])

    def get_entities_from_content(self):
        """Get entities directly from content field."""
        data = self.graph_content
        entities = data.get('entities', []) if isinstance(data, dict) else []
        
        return entities
    
    def get_relations_from_content(self):
        """Get relations directly from content field."""
        data = self.graph_content
        relations = data.get('relations', []) if isinstance(data, dict) else []
        
        return relations
    
    def get_all_entities(self, session=None):
        """
        Get all entities, preferring database entities if available.
        
        If no database entities exist, falls back to content entities.
        If session is provided, queries database entities, otherwise returns content entities.
        """
        if session:
            db_entities = session.query(Entity).filter_by(graph_id=self.id).all()
            if db_entities:
                return [entity.to_dict() for entity in db_entities]
        
        return self.get_entities_from_content()
    
    def get_all_relations(self, session=None):
        """
        Get all relations, preferring database relations if available.
        
        If no database relations exist, falls back to content relations.
        If session is provided, queries database relations, otherwise returns content relations.
        """
        if session:
            db_relations = session.query(Relation).filter_by(graph_id=self.id).all()
            if db_relations:
                return [relation.to_dict() for relation in db_relations]
        
        return self.get_relations_from_content()
    
    def to_dict(self):
        """Convert to dictionary representation."""
        result = {
            "id": self.id,
            "filename": self.filename,
            "creation_timestamp": self.creation_timestamp.isoformat(),
            "entity_count": self.entity_count,
            "relation_count": self.relation_count,
        }
        
        return result
        
    @classmethod
    def from_dict(cls, data):
        """Create a KnowledgeGraph instance from a dictionary representation."""
        kg = cls()
        kg.filename = data.get('filename')
        
        # Store content as JSON
        kg.content = json.dumps(data)
        
        return kg


class Entity(Base):
    """Model for storing knowledge graph entities."""
    __tablename__ = "entities"

    id = Column(Integer, primary_key=True, index=True)
    graph_id = Column(Integer, ForeignKey("knowledge_graphs.id"))
    entity_id = Column(String(255), index=True)  # Original entity ID in the graph
    type = Column(String(255))
    name = Column(String(255))
    properties = Column(JSON)

    # Relationships
    graph = relationship("KnowledgeGraph", back_populates="entities")
    source_relations = relationship("Relation", foreign_keys="Relation.source_id", back_populates="source")
    target_relations = relationship("Relation", foreign_keys="Relation.target_id", back_populates="target")
    
    # Add a composite unique constraint to ensure entity_id is unique per graph
    __table_args__ = (
        UniqueConstraint('graph_id', 'entity_id', name='uix_entity_graph_id_entity_id'),
    )
    
    def to_dict(self):
        """Convert to dictionary representation."""
        result = {
            "id": self.entity_id,
            "type": self.type,
            "name": self.name,
            "properties": self.properties or {}
        }
        
        return result
    
    @classmethod
    def from_dict(cls, data, graph_id):
        """Create an Entity instance from a dictionary."""
        entity = cls()
        entity.graph_id = graph_id
        entity.entity_id = data.get('id')
        entity.type = data.get('type')
        entity.name = data.get('name')
        entity.properties = data.get('properties')
        
        return entity


class Relation(Base):
    """Model for storing knowledge graph relations."""
    __tablename__ = "relations"

    id = Column(Integer, primary_key=True, index=True)
    graph_id = Column(Integer, ForeignKey("knowledge_graphs.id"))
    relation_id = Column(String(255), index=True)  # Original relation ID in the graph
    type = Column(String(255))
    source_id = Column(Integer, ForeignKey("entities.id"))
    target_id = Column(Integer, ForeignKey("entities.id"))
    properties = Column(JSON)

    # Relationships
    graph = relationship("KnowledgeGraph", back_populates="relations")
    source = relationship("Entity", foreign_keys=[source_id], back_populates="source_relations")
    target = relationship("Entity", foreign_keys=[target_id], back_populates="target_relations")
    
    # Add a composite unique constraint to ensure relation_id is unique per graph
    __table_args__ = (
        UniqueConstraint('graph_id', 'relation_id', name='uix_relation_graph_id_relation_id'),
    )
    
    def to_dict(self):
        """Convert to dictionary representation."""
        result = {
            "id": self.relation_id,
            "type": self.type,
            "source": self.source.entity_id if self.source else None,
            "target": self.target.entity_id if self.target else None,
            "properties": self.properties or {}
        }
        
        return result
    
    @classmethod
    def from_dict(cls, data, graph_id, source_entity=None, target_entity=None):
        """Create a Relation instance from a dictionary."""
        relation = cls()
        relation.graph_id = graph_id
        relation.relation_id = data.get('id')
        relation.type = data.get('type')
        
        # Set source and target
        if source_entity:
            relation.source_id = source_entity.id
        
        if target_entity:
            relation.target_id = target_entity.id
        
        # Set properties
        relation.properties = data.get('properties')
        
        return relation


class PromptReconstruction(Base):
    """Model for storing prompt reconstruction results."""
    __tablename__ = "prompt_reconstructions"

    id = Column(Integer, primary_key=True)
    knowledge_graph_id = Column(Integer, ForeignKey("knowledge_graphs.id"), nullable=False)
    relation_id = Column(String(255), nullable=False)
    reconstructed_prompt = Column(Text)
    dependencies = Column(JSON)
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))

    # Relationships
    knowledge_graph = relationship("KnowledgeGraph", back_populates="prompt_reconstructions")
    perturbation_tests = relationship("PerturbationTest", back_populates="prompt_reconstruction")

    def to_dict(self):
        return {
            "id": self.id,
            "knowledge_graph_id": self.knowledge_graph_id,
            "relation_id": self.relation_id,
            "reconstructed_prompt": self.reconstructed_prompt,
            "dependencies": self.dependencies,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None
        }


class PerturbationTest(Base):
    """Model for storing perturbation test results."""
    __tablename__ = "perturbation_tests"

    id = Column(Integer, primary_key=True)
    knowledge_graph_id = Column(Integer, ForeignKey("knowledge_graphs.id"), nullable=False)
    prompt_reconstruction_id = Column(Integer, ForeignKey("prompt_reconstructions.id"), nullable=False)
    relation_id = Column(String(255), nullable=False)
    perturbation_type = Column(String(50), nullable=False)  # e.g., 'entity_removal', 'relation_removal'
    perturbation_set_id = Column(String(64), nullable=False, index=True)
    test_result = Column(JSON)
    perturbation_score = Column(Float)
    test_metadata = Column(JSON)
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))

    # Relationships
    knowledge_graph = relationship("KnowledgeGraph", back_populates="perturbation_tests")
    prompt_reconstruction = relationship("PromptReconstruction", back_populates="perturbation_tests")

    def to_dict(self):
        return {
            "id": self.id,
            "knowledge_graph_id": self.knowledge_graph_id,
            "prompt_reconstruction_id": self.prompt_reconstruction_id,
            "relation_id": self.relation_id,
            "perturbation_type": self.perturbation_type,
            "perturbation_set_id": self.perturbation_set_id,
            "test_result": self.test_result,
            "perturbation_score": self.perturbation_score,
            "test_metadata": self.test_metadata,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None
        }


class CausalAnalysis(Base):
    """Model for storing causal analysis results."""
    __tablename__ = "causal_analyses"

    id = Column(Integer, primary_key=True)
    knowledge_graph_id = Column(Integer, ForeignKey("knowledge_graphs.id"), nullable=False)
    perturbation_set_id = Column(String(64), nullable=False, index=True)
    # Analysis method and results
    analysis_method = Column(String(50), nullable=False)  # e.g., 'graph', 'component', 'dowhy'
    analysis_result = Column(JSON)  # Store the full analysis result
    causal_score = Column(Float)  # Store the numerical causal score
    analysis_metadata = Column(JSON)  # Store additional metadata about the analysis
    # Timestamps
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))

    # Relationships
    knowledge_graph = relationship("KnowledgeGraph", back_populates="causal_analyses")

    # Indexes
    __table_args__ = (
        Index("idx_causal_analyses_kgid", "knowledge_graph_id"),
        Index("idx_causal_analyses_method", "analysis_method"),
        Index("idx_causal_analyses_setid", "perturbation_set_id"),
    )

    def to_dict(self):
        return {
            "id": self.id,
            "knowledge_graph_id": self.knowledge_graph_id,
            "perturbation_set_id": self.perturbation_set_id,
            "analysis_method": self.analysis_method,
            "analysis_result": self.analysis_result,
            "causal_score": self.causal_score,
            "analysis_metadata": self.analysis_metadata,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None
        }


class ObservabilityConnection(Base):
    """Model for storing AI observability platform connections."""
    __tablename__ = "observability_connections"

    id = Column(Integer, primary_key=True, index=True)
    connection_id = Column(String(36), unique=True, index=True, default=lambda: str(uuid.uuid4()))
    platform = Column(String(50), nullable=False)  # langfuse, langsmith, etc.
    public_key = Column(Text, nullable=False)  # Encrypted API key
    secret_key = Column(Text, nullable=True)  # Encrypted secret key (for Langfuse)
    host = Column(String(255), nullable=True)  # Host URL
    projects = Column(JSON, nullable=True)  # Available projects from the platform
    status = Column(String(50), default="connected")
    connected_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    last_sync = Column(DateTime, nullable=True)
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))

    # Relationships
    fetched_traces = relationship("FetchedTrace", back_populates="connection", cascade="all, delete-orphan")

    def to_dict(self):
        return {
            "id": self.connection_id,
            "platform": self.platform,
            "status": self.status,
            "connected_at": self.connected_at.isoformat() if self.connected_at else None,
            "last_sync": self.last_sync.isoformat() if self.last_sync else None,
            "host": self.host,
            "projects": self.projects or []
        }


class FetchedTrace(Base):
    """Model for storing fetched traces from observability platforms."""
    __tablename__ = "fetched_traces"

    id = Column(Integer, primary_key=True, index=True)
    trace_id = Column(String(255), nullable=False, index=True)  # Original trace ID from platform
    name = Column(String(255), nullable=False)
    platform = Column(String(50), nullable=False)
    connection_id = Column(String(36), ForeignKey("observability_connections.connection_id"), nullable=False)
    project_name = Column(String(255), nullable=True, index=True)  # Project name for LangSmith, null for Langfuse
    data = Column(SafeJSON, nullable=True)  # Full trace data
    fetched_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    imported = Column(Boolean, default=False)
    imported_at = Column(DateTime, nullable=True)
    imported_trace_id = Column(String(36), nullable=True)  # Reference to imported trace

    # Relationships
    connection = relationship("ObservabilityConnection", back_populates="fetched_traces")

    __table_args__ = (
        UniqueConstraint('trace_id', 'connection_id', name='uix_fetched_trace_id_connection'),
    )

    def _extract_generated_timestamp(self):
        """Extract the actual generated timestamp from trace data based on platform."""
        if not self.data:
            return None
        
        if self.platform == "langfuse":
            # For Langfuse, find the earliest timestamp from traces
            traces = self.data.get("traces", [])
            if traces:
                timestamps = []
                for trace in traces:
                    if isinstance(trace, dict):
                        # Check for various timestamp fields in Langfuse traces
                        for ts_field in ["timestamp", "startTime", "createdAt"]:
                            if ts_field in trace:
                                timestamps.append(trace[ts_field])
                                break
                if timestamps:
                    return min(timestamps)
            
            # Fallback to session info or other timestamps
            session_info = self.data.get("session_info", {})
            if session_info and "createdAt" in session_info:
                return session_info["createdAt"]
            
            # Other fallback fields at top level
            for field in ["timestamp", "createdAt", "startTime"]:
                if field in self.data:
                    return self.data[field]
        
        elif self.platform == "langsmith":
            # For LangSmith, find the earliest start_time from traces
            traces = self.data.get("traces", [])
            if traces:
                start_times = []
                for trace in traces:
                    if isinstance(trace, dict) and "start_time" in trace:
                        start_times.append(trace["start_time"])
                if start_times:
                    return min(start_times)
            
            # Fallback to other timestamp fields
            for field in ["timestamp", "start_time", "created_at"]:
                if field in self.data:
                    return self.data[field]
        
        return None

    def to_dict(self, preview=True):
        data = self.data
        original_stats = {}
        
        if data:
            # Calculate original data statistics
            import json
            original_json_str = json.dumps(data, ensure_ascii=False)
            original_stats = {
                "original_character_count": len(original_json_str),
                "original_line_count": original_json_str.count('\n') + 1,
                "original_size_kb": round(len(original_json_str) / 1024, 2)
            }
            
            if preview:
                # Truncate long strings to prevent browser crashes but preserve full structure
                from backend.routers.observability import truncate_long_strings
                data = truncate_long_strings(data, max_string_length=500)
        
        # Extract generated timestamp
        generated_timestamp = self._extract_generated_timestamp()
        
        result = {
            "id": self.trace_id,
            "name": self.name,
            "platform": self.platform,
            "fetched_at": self.fetched_at.isoformat() if self.fetched_at else None,
            "generated_timestamp": generated_timestamp,
            "imported": self.imported,
            "imported_at": self.imported_at.isoformat() if self.imported_at else None,
            "data": data
        }
        
        # Add original statistics to the result
        result.update(original_stats)
        
        return result
    
    def get_full_data(self):
        """Get full original data for download (no limitations)"""
        return {
            "id": self.trace_id,
            "name": self.name,
            "platform": self.platform,
            "fetched_at": self.fetched_at.isoformat() if self.fetched_at else None,
            "imported": self.imported,
            "imported_at": self.imported_at.isoformat() if self.imported_at else None,
            "data": self.data  # Full original data
        }