import uuid
from typing import Annotated, Literal, Any, cast
from datetime import datetime, timezone
from fastapi import Depends 
from functools import lru_cache
from pydantic import BaseModel
from sqlmodel import DateTime, Field, SQLModel, select
from sqlmodel.ext.asyncio.session import AsyncSession
from sqlmodel.main import SQLModelMetaclass
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from typing_extensions import Self

from maur.settings import MysqlSettings, PostgresSettings


@lru_cache
def engine() -> AsyncEngine:
    sqlalchemy_uri = None

    if settings := MysqlSettings.from_env_or_none():
        sqlalchemy_uri = settings.db_uri.encoded_string()

    if settings := PostgresSettings.from_env_or_none():
        sqlalchemy_uri = settings.db_uri.encoded_string().replace("postgresql://", "postgresql+asyncpg://")

    if sqlalchemy_uri is None:
        raise ValueError("Found no database settings. Takk should set this for you, but remember to set the DB_URI.")

    return create_async_engine(sqlalchemy_uri)

async def session():
    async with AsyncSession(engine()) as session:
        yield session
        await session.commit()


SessionDep = Annotated[AsyncSession, Depends(session)]



class Model(SQLModel, metaclass=SQLModelMetaclass):
    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)

    created_at: datetime = Field(default_factory=lambda: datetime.now(tz=timezone.utc), sa_type=cast(type[Any], DateTime(timezone=True)))
    updated_at: datetime = Field(default_factory=lambda: datetime.now(tz=timezone.utc), sa_type=cast(type[Any], DateTime(timezone=True)))
    deleted_at: datetime | None = Field(default=None, sa_type=cast(type[Any], DateTime(timezone=True)))


    @classmethod
    async def get(cls, id: uuid.UUID, session: AsyncSession) -> Self | None:
        query = await session.exec(
            select(cls).where(cls.id == id, cls.deleted_at == None)
        )
        value = query.first()
        if value is None:
            return value

        return value

    async def delete(self, session: AsyncSession, is_soft: bool = True) -> None:
        if is_soft:
            self.deleted_at = datetime.now(tz=timezone.utc)
            session.add(self)
            await session.flush()
        else:
            await session.delete(self)


    async def insert(self, session: AsyncSession) -> Self:
        session.add(self)
        await session.flush()
        return self

    async def update(self, session: AsyncSession, new_values: BaseModel | dict[str, Any] | None = None) -> Self:
        if new_values:
            self.sqlmodel_update(new_values)
        self.updated_at = datetime.now(tz=timezone.utc)
        session.add(self)
        await session.flush()
        return self



CodingTaskStatus = Literal["pending", "in_progress", "completed", "failed"]

class CodingTask(Model, table=True):
    source: str
    source_id: str = Field(index=True)

    repo_branch: str = Field(default="main")

    prompt: str

    result_branch: str | None = Field(default=None)
    result_url: str | None = Field(default=None)

    error: str | None = Field(default=None)

    tokens_input: int = Field(default=0)
    tokens_output: int = Field(default=0)
    tokens_cache_read: int = Field(default=0)
    tokens_cache_write: int = Field(default=0)
    model_used: str | None = Field(default=None)

    started_at: datetime | None = Field(default=None, sa_type=cast(type[Any], DateTime(timezone=True)))
    finnished_at: datetime | None = Field(default=None, sa_type=cast(type[Any], DateTime(timezone=True)))


    @property
    def status(self) -> CodingTaskStatus:
        if self.finnished_at:
            if self.error:
                return "failed"
            else:
                return "completed"

        if self.started_at:
            return "in_progress"

        return "pending"

