import asyncio
import os
from typing import Any, Optional
import asyncpg

DATABASE_URL = os.getenv(
    "DATABASE_URL",
    "postgresql://asp_user:asp_password@localhost:5432/asp",
)

PIPE_SCHEMA = """
CREATE TABLE IF NOT EXISTS pipeline_steps (
    case_id     TEXT NOT NULL,
    step        TEXT NOT NULL,
    status      TEXT NOT NULL CHECK (status IN ('partial', 'complete')),
    content     TEXT NOT NULL,
    updated_at  TIMESTAMPTZ NOT NULL DEFAULT now(),
    PRIMARY KEY (case_id, step)
);
"""


class CheckpointDB:
    def __init__(
        self,
        url: str = DATABASE_URL,
        min_size: int = 1,
        max_size: int = 10,
        timeout: float = 5,
        command_timeout: float = 30,
    ):
        self.url = url
        self.min_size = min_size
        self.max_size = max_size
        self.timeout = timeout
        self.command_timeout = command_timeout
        self._pool: Optional[asyncpg.Pool] = None
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._lock: Optional[asyncio.Lock] = None
        self._lock_loop: Optional[asyncio.AbstractEventLoop] = None

    def lock_for(self, loop: asyncio.AbstractEventLoop) -> asyncio.Lock:
        if self._lock is None or self._lock_loop is not loop:
            self._lock, self._lock_loop = asyncio.Lock(), loop
        return self._lock

    async def pool(self) -> asyncpg.Pool:
        loop = asyncio.get_running_loop()
        if self._pool is not None and self._loop is loop:
            return self._pool
        async with self.lock_for(loop):
            if self._pool is None or self._loop is not loop:
                self.discard_pool()
                pool = await asyncpg.create_pool(
                    self.url,
                    min_size=self.min_size,
                    max_size=self.max_size,
                    timeout=self.timeout,
                    command_timeout=self.command_timeout,
                )
                async with pool.acquire() as conn:
                    await conn.execute(PIPE_SCHEMA)
                self._pool, self._loop = pool, loop
            return self._pool

    def discard_pool(self) -> None:
        stale, self._pool, self._loop = self._pool, None, None
        if stale is not None:
            try:
                stale.terminate()
            except Exception:
                pass

    async def close(self) -> None:
        pool, self._pool, self._loop = self._pool, None, None
        if pool is not None:
            await pool.close()

    async def save_step(self, case_id: str, step: str, content: str, status: str) -> bool:
        pool = await self.pool()
        async with pool.acquire() as conn:
            row = await conn.fetchval(
                """
                INSERT INTO pipeline_steps (case_id, step, status, content, updated_at)
                VALUES ($1, $2, $3, $4, now())
                ON CONFLICT (case_id, step) DO UPDATE
                   SET status = EXCLUDED.status,
                       content = EXCLUDED.content,
                       updated_at = now()
                 WHERE NOT (pipeline_steps.status = 'complete' AND EXCLUDED.status = 'partial')
                RETURNING 1
                """,
                case_id, step, status, content
            )
        return row is not None

    async def load_steps(self, case_id: str) -> dict[str, dict[str, Any]]:
        pool = await self.pool()
        async with pool.acquire() as conn:
            rows = await conn.fetch(
                "SELECT step, status, content, updated_at FROM pipeline_steps WHERE case_id = $1",
                case_id
            )
        return {
            r["step"]: {
                "status": r["status"],
                "content": r["content"],
                "updated_at": r["updated_at"].isoformat()
            }
            for r in rows
        }

    async def delete_steps(self, case_id: str, steps: list[str]) -> None:
        pool = await self.pool()
        async with pool.acquire() as conn:
            await conn.execute(
                "DELETE FROM pipeline_steps WHERE case_id = $1 AND step = ANY($2::text[])",
                case_id, steps
            )

    async def list_cases(self, limit: int = 25) -> list[dict[str, Any]]:
        pool = await self.pool()
        async with pool.acquire() as conn:
            rows = await conn.fetch(
                """
                SELECT case_id,
                       max(updated_at) AS updated_at,
                       string_agg(step || ':' || status, ',' ORDER BY step) AS steps,
                       max(content) FILTER (WHERE step = 'meta') AS meta
                  FROM pipeline_steps
                 GROUP BY case_id
                 ORDER BY max(updated_at) DESC
                 LIMIT $1
                """,
                limit
            )
        return [
            {
                "case_id": r["case_id"],
                "updated_at": r["updated_at"].isoformat(),
                "steps": r["steps"],
                "meta": r["meta"]
            }
            for r in rows
        ]

    async def forget_case(self, case_id: str) -> int:
        pool = await self.pool()
        async with pool.acquire() as conn:
            n = await conn.fetchval(
                "WITH d AS (DELETE FROM pipeline_steps WHERE case_id = $1 RETURNING 1) SELECT count(*) FROM d",
                case_id
            )
        return int(n or 0)

    async def delete_cases_older_than(self, days: int) -> int:
        pool = await self.pool()
        async with pool.acquire() as conn:
            ids = [
                r["case_id"]
                for r in await conn.fetch(
                    """
                    DELETE FROM pipeline_steps WHERE case_id IN (
                        SELECT case_id FROM pipeline_steps
                         GROUP BY case_id
                        HAVING max(updated_at) < now() - make_interval(days => $1)
                    ) RETURNING case_id
                    """,
                    int(days)
                )
            ]
        return len(set(ids))

    async def delete_all_cases(self) -> int:
        pool = await self.pool()
        async with pool.acquire() as conn:
            rows = await conn.fetch("DELETE FROM pipeline_steps RETURNING case_id")
        return len({r["case_id"] for r in rows})


checkpoint_db = CheckpointDB()