"""
数据库模型定义
定义工作流相关的所有数据表结构
"""

from datetime import datetime
from sqlalchemy import Column, String, Text, DateTime, ForeignKey, Integer, Boolean
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship

# 创建基类
Base = declarative_base()


class Workflow(Base):
    """工作流模型"""

    __tablename__ = "workflows"

    name = Column(String, primary_key=True)
    base_name = Column(String, nullable=False)
    version = Column(String, nullable=False)
    workflow_json = Column(Text, nullable=False)

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "base_name": self.base_name,
            "version": self.version,
            "workflow_json": self.workflow_json,
        }


class WorkflowRun(Base):
    """工作流运行记录模型"""

    __tablename__ = "workflow_run"

    id = Column(String, primary_key=True)
    workflow_name = Column(String, nullable=False)
    prompt_id = Column(String)
    client_id = Column(String)
    status = Column(String, nullable=False, default="pending")
    server_url = Column(String)
    created_at = Column(DateTime, default=datetime.utcnow)
    started_at = Column(DateTime)
    completed_at = Column(DateTime)
    error_message = Column(Text)
    workflow_json = Column(Text, nullable=False)
    api_spec = Column(Text, nullable=False)
    request_data = Column(Text, nullable=False)
    result = Column(Text)

    # 关联关系
    nodes = relationship(
        "WorkflowRunNode", back_populates="workflow_run", cascade="all, delete-orphan"
    )

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "workflow_name": self.workflow_name,
            "prompt_id": self.prompt_id,
            "client_id": self.client_id,
            "status": self.status,
            "server_url": self.server_url,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": (
                self.completed_at.isoformat() if self.completed_at else None
            ),
            "error_message": self.error_message,
            "workflow_json": self.workflow_json,
            "api_spec": self.api_spec,
            "request_data": self.request_data,
            "result": self.result,
        }


class WorkflowRunNode(Base):
    """工作流运行节点模型"""

    __tablename__ = "workflow_run_nodes"

    id = Column(String, primary_key=True)
    workflow_run_id = Column(String, ForeignKey("workflow_run.id"), nullable=False)
    node_id = Column(String, nullable=False)
    node_type = Column(String, nullable=False)
    status = Column(String, nullable=False, default="pending")
    started_at = Column(DateTime)
    completed_at = Column(DateTime)
    output_data = Column(Text)
    error_message = Column(Text)

    # 关联关系
    workflow_run = relationship("WorkflowRun", back_populates="nodes")

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "workflow_run_id": self.workflow_run_id,
            "node_id": self.node_id,
            "node_type": self.node_type,
            "status": self.status,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": (
                self.completed_at.isoformat() if self.completed_at else None
            ),
            "output_data": self.output_data,
            "error_message": self.error_message,
        }


class ComfyUIServer(Base):
    """ComfyUI 服务器模型"""

    __tablename__ = "comfyui_servers"

    name = Column(String, primary_key=True)
    http_url = Column(String, nullable=False)
    ws_url = Column(String, nullable=False)
    status = Column(String, nullable=False, default="offline")
    last_health_check = Column(DateTime)
    current_tasks = Column(Integer, default=0)
    max_concurrent_tasks = Column(Integer, default=1)
    capabilities = Column(Text)  # JSON 字符串
    server_metadata = Column(Text)  # JSON 字符串
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "http_url": self.http_url,
            "ws_url": self.ws_url,
            "status": self.status,
            "last_health_check": self.last_health_check.isoformat() if self.last_health_check else None,
            "current_tasks": self.current_tasks,
            "max_concurrent_tasks": self.max_concurrent_tasks,
            "capabilities": self.capabilities,
            "server_metadata": self.server_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,
        }