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

from sqlmodel import SQLModel, Field, select

from backend.app.db.connection import get_db_session


class BulkAddressUpload(SQLModel, table=True):
    __tablename__ = "bulk_address_uploads"

    id: UUID = Field(default_factory=uuid4, primary_key=True, index=True)
    portfolio_id: str = Field(nullable=False)
    user_id: str = Field(nullable=False)
    s3_bucket: str = Field(nullable=False)
    s3_key: str = Field(nullable=False)
    filename: str = Field(nullable=False)
    status: str = Field(default="ready_for_processing", nullable=False)
    task_id: Optional[UUID] = Field(default=None)
    combined_output_s3_uri: Optional[str] = Field(default=None)
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)


def set_combined_output_s3_uri(task_id: UUID, s3_uri: str) -> None:
    now = datetime.now(timezone.utc)
    with get_db_session() as session:
        row = session.exec(
            select(BulkAddressUpload).where(BulkAddressUpload.task_id == task_id)
        ).first()
        if not row:
            raise ValueError(f"No bulk_address_uploads row for task_id {task_id}")
        row.combined_output_s3_uri = s3_uri
        row.updated_at = now
        session.add(row)
        session.commit()