import uuid
from enum import Enum
from typing import Annotated, List
from datetime import datetime
from fastapi import Depends
from sqlmodel import SQLModel, Session, Field, create_engine, Relationship
from sqlalchemy import Column, DateTime, func


class User(SQLModel, table=True):
    github_id: str = Field(primary_key=True, index=True)
    display_name: str = Field(unique=True, index=True)
    github_user_json: str
    github_emails_json: str
    ab_api_key: str = Field(unique=True, index=True)
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), server_default=func.now(), index=True
        ),
    )

    agents: list["Agent"] = Relationship(back_populates="user")
    assessments: list["Assessment"] = Relationship(back_populates="user")


class AgentAssessmentLink(SQLModel, table=True):
    agent_id: uuid.UUID | None = Field(
        default=None, foreign_key="agent.id", primary_key=True
    )
    assessment_id: uuid.UUID | None = Field(
        default=None, foreign_key="assessment.id", primary_key=True
    )


class AgentDeployType(str, Enum):
    HOSTED = "hosted"
    REMOTE = "remote"
    PROXIED = "proxied"


class AgentHostedStatus(str, Enum):
    PENDING = "pending"  # or not hosted
    STARTING = "starting"
    RUNNING = "running"
    ERROR = "error"


class Agent(SQLModel, table=True):
    id: uuid.UUID | None = Field(default_factory=uuid.uuid4, primary_key=True)
    name: str = Field(index=True)
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), server_default=func.now(), index=True
        ),
    )
    # agent types
    is_green: bool = Field(index=True)
    deploy_type: AgentDeployType = Field(index=True)
    ctrl_url: str | None
    agent_url: str | None  # mainly used for proxied agents
    # agent config
    secret: str
    inject_litellm_proxy_api: bool
    # -- three ways to instantiate hosted agents: git, docker, description_prompt
    git_url: str | None
    git_branch: str | None
    docker_image_url: str | None
    description_prompt: str | None
    # status
    hosted_status: AgentHostedStatus = Field(
        default=AgentHostedStatus.PENDING, index=True
    )
    hosted_error_msg: str | None
    gcp_docker_image_url: str | None
    gcp_cloudrun_service_name: str | None

    user_id: str = Field(foreign_key="user.github_id", index=True)
    user: "User" = Relationship(back_populates="agents")

    agent_checks: list["AgentCheck"] = Relationship(back_populates="agent")
    assessments: list["Assessment"] = Relationship(
        back_populates="agents", link_model=AgentAssessmentLink
    )
    # FIXME: do we need relations that links to results?


class AgentCheck(SQLModel, table=True):
    # only check when initiated and also when requested for assessment
    id: uuid.UUID | None = Field(default_factory=uuid.uuid4, primary_key=True)
    agent_id: uuid.UUID = Field(foreign_key="agent.id", index=True)
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), server_default=func.now(), index=True
        ),
    )
    is_ctrl_reachable: bool | None
    agent_count: int | None
    agent_url: str | None
    card_url: str | None
    card_content: str | None

    agent: Agent = Relationship(back_populates="agent_checks")


class AssessmentStatus(str, Enum):
    PENDING = "pending"
    REJECTED = "rejected"
    STARTING = "starting"
    RUNNING = "running"
    COMPLETED = "completed"
    ERROR = "error"


class Assessment(SQLModel, table=True):
    id: uuid.UUID | None = Field(default_factory=uuid.uuid4, primary_key=True)
    batch_first_id: uuid.UUID | None = Field(index=True)
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), server_default=func.now(), index=True
        ),
    )
    submitted_by: str = Field(foreign_key="user.github_id", index=True)  # github_id

    green_agent_id: str = Field(index=True)  # no foreign key for now
    agent_id_list_json: str  # list of agent IDs in JSON array format
    config: str
    repeat_n: int
    repeat_idx: int
    status: AssessmentStatus = Field(default=AssessmentStatus.PENDING, index=True)

    agents: List[Agent] = Relationship(
        back_populates="assessments", link_model=AgentAssessmentLink
    )
    log_entries: List["LogEntry"] = Relationship(back_populates="assessment")
    user: "User" = Relationship(back_populates="assessments")
    # FIXME: do we need relations that links to results?


class LogEntry(SQLModel, table=True):
    id: uuid.UUID | None = Field(default_factory=uuid.uuid4, primary_key=True)
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), server_default=func.now(), index=True
        ),
    )
    source: str | None = Field(index=True)  # which agent, or system
    channel: str | None = Field(index=True)  # llm gateway, mcpcp, agent, sdk, etc.
    content: str

    assessment_id: uuid.UUID = Field(foreign_key="assessment.id", index=True)
    assessment: Assessment = Relationship(back_populates="log_entries")


class Result(SQLModel, table=True):
    id: uuid.UUID | None = Field(default_factory=uuid.uuid4, primary_key=True)
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), server_default=func.now(), index=True
        ),
    )

    assessment_id: uuid.UUID = Field(
        foreign_key="assessment.id", index=True, primary_key=True
    )
    agent_id: uuid.UUID = Field(foreign_key="agent.id", index=True, primary_key=True)
    metric_name: str = Field(index=True, primary_key=True)

    value_type: str = Field(index=True)  # "numeric" or "text"
    float_value: float | None
    int_value: int | None
    text_value: str | None


sqlite_file_name = "database.db"
sqlite_url = f"sqlite:///{sqlite_file_name}"


connect_args = {"check_same_thread": False}
engine = create_engine(sqlite_url, connect_args=connect_args)


def create_db_and_tables():
    SQLModel.metadata.create_all(engine)


def get_session():
    with Session(engine) as session:
        yield session


SessionDep = Annotated[Session, Depends(get_session)]
