from datetime import datetime
from typing import Any

from sqlalchemy import Column
from sqlalchemy.dialects.postgresql import JSONB
from sqlmodel import Field, SQLModel


class MyDoc(SQLModel, table=True):
    __tablename__ = "my_doc"

    id: str = Field(primary_key=True)
    category: str = Field(index=True)
    type: str = Field(index=True)
    provider: str = Field(default="", index=True)

    url: str = ""
    title: str = ""
    summary: str = ""

    content_len: int = 0
    content_s3_bucket: str = ""
    content_s3_key: str = ""

    sentiment_state: str = Field(default="NONE", index=True)
    sentiment_list: list[dict[str, Any]] = Field(
        default_factory=list,
        sa_column=Column(JSONB, nullable=False),
    )
    meta: dict[str, Any] = Field(
        default_factory=dict,
        sa_column=Column(JSONB, nullable=False),
    )

    dt_publish: datetime | None = None
    dt_create: datetime = Field(default_factory=datetime.utcnow)
    dt_update: datetime = Field(default_factory=datetime.utcnow)

import os
from collections.abc import Generator

from sqlmodel import Session, SQLModel, create_engine


DATABASE_URL = os.environ["MY_DB_URL"]

engine = create_engine(
    DATABASE_URL,
    pool_pre_ping=True,
    pool_size=5,
    max_overflow=10,
    echo=False,
)


def create_db_and_tables() -> None:
    SQLModel.metadata.create_all(engine)


def get_session() -> Generator[Session, None, None]:
    with Session(engine) as session:
        yield session

import os

import boto3


AWS_REGION = os.environ.get("MY_AWS_REGION", "ap-northeast-2")
S3_BUCKET = os.environ["MY_S3_BUCKET"]

s3_client = boto3.client("s3", region_name=AWS_REGION)


def put_text_to_s3(*, key: str, text: str, content_type: str = "text/markdown") -> tuple[str, str]:
    s3_client.put_object(
        Bucket=S3_BUCKET,
        Key=key,
        Body=text.encode("utf-8"),
        ContentType=f"{content_type}; charset=utf-8",
    )

    return S3_BUCKET, key

from datetime import datetime

from sqlalchemy.dialects.postgresql import insert
from sqlmodel import Session, select

from models import MyDoc


def get_my_doc(session: Session, doc_id: str) -> MyDoc | None:
    return session.get(MyDoc, doc_id)


def list_pending_docs(session: Session, limit: int = 100) -> list[MyDoc]:
    stmt = (
        select(MyDoc)
        .where(MyDoc.sentiment_state == "NONE")
        .order_by(MyDoc.dt_publish.desc())
        .limit(limit)
    )
    return list(session.exec(stmt).all())


def update_analyzed_doc(
    session: Session,
    doc_id: str,
    *,
    summary: str,
    sentiment_list: list[dict],
) -> MyDoc | None:
    doc = session.get(MyDoc, doc_id)
    if doc is None:
        return None

    doc.summary = summary
    doc.sentiment_list = sentiment_list
    doc.sentiment_state = "ANALYZED"
    doc.dt_update = datetime.utcnow()

    session.add(doc)
    session.commit()
    session.refresh(doc)
    return doc


def upsert_my_doc(session: Session, data: dict) -> MyDoc:
    now = datetime.utcnow()
    data["dt_update"] = now

    stmt = insert(MyDoc).values(**data)
    update_data = {key: value for key, value in data.items() if key not in {"id", "dt_create"}}

    stmt = stmt.on_conflict_do_update(
        index_elements=[MyDoc.id],
        set_=update_data,
    )

    session.exec(stmt)
    session.commit()

    doc = session.get(MyDoc, data["id"])
    if doc is None:
        raise RuntimeError(f"MyDoc upsert failed: {data['id']}")

    return doc