agent_management/database.py

"""
Database models and session management for Agent Manager.
Supports both SQLite (development) and PostgreSQL (production).
"""

from datetime import datetime
from typing import Optional, Dict, Any
from sqlalchemy import (
    create_engine, Column, Integer, String, DateTime, 
    Boolean, JSON, Float, ForeignKey, Text, Enum as SQLEnum
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship, Session
import enum
import os

# Database URL - PostgreSQL (hardcoded)
DATABASE_URL = "postgresql://taiji:By%40123456.@taijipda.postgres.database.azure.com:5432/taijiagnet"

engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


class AgentType(str, enum.Enum):
    """Agent type enumeration"""
    PLATFORM = "platform"
    CUSTOM = "custom"


class AgentStatus(str, enum.Enum):
    """Agent status enumeration"""
    PENDING = "pending"
    RUNNING = "running"
    STOPPED = "stopped"
    FAILED = "failed"
    SCALING = "scaling"


class Template(Base):
    """Template model for agent templates"""
    __tablename__ = "templates"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100), unique=True, nullable=False, index=True)
    display_name = Column(String(200), nullable=False)
    description = Column(Text)
    agent_type = Column(SQLEnum(AgentType), nullable=False, index=True)
    
    # Framework configuration (NEW)
    agent_framework = Column(String(50), default="langchain")  # "langchain", "mcp", "a2a"
    
    # Image configuration
    image = Column(String(500), nullable=False)
    port = Column(Integer, nullable=True)
    
    # Environment variable requirements (JSON format)
    # {"required": {"KEY": "description"}, "optional": {"KEY": "description"}}
    env_requirements = Column(JSON, default={})
    
    # Tools configuration (NEW) - JSON format for tool definitions
    tools_config = Column(JSON, default={})
    
    # Model configuration defaults (NEW)
    default_model_provider = Column(String(100))  # e.g., "openai", "azure-openai"
    default_model_name = Column(String(200))  # e.g., "gpt-4", "claude-3"
    
    # Resource configuration (for platform agents, fixed by admin)
    cpu_request = Column(String(20))  # e.g., "100m"
    cpu_limit = Column(String(20))    # e.g., "500m"
    memory_request = Column(String(20))  # e.g., "128Mi"
    memory_limit = Column(String(20))    # e.g., "512Mi"
    
    # Scaling configuration defaults
    min_replicas = Column(Integer, default=1)
    max_replicas = Column(Integer, default=3)
    target_cpu_utilization = Column(Integer, default=80)
    
    # Metadata
    is_active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    created_by = Column(String(100))
    
    # Relationships
    agents = relationship("Agent", back_populates="template")


class Agent(Base):
    """Agent instance model"""
    __tablename__ = "agents"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100), unique=True, nullable=False, index=True)
    display_name = Column(String(200))
    
    # Template reference (nullable for backward compatibility with non-template agents)
    template_id = Column(Integer, ForeignKey("templates.id"), nullable=True)
    template = relationship("Template", back_populates="agents")
    
    # Ownership and organization
    owner_id = Column(String(100), nullable=False, index=True)
    channel_id = Column(String(100), index=True)
    tenant_id = Column(String(100), index=True)
    
    # Agent configuration
    agent_type = Column(SQLEnum(AgentType), nullable=False, index=True)
    status = Column(SQLEnum(AgentStatus), default=AgentStatus.PENDING, index=True)
    
    # Framework configuration (NEW)
    agent_framework = Column(String(50), default="langchain")  # "langchain", "mcp", "a2a"
    
    # Environment variables (encrypted in production)
    environment_vars = Column(JSON, default={})
    
    # Tools configuration (NEW) - Instance-level tools override
    tools_config = Column(JSON, default={})
    tool_endpoint = Column(String(500))  # External tool endpoint URL
    tool_api_key = Column(String(500))   # Encrypted tool API key
    
    # Model configuration (NEW) - Instance-level model settings
    model_provider = Column(String(100))  # e.g., "openai", "azure-openai"
    model_name = Column(String(200))      # e.g., "gpt-4"
    model_endpoint = Column(String(500))  # Model API endpoint
    model_api_key = Column(String(500))   # Encrypted model API key
    
    # Storage configuration (NEW) - For agents that need storage
    storage_connection_string = Column(String(1000))  # Encrypted storage connection
    storage_account_name = Column(String(200))
    
    # Resource configuration (for custom agents)
    cpu_request = Column(String(20))
    cpu_limit = Column(String(20))
    memory_request = Column(String(20))
    memory_limit = Column(String(20))
    
    # Scaling configuration
    min_replicas = Column(Integer, default=1)
    max_replicas = Column(Integer, default=3)
    target_cpu_utilization = Column(Integer, default=80)
    current_replicas = Column(Integer, default=0)
    
    # Kubernetes resources
    deployment_name = Column(String(100))
    service_name = Column(String(100))
    service_url = Column(String(500))
    namespace = Column(String(100), default="ai-agents")
    
    # Access information (NEW) - LoadBalancer and DNS details
    external_ip = Column(String(100), nullable=True)  # LoadBalancer external IP
    domain = Column(String(500), nullable=True)  # Full DNS domain (e.g., agent-name.taijiagnet.com)
    ip_url = Column(String(500), nullable=True)  # HTTP URL via IP (http://x.x.x.x:80)
    domain_url = Column(String(500), nullable=True)  # HTTP URL via domain
    recommended_url = Column(String(500), nullable=True)  # Recommended access URL
    
    # Metadata
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    last_accessed_at = Column(DateTime, nullable=True)
    
    # Relationships
    metrics = relationship("AgentMetric", back_populates="agent", cascade="all, delete-orphan")


class Quota(Base):
    """Resource quota model"""
    __tablename__ = "quotas"

    id = Column(Integer, primary_key=True, index=True)
    
    # Quota owner (hierarchical: admin -> channel -> tenant)
    owner_type = Column(String(20), nullable=False)  # "admin", "channel", "tenant"
    owner_id = Column(String(100), nullable=False, index=True)
    channel_id = Column(String(100), index=True)
    
    # For platform agents: quota is Pod count
    platform_pod_quota = Column(Integer, default=0)
    platform_pod_used = Column(Integer, default=0)
    
    # For custom agents: quota is CPU/Memory totals
    custom_cpu_quota = Column(Float, default=0.0)  # in cores
    custom_cpu_used = Column(Float, default=0.0)
    custom_memory_quota = Column(Float, default=0.0)  # in GB
    custom_memory_used = Column(Float, default=0.0)
    
    # Metadata
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # Unique constraint
    __table_args__ = (
        # Unique per owner
        # UniqueConstraint('owner_type', 'owner_id', name='uq_quota_owner'),
    )


class AgentMetric(Base):
    """Agent metrics tracking"""
    __tablename__ = "agent_metrics"

    id = Column(Integer, primary_key=True, index=True)
    agent_id = Column(Integer, ForeignKey("agents.id"), nullable=False, index=True)
    agent = relationship("Agent", back_populates="metrics")
    
    # Timestamp
    timestamp = Column(DateTime, default=datetime.utcnow, index=True)
    
    # Resource metrics
    cpu_usage = Column(Float)  # in cores
    memory_usage = Column(Float)  # in MB
    network_rx_bytes = Column(Integer, default=0)
    network_tx_bytes = Column(Integer, default=0)
    
    # Replica count
    replica_count = Column(Integer, default=0)
    
    # Request metrics
    request_count = Column(Integer, default=0)
    error_count = Column(Integer, default=0)


# ============================================================================
# Heicode Integration Models (NEW)
# ============================================================================

class DeploymentStatus(str, enum.Enum):
    """Deployment status for Heicode integration"""
    PENDING = "pending"
    RUNNING = "running"
    STOPPED = "stopped"
    FAILED = "failed"


class RiskLevel(str, enum.Enum):
    """Risk level for deployments"""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"


class BillingProvider(str, enum.Enum):
    """Model gateway provider"""
    NEWAPI = "newapi"
    LITELLM = "litellm"


class Deployment(Base):
    """Deployment model for Heicode integration"""
    __tablename__ = "deployments"

    id = Column(Integer, primary_key=True, index=True)
    deployment_id = Column(String(100), unique=True, nullable=False, index=True)

    # Ownership
    user_id = Column(String(100), nullable=False, index=True)
    binding_scope = Column(String(200), nullable=False, index=True)
    correlation_id = Column(String(100), index=True)

    # Deployment configuration
    orchestration_plan = Column(Text, nullable=False)
    risk_level = Column(SQLEnum(RiskLevel), nullable=False)
    approval_token = Column(Text)

    # Budget
    budget_max_usd = Column(Float)
    budget_consumed_usd = Column(Float, default=0.0)
    budget_alert_threshold_pct = Column(Integer, default=80)

    # Model gateway configuration
    billing_provider = Column(SQLEnum(BillingProvider), nullable=False)
    default_model_id = Column(String(200), nullable=False)
    allowed_model_ids = Column(JSON, nullable=False)
    secret_ref = Column(String(500))

    # Resource grants
    resource_grants = Column(JSON, default=[])

    # Status
    status = Column(SQLEnum(DeploymentStatus), nullable=False, default=DeploymentStatus.PENDING, index=True)
    phase = Column(String(100))
    error_message = Column(Text)

    # Kubernetes resources
    namespace = Column(String(100), nullable=False)
    configmap_name = Column(String(100))

    # Metadata
    created_at = Column(DateTime, default=datetime.utcnow, index=True)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    stopped_at = Column(DateTime)

    # Relationships
    agent_instances = relationship("AgentInstance", back_populates="deployment", cascade="all, delete-orphan")
    events = relationship("Event", back_populates="deployment", cascade="all, delete-orphan")


class AgentInstance(Base):
    """Agent instance model for Heicode deployments"""
    __tablename__ = "agent_instances"

    id = Column(Integer, primary_key=True, index=True)
    agent_instance_id = Column(String(100), unique=True, nullable=False, index=True)
    deployment_id = Column(String(100), ForeignKey("deployments.deployment_id", ondelete="CASCADE"), nullable=False, index=True)

    # Agent configuration
    role = Column(String(100), nullable=False)
    image = Column(String(500), nullable=False)
    phase = Column(String(100))

    # Kubernetes resources
    namespace = Column(String(100), nullable=False)
    pod_name = Column(String(100), nullable=False)
    service_account = Column(String(100))

    # Status
    status = Column(SQLEnum(DeploymentStatus), nullable=False, default=DeploymentStatus.PENDING, index=True)
    error_message = Column(Text)

    # Metadata
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationships
    deployment = relationship("Deployment", back_populates="agent_instances")


class Event(Base):
    """Event model for deployment events"""
    __tablename__ = "events"

    id = Column(Integer, primary_key=True, index=True)
    event_id = Column(String(100), unique=True, nullable=False, index=True)
    deployment_id = Column(String(100), ForeignKey("deployments.deployment_id", ondelete="CASCADE"), nullable=False, index=True)
    agent_instance_id = Column(String(100), ForeignKey("agent_instances.agent_instance_id", ondelete="SET NULL"))

    # Event details
    event_type = Column(String(100), nullable=False, index=True)
    correlation_id = Column(String(100))
    payload = Column(JSON)

    # Metadata
    occurred_at = Column(DateTime, default=datetime.utcnow, index=True)

    # Relationships
    deployment = relationship("Deployment", back_populates="events")


class AuditLog(Base):
    """Audit log model for tracking all operations"""
    __tablename__ = "audit_logs"

    id = Column(Integer, primary_key=True, index=True)
    audit_id = Column(String(100), unique=True, nullable=False, index=True)

    # Actor
    actor = Column(String(200), nullable=False, index=True)
    user_id = Column(String(100), index=True)
    binding_scope = Column(String(200), index=True)

    # Action
    action = Column(String(100), nullable=False, index=True)
    resource_type = Column(String(50), nullable=False)
    resource_id = Column(String(100))

    # Request details
    correlation_id = Column(String(100))
    request_payload = Column(JSON)

    # Result
    result = Column(String(50), nullable=False, index=True)
    error_code = Column(String(50))
    error_message = Column(Text)

    # Metadata
    occurred_at = Column(DateTime, default=datetime.utcnow, index=True)
    ip_address = Column(String(50))
    user_agent = Column(Text)


# ============================================================================
# Sub-mode Runtime Internal Models
# ============================================================================

class SwarmStatus(str, enum.Enum):
    """Swarm status enumeration"""
    INITIALIZING = "initializing"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    STOPPED = "stopped"


class SwarmAgentStatus(str, enum.Enum):
    """Swarm agent status enumeration"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


class Swarm(Base):
    """Internal runtime run model used by Heicode sub-mode execution."""
    __tablename__ = "swarms"

    id = Column(Integer, primary_key=True, index=True)
    swarm_id = Column(String(100), unique=True, nullable=False, index=True)

    # Task information
    task_description = Column(Text, nullable=False)
    project_context = Column(JSON)  # {repo_url, branch, language, framework}

    # Orchestration configuration
    orchestration_strategy = Column(String(50), default="sequential")  # sequential, parallel, hybrid
    max_iterations = Column(Integer, default=3)
    timeout_minutes = Column(Integer, default=30)

    # Status
    status = Column(SQLEnum(SwarmStatus), default=SwarmStatus.INITIALIZING, index=True)
    phase = Column(String(50))  # planning, coding, reviewing, testing
    progress = Column(Integer, default=0)  # 0-100

    # Results
    artifacts = Column(JSON, default=[])  # Generated code, documents, etc.
    error_message = Column(Text)

    # Metrics
    total_messages = Column(Integer, default=0)
    tokens_used = Column(Integer, default=0)

    # Callback
    callback_url = Column(String(500))
    callback_method = Column(String(10), default="POST")

    # Ownership
    owner_id = Column(String(100), nullable=False, index=True)

    # Timestamps
    created_at = Column(DateTime, default=datetime.utcnow, index=True)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    completed_at = Column(DateTime)

    # Relationships
    swarm_agents = relationship("SwarmAgent", back_populates="swarm", cascade="all, delete-orphan")
    swarm_messages = relationship("SwarmMessage", back_populates="swarm", cascade="all, delete-orphan")


class SwarmAgent(Base):
    """Internal runtime agent model for sub-mode execution."""
    __tablename__ = "swarm_agents"

    id = Column(Integer, primary_key=True, index=True)
    agent_id = Column(String(100), unique=True, nullable=False, index=True)
    swarm_id = Column(String(100), ForeignKey("swarms.swarm_id", ondelete="CASCADE"), nullable=False, index=True)

    # Agent configuration
    role = Column(String(100), nullable=False)  # architect, coder, reviewer, tester
    template = Column(String(100), nullable=False)  # a2a_litellm_agent, code_manager_agent
    model = Column(String(200))  # gpt-4, claude-3, etc.
    capabilities = Column(JSON, default=[])  # ["design", "coding", "review"]
    system_prompt = Column(Text)

    # Kubernetes resources
    namespace = Column(String(100), nullable=False)
    pod_name = Column(String(100), nullable=False)
    service_url = Column(String(500))
    external_ip = Column(String(100))

    # Status
    status = Column(SQLEnum(SwarmAgentStatus), default=SwarmAgentStatus.PENDING, index=True)
    current_task = Column(Text)
    output = Column(Text)

    # Timestamps
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationships
    swarm = relationship("Swarm", back_populates="swarm_agents")


class SwarmMessage(Base):
    """Internal runtime message model for sub-mode execution."""
    __tablename__ = "swarm_messages"

    id = Column(Integer, primary_key=True, index=True)
    message_id = Column(String(100), unique=True, nullable=False, index=True)
    swarm_id = Column(String(100), ForeignKey("swarms.swarm_id", ondelete="CASCADE"), nullable=False, index=True)

    # Message information
    from_agent_id = Column(String(100), index=True)  # NULL for orchestrator
    to_agent_id = Column(String(100), index=True)  # NULL for broadcast
    message_type = Column(String(50))  # task, response, broadcast, artifact
    content = Column(Text)
    message_metadata = Column("metadata", JSON)

    # Timestamp
    created_at = Column(DateTime, default=datetime.utcnow, index=True)

    # Relationships
    swarm = relationship("Swarm", back_populates="swarm_messages")


# Database initialization
def init_db():
    """Initialize database tables"""
    Base.metadata.create_all(bind=engine)


def get_db() -> Session:
    """Get database session (FastAPI dependency)"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def parse_resource_string(resource_str: str) -> float:
    """Parse Kubernetes resource string to float
    
    Examples:
        "100m" -> 0.1 (cores)
        "2" -> 2.0 (cores)
        "128Mi" -> 128.0 (MB)
        "1Gi" -> 1024.0 (MB)
    """
    if not resource_str:
        return 0.0
    
    resource_str = resource_str.strip()
    
    # CPU resources
    if resource_str.endswith('m'):
        return float(resource_str[:-1]) / 1000.0
    
    # Memory resources
    if resource_str.endswith('Mi'):
        return float(resource_str[:-2])
    elif resource_str.endswith('Gi'):
        return float(resource_str[:-2]) * 1024.0
    elif resource_str.endswith('Ki'):
        return float(resource_str[:-2]) / 1024.0
    
    # Plain number
    try:
        return float(resource_str)
    except ValueError:
        return 0.0


# Initialize database on import
init_db()