import uuid
from datetime import datetime, timezone
from typing import Optional

from sqlalchemy import (
    JSON,
    Column,
    DateTime,
    Float,
    ForeignKey,
    Integer,
    String,
    Text,
    UniqueConstraint,
)
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import relationship

from app.db import Base


def now_utc() -> datetime:
    return datetime.now(timezone.utc)


class Sample(Base):
    __tablename__ = "samples"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    repo = Column(String(255), nullable=False)
    commit_sha = Column(String(40), nullable=False)
    language = Column(String(50), nullable=False)
    diff = Column(Text, nullable=False)
    before_context = Column(JSON, nullable=True)
    after_context = Column(JSON, nullable=True)
    created_at = Column(DateTime, default=now_utc, nullable=False)

    defects = relationship("Defect", back_populates="sample", cascade="all, delete-orphan")
    runs = relationship("ExperimentRun", back_populates="sample")


class Defect(Base):
    __tablename__ = "defects"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    sample_id = Column(UUID(as_uuid=True), ForeignKey("samples.id"), nullable=False)
    defect_type = Column(String(100), nullable=False)
    language = Column(String(50), nullable=False)
    line_start = Column(Integer, nullable=True)
    line_end = Column(Integer, nullable=True)
    description = Column(Text, nullable=True)
    reference_fix = Column(Text, nullable=True)
    created_at = Column(DateTime, default=now_utc, nullable=False)

    sample = relationship("Sample", back_populates="defects")


class PromptTemplate(Base):
    __tablename__ = "prompt_templates"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    strategy_id = Column(String(100), nullable=False)
    level = Column(String(20), nullable=False)
    created_at = Column(DateTime, default=now_utc, nullable=False)

    versions = relationship(
        "PromptTemplateVersion",
        back_populates="template",
        cascade="all, delete-orphan",
        order_by="PromptTemplateVersion.version_number",
    )

    __table_args__ = (UniqueConstraint("strategy_id", "level", name="uix_strategy_level"),)


class PromptTemplateVersion(Base):
    __tablename__ = "prompt_template_versions"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    template_id = Column(UUID(as_uuid=True), ForeignKey("prompt_templates.id"), nullable=False)
    version_number = Column(Integer, nullable=False)
    body = Column(Text, nullable=False)
    variables_schema = Column(JSON, nullable=False, default=dict)
    created_at = Column(DateTime, default=now_utc, nullable=False)

    template = relationship("PromptTemplate", back_populates="versions")
    runs = relationship("ExperimentRun", back_populates="template_version")


class Experiment(Base):
    __tablename__ = "experiments"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    name = Column(String(255), nullable=False)
    models = Column(JSON, nullable=False)
    levels = Column(JSON, nullable=False)
    sample_ids = Column(JSON, nullable=False)
    repeats = Column(Integer, nullable=False, default=3)
    sampling_params = Column(JSON, nullable=False, default=dict)
    status = Column(String(20), default="pending", nullable=False)
    created_at = Column(DateTime, default=now_utc, nullable=False)

    runs = relationship("ExperimentRun", back_populates="experiment", cascade="all, delete-orphan")


class ExperimentRun(Base):
    __tablename__ = "experiment_runs"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    run_id = Column(String(64), unique=True, nullable=False, default=lambda: str(uuid.uuid4()))
    experiment_id = Column(UUID(as_uuid=True), ForeignKey("experiments.id"), nullable=False)
    sample_id = Column(UUID(as_uuid=True), ForeignKey("samples.id"), nullable=False)
    model_id = Column(String(100), nullable=False)
    template_version_id = Column(UUID(as_uuid=True), ForeignKey("prompt_template_versions.id"), nullable=False)
    repeat_index = Column(Integer, nullable=False)
    status = Column(String(20), default="pending", nullable=False)
    retry_count = Column(Integer, default=0, nullable=False)
    sampling_params = Column(JSON, nullable=False, default=dict)
    started_at = Column(DateTime, nullable=True)
    completed_at = Column(DateTime, nullable=True)
    created_at = Column(DateTime, default=now_utc, nullable=False)

    experiment = relationship("Experiment", back_populates="runs")
    sample = relationship("Sample", back_populates="runs")
    template_version = relationship("PromptTemplateVersion", back_populates="runs")
    result = relationship("Result", back_populates="run", uselist=False, cascade="all, delete-orphan")


class Result(Base):
    __tablename__ = "results"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    run_id = Column(UUID(as_uuid=True), ForeignKey("experiment_runs.id"), nullable=False, unique=True)
    raw_output = Column(Text, nullable=True)
    token_usage = Column(JSON, nullable=True)
    latency_ms = Column(Float, nullable=True)
    parsed_findings = Column(JSON, nullable=True)
    detection_rate = Column(Float, nullable=True)
    false_positive_rate = Column(Float, nullable=True)
    coverage_rate = Column(Float, nullable=True)
    stability_score = Column(Float, nullable=True)
    likert_score = Column(Integer, nullable=True)
    created_at = Column(DateTime, default=now_utc, nullable=False)
    updated_at = Column(DateTime, default=now_utc, onupdate=now_utc, nullable=False)

    run = relationship("ExperimentRun", back_populates="result")