from sqlalchemy import Boolean, DateTime, Float, ForeignKey, Integer, JSON, String, Text
from sqlalchemy.orm import Mapped, mapped_column, relationship

from app.core.database import Base
from app.models.base import TimestampMixin


class LLMProvider(TimestampMixin, Base):
    __tablename__ = "ai_llm_provider"

    name: Mapped[str] = mapped_column(String(100))
    code: Mapped[str] = mapped_column(String(80), unique=True, index=True)
    provider_type: Mapped[str] = mapped_column(String(40), default="openai_compatible")
    base_url: Mapped[str] = mapped_column(String(300), default="")
    api_key: Mapped[str] = mapped_column(String(500), default="")
    status: Mapped[str] = mapped_column(String(20), default="enabled")
    models: Mapped[list["LLMModel"]] = relationship("LLMModel", back_populates="provider")


class LLMModel(TimestampMixin, Base):
    __tablename__ = "ai_llm_model"

    provider_id: Mapped[str] = mapped_column(String(32), ForeignKey("ai_llm_provider.id"), index=True)
    name: Mapped[str] = mapped_column(String(120))
    display_name: Mapped[str] = mapped_column(String(120), default="")
    context_length: Mapped[int] = mapped_column(Integer, default=8192)
    supports_streaming: Mapped[bool] = mapped_column(Boolean, default=True)
    supports_function_call: Mapped[bool] = mapped_column(Boolean, default=False)
    default_temperature: Mapped[float] = mapped_column(Float, default=0.7)
    default_max_tokens: Mapped[int] = mapped_column(Integer, default=2048)
    status: Mapped[str] = mapped_column(String(20), default="enabled")
    provider: Mapped[LLMProvider] = relationship("LLMProvider", back_populates="models")


class Agent(TimestampMixin, Base):
    __tablename__ = "ai_agent"

    name: Mapped[str] = mapped_column(String(100))
    code: Mapped[str] = mapped_column(String(80), unique=True, index=True)
    description: Mapped[str] = mapped_column(Text, default="")
    avatar: Mapped[str] = mapped_column(String(300), default="")
    status: Mapped[str] = mapped_column(String(20), default="draft")
    persona: Mapped[dict] = mapped_column(JSON, default=dict)
    system_prompt: Mapped[str] = mapped_column(Text, default="")
    model_id: Mapped[str | None] = mapped_column(String(32), ForeignKey("ai_llm_model.id"), nullable=True)
    temperature: Mapped[float] = mapped_column(Float, default=0.7)
    max_tokens: Mapped[int] = mapped_column(Integer, default=2048)
    tools: Mapped[list] = mapped_column(JSON, default=list)
    knowledge_base_ids: Mapped[list] = mapped_column(JSON, default=list)
    enable_memory: Mapped[bool] = mapped_column(Boolean, default=True)
    memory_window: Mapped[int] = mapped_column(Integer, default=10)
    model: Mapped[LLMModel | None] = relationship("LLMModel")


class AgentConversation(TimestampMixin, Base):
    __tablename__ = "ai_agent_conversation"

    agent_id: Mapped[str] = mapped_column(String(32), ForeignKey("ai_agent.id"), index=True)
    user_id: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True)
    title: Mapped[str] = mapped_column(String(160), default="新建对话")
    total_tokens: Mapped[int] = mapped_column(Integer, default=0)
    agent: Mapped[Agent] = relationship("Agent")
    messages: Mapped[list["AgentMessage"]] = relationship("AgentMessage", back_populates="conversation")


class AgentMessage(TimestampMixin, Base):
    __tablename__ = "ai_agent_message"

    conversation_id: Mapped[str] = mapped_column(String(32), ForeignKey("ai_agent_conversation.id"), index=True)
    role: Mapped[str] = mapped_column(String(20))
    content: Mapped[str] = mapped_column(Text, default="")
    status: Mapped[str] = mapped_column(String(20), default="completed")
    reasoning_steps: Mapped[list] = mapped_column(JSON, default=list)
    total_tokens: Mapped[int] = mapped_column(Integer, default=0)
    elapsed_time: Mapped[int] = mapped_column(Integer, default=0)
    conversation: Mapped[AgentConversation] = relationship("AgentConversation", back_populates="messages")


class Workflow(TimestampMixin, Base):
    __tablename__ = "ai_workflow"

    name: Mapped[str] = mapped_column(String(100))
    code: Mapped[str] = mapped_column(String(80), unique=True, index=True)
    description: Mapped[str] = mapped_column(Text, default="")
    status: Mapped[str] = mapped_column(String(20), default="draft")
    version: Mapped[int] = mapped_column(Integer, default=1)
    published_version: Mapped[int | None] = mapped_column(Integer, nullable=True)
    definition: Mapped[dict] = mapped_column(JSON, default=dict)
    published_definition: Mapped[dict] = mapped_column(JSON, default=dict)
    input_variables: Mapped[list] = mapped_column(JSON, default=list)
    output_variables: Mapped[list] = mapped_column(JSON, default=list)
    run_count: Mapped[int] = mapped_column(Integer, default=0)


class WorkflowRun(TimestampMixin, Base):
    __tablename__ = "ai_workflow_run"

    workflow_id: Mapped[str] = mapped_column(String(32), ForeignKey("ai_workflow.id"), index=True)
    status: Mapped[str] = mapped_column(String(20), default="running")
    trigger_type: Mapped[str] = mapped_column(String(40), default="api")
    inputs: Mapped[dict] = mapped_column(JSON, default=dict)
    outputs: Mapped[dict] = mapped_column(JSON, default=dict)
    execution_log: Mapped[list] = mapped_column(JSON, default=list)
    error_message: Mapped[str] = mapped_column(Text, default="")
    total_tokens: Mapped[int] = mapped_column(Integer, default=0)
    elapsed_time: Mapped[int] = mapped_column(Integer, default=0)
    started_at: Mapped[object | None] = mapped_column(DateTime, nullable=True)
    completed_at: Mapped[object | None] = mapped_column(DateTime, nullable=True)
    workflow: Mapped[Workflow] = relationship("Workflow")


class KnowledgeBase(TimestampMixin, Base):
    __tablename__ = "ai_knowledge_base"

    name: Mapped[str] = mapped_column(String(120))
    code: Mapped[str] = mapped_column(String(80), unique=True, index=True)
    description: Mapped[str] = mapped_column(Text, default="")
    status: Mapped[str] = mapped_column(String(20), default="enabled")


class AgentTeam(TimestampMixin, Base):
    __tablename__ = "ai_agent_team"

    name: Mapped[str] = mapped_column(String(120))
    code: Mapped[str] = mapped_column(String(80), unique=True, index=True)
    description: Mapped[str] = mapped_column(Text, default="")
    mode: Mapped[str] = mapped_column(String(20), default="sequential")
    members: Mapped[list] = mapped_column(JSON, default=list)
    status: Mapped[str] = mapped_column(String(20), default="enabled")


class CollaborationRun(TimestampMixin, Base):
    __tablename__ = "ai_collaboration_run"

    team_id: Mapped[str] = mapped_column(String(32), ForeignKey("ai_agent_team.id"), index=True)
    task: Mapped[str] = mapped_column(Text, default="")
    status: Mapped[str] = mapped_column(String(20), default="running")
    messages: Mapped[list] = mapped_column(JSON, default=list)
    final_answer: Mapped[str] = mapped_column(Text, default="")
    elapsed_time: Mapped[int] = mapped_column(Integer, default=0)
    team: Mapped[AgentTeam] = relationship("AgentTeam")