from sqlalchemy import create_engine, Column, String, Enum, func
from sqlalchemy.orm import sessionmaker, declarative_base
import uuid
import hashlib
from typing import Literal

# Define the base for ORM mapping
Base = declarative_base()

# Define the Images table as a Python class
class Image(Base):
    __tablename__ = 'images'

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    prompt = Column(String, nullable=False)
    model = Column(Enum("dall-e-3", "dall-e-2", name="model_enum"), nullable=False)
    filename = Column(String, nullable=False)
    prompt_hash = Column(String, nullable=False)
    noun = Column(String, nullable=False)

# Create a database engine and session factory
def create_db(db_path: str):
    engine = create_engine(f'sqlite:///{db_path}')
    Base.metadata.create_all(engine)
    return engine

def hash_image_row(
        prompt: str,
        model: str,
        size: str,
        quality: str,
) -> str:
    # Create a unique hash from the concatenated parameters
    concat_string = f"{prompt}:{model}:{size}:{quality}"
    unique_hash = hashlib.sha256(concat_string.encode()).hexdigest()
    return unique_hash


def insert_image_row(
        session,
        image_id: str,
        prompt: str,
        hash: str,
        model: Literal["dall-e-3", "dall-e-2"] = "dall-e-3",
        size: Literal["1024x1024"] = "1024x1024",
        quality: Literal["standard", "hd"] = "standard",
):
    image = Image(
        id=image_id,
        prompt=prompt,
        filename=f"{image_id}.png",
        hash=hash
    )
    session.add(image)
    session.commit()


def count_hash(session, hash_value: str):
    return session.query(func.count(Image.hash)).filter(Image.hash == hash_value).scalar()


if __name__ == "__main__":
    # Example usage
    db_path = "../output/images.db"
    engine = create_db(db_path)
    Session = sessionmaker(bind=engine)
    session = Session()

    prompt = "A futuristic cityscape"
    model =  "dall-e-3"
    size = "1024x1024"
    quality = "standard"
    hash = hash_image_row(prompt, model, size, quality)
    # Insert an example record
    insert_image_row(
        session,
        image_id=str(uuid.uuid4()),
        prompt=prompt,
        model=model,
        size=size,
        quality="standard",
        hash=hash
    )

    session.close()

