from __future__ import annotations

import enum
import json
from datetime import UTC, datetime
from typing import Any

from sqlalchemy import DateTime, Float, ForeignKey, Integer, String, Text, UniqueConstraint
from sqlalchemy import Enum as SAEnum
from sqlalchemy.orm import Mapped, mapped_column, relationship

from evalkit.server.db.base import Base


class EvaluationStatus(enum.StrEnum):
    pending = "pending"
    running = "running"
    complete = "complete"
    failed = "failed"


class ProjectModel(Base):
    __tablename__ = "projects"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    name: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
    description: Mapped[str] = mapped_column(Text, default="")
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))

    evaluations: Mapped[list[EvaluationModel]] = relationship(back_populates="project", cascade="all, delete-orphan")
    traces: Mapped[list[TraceModel]] = relationship(back_populates="project", cascade="all, delete-orphan")


class EvaluationModel(Base):
    __tablename__ = "evaluations"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    project_id: Mapped[int] = mapped_column(ForeignKey("projects.id"), nullable=False)
    suite_name: Mapped[str] = mapped_column(String(100), nullable=False)
    status: Mapped[EvaluationStatus] = mapped_column(SAEnum(EvaluationStatus), default=EvaluationStatus.pending)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    average_score: Mapped[float | None] = mapped_column(Float, nullable=True)

    project: Mapped[ProjectModel] = relationship(back_populates="evaluations")
    results: Mapped[list[EvalResultModel]] = relationship(back_populates="evaluation", cascade="all, delete-orphan")
    tags: Mapped[list[EvaluationTagModel]] = relationship(back_populates="evaluation", cascade="all, delete-orphan")


class EvalResultModel(Base):
    __tablename__ = "eval_results"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    evaluation_id: Mapped[int] = mapped_column(ForeignKey("evaluations.id"), nullable=False)
    test_case_json: Mapped[str] = mapped_column(Text, nullable=False)
    overall_score: Mapped[float] = mapped_column(Float, nullable=False)

    evaluation: Mapped[EvaluationModel] = relationship(back_populates="results")
    metric_scores: Mapped[list[MetricScoreModel]] = relationship(
        back_populates="eval_result", cascade="all, delete-orphan"
    )

    def get_test_case_dict(self) -> dict[str, Any]:
        return json.loads(self.test_case_json)


class MetricScoreModel(Base):
    __tablename__ = "metric_scores"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    eval_result_id: Mapped[int] = mapped_column(ForeignKey("eval_results.id"), nullable=False)
    metric_name: Mapped[str] = mapped_column(String(100), nullable=False)
    score: Mapped[float] = mapped_column(Float, nullable=False)
    passed: Mapped[bool] = mapped_column(Integer, nullable=False)  # SQLite stores bools as int
    reason: Mapped[str] = mapped_column(Text, default="")

    eval_result: Mapped[EvalResultModel] = relationship(back_populates="metric_scores")


class EvaluationTagModel(Base):
    __tablename__ = "evaluation_tags"
    __table_args__ = (UniqueConstraint("evaluation_id", "key", name="uq_eval_tag_key"),)

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    evaluation_id: Mapped[int] = mapped_column(ForeignKey("evaluations.id"), nullable=False)
    key: Mapped[str] = mapped_column(String(100), nullable=False)
    value: Mapped[str] = mapped_column(String(500), nullable=False)

    evaluation: Mapped[EvaluationModel] = relationship(back_populates="tags")


class TraceModel(Base):
    __tablename__ = "traces"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    project_id: Mapped[int] = mapped_column(ForeignKey("projects.id"), nullable=False)
    timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    test_case_json: Mapped[str] = mapped_column(Text, nullable=False)
    metadata_json: Mapped[str] = mapped_column(Text, default="{}")
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))

    project: Mapped[ProjectModel] = relationship(back_populates="traces")
    scores: Mapped[list[TraceScoreModel]] = relationship(back_populates="trace", cascade="all, delete-orphan")

    def get_test_case(self) -> dict[str, Any]:
        return json.loads(self.test_case_json)

    def get_metadata(self) -> dict[str, str]:
        return json.loads(self.metadata_json)


class TraceScoreModel(Base):
    __tablename__ = "trace_scores"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    trace_id: Mapped[int] = mapped_column(ForeignKey("traces.id"), nullable=False)
    metric_name: Mapped[str] = mapped_column(String(100), nullable=False)
    score: Mapped[float] = mapped_column(Float, nullable=False)
    passed: Mapped[bool] = mapped_column(Integer, nullable=False)
    reason: Mapped[str] = mapped_column(Text, default="")

    trace: Mapped[TraceModel] = relationship(back_populates="scores")
