# Trouve-KA — dépôts PostgreSQL
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai

"""Dépôts PostgreSQL de Trouve-KA.

Le frontier vit dans Postgres : les workers réclament des lots d'URLs via
FOR UPDATE SKIP LOCKED, ce qui permet N workers concurrents sans coordination
externe (CLAUDE.md §13 : pas de complexité distribuée prématurée).
"""

from datetime import UTC, datetime, timedelta
from typing import Any

import asyncpg


def _dsn(url: str) -> str:
    # asyncpg accepte postgresql:// mais pas postgresql+asyncpg://
    return url.replace("postgresql+asyncpg://", "postgresql://")


class Database:
    def __init__(self, database_url: str, *, pool_min: int = 2, pool_max: int = 10):
        self._url = _dsn(database_url)
        self._pool_min = pool_min
        self._pool_max = pool_max
        self.pool: asyncpg.Pool | None = None

    async def connect(self) -> None:
        self.pool = await asyncpg.create_pool(
            self._url, min_size=self._pool_min, max_size=self._pool_max, command_timeout=30
        )

    async def close(self) -> None:
        if self.pool:
            await self.pool.close()
            self.pool = None

    # ------------------------------------------------------------------ domaines

    async def upsert_domain(self, domain: str, *, is_seed: bool = False) -> int:
        row = await self.pool.fetchrow(
            """
            INSERT INTO domains (domain, is_seed) VALUES ($1, $2)
            ON CONFLICT (domain) DO UPDATE SET is_seed = domains.is_seed OR EXCLUDED.is_seed
            RETURNING id
            """,
            domain,
            is_seed,
        )
        return row["id"]

    async def get_domain(self, domain: str) -> asyncpg.Record | None:
        return await self.pool.fetchrow("SELECT * FROM domains WHERE domain = $1", domain)

    async def get_domain_by_id(self, domain_id: int) -> asyncpg.Record | None:
        return await self.pool.fetchrow("SELECT * FROM domains WHERE id = $1", domain_id)

    async def save_robots(
        self, domain_id: int, body: str | None, status: str, crawl_delay_ms: int | None
    ) -> None:
        await self.pool.execute(
            """
            UPDATE domains SET robots_body = $2, robots_status = $3,
                   robots_fetched_at = now(), crawl_delay_ms = $4
            WHERE id = $1
            """,
            domain_id,
            body,
            status,
            crawl_delay_ms,
        )

    async def block_domain(self, domain: str) -> None:
        await self.pool.execute(
            "UPDATE domains SET blocked = TRUE WHERE domain = $1", domain
        )
        await self.pool.execute(
            """
            UPDATE frontier_items SET status = 'blocked'
            WHERE status IN ('pending', 'in_progress')
              AND url_id IN (SELECT u.id FROM urls u JOIN domains d ON d.id = u.domain_id
                             WHERE d.domain = $1)
            """,
            domain,
        )

    async def update_domain_after_page(
        self, domain_id: int, page_quebec_score: float, language: str | None
    ) -> None:
        """Met à jour le score Québec du domaine (moyenne mobile) et ses stats de langue."""
        lang = language or "unknown"
        await self.pool.execute(
            """
            UPDATE domains SET
                page_count = page_count + 1,
                last_crawled_at = now(),
                quebec_score = (quebec_score * page_count + $2) / (page_count + 1),
                language_stats = jsonb_set(
                    language_stats, ARRAY[$3],
                    (COALESCE(language_stats->>$3, '0')::int + 1)::text::jsonb)
            WHERE id = $1
            """,
            domain_id,
            page_quebec_score,
            lang,
        )

    async def record_domain_links(self, from_domain_id: int, to_domain_counts: dict[int, int]) -> None:
        if not to_domain_counts:
            return
        await self.pool.executemany(
            """
            INSERT INTO domain_links (from_domain_id, to_domain_id, link_count)
            VALUES ($1, $2, $3)
            ON CONFLICT (from_domain_id, to_domain_id)
            DO UPDATE SET link_count = domain_links.link_count + EXCLUDED.link_count,
                          updated_at = now()
            """,
            [(from_domain_id, to_id, n) for to_id, n in to_domain_counts.items() if to_id != from_domain_id],
        )

    # ------------------------------------------------------------------ frontier

    async def enqueue_url(
        self,
        url: str,
        domain: str,
        *,
        priority: float = 0.5,
        depth: int = 0,
        source_url_id: int | None = None,
        is_seed: bool = False,
        max_urls_per_domain: int = 5000,
    ) -> int | None:
        """Ajoute une URL au frontier si inconnue. Retourne url_id si ajoutée, None sinon."""
        async with self.pool.acquire() as conn:
            async with conn.transaction():
                domain_id = await conn.fetchval(
                    """
                    INSERT INTO domains (domain, is_seed) VALUES ($1, $2)
                    ON CONFLICT (domain) DO UPDATE SET is_seed = domains.is_seed OR EXCLUDED.is_seed
                    RETURNING id
                    """,
                    domain,
                    is_seed,
                )
                dom = await conn.fetchrow(
                    "SELECT blocked, page_count FROM domains WHERE id = $1", domain_id
                )
                if dom["blocked"]:
                    return None
                if not is_seed and dom["page_count"] >= max_urls_per_domain:
                    return None
                url_id = await conn.fetchval(
                    """
                    INSERT INTO urls (url, domain_id) VALUES ($1, $2)
                    ON CONFLICT (url) DO NOTHING RETURNING id
                    """,
                    url,
                    domain_id,
                )
                if url_id is None:
                    return None  # URL déjà connue : pas de doublon dans le frontier
                await conn.execute(
                    """
                    INSERT INTO frontier_items (url_id, priority, depth, source_url_id)
                    VALUES ($1, $2, $3, $4) ON CONFLICT (url_id) DO NOTHING
                    """,
                    url_id,
                    priority,
                    depth,
                    source_url_id,
                )
                return url_id

    async def claim_batch(self, worker_id: str, limit: int = 20) -> list[asyncpg.Record]:
        """Réclame un lot d'URLs prêtes à crawler (SKIP LOCKED, multi-workers sûr)."""
        return await self.pool.fetch(
            """
            WITH claimed AS (
                SELECT f.id FROM frontier_items f
                WHERE f.status = 'pending' AND f.next_crawl_at <= now()
                ORDER BY f.priority DESC, f.next_crawl_at
                LIMIT $2
                FOR UPDATE SKIP LOCKED
            )
            UPDATE frontier_items f SET status = 'in_progress', locked_by = $1, locked_at = now()
            FROM claimed
            WHERE f.id = claimed.id
            RETURNING f.id, f.url_id, f.priority, f.depth, f.retries,
                      f.last_crawled_at, f.next_crawl_at,
                      (SELECT url FROM urls WHERE id = f.url_id) AS url,
                      (SELECT domain_id FROM urls WHERE id = f.url_id) AS domain_id
            """,
            worker_id,
            limit,
        )

    async def release_item(
        self,
        url_id: int,
        *,
        status: str,
        error_code: str | None = None,
        next_crawl_at: datetime | None = None,
        priority: float | None = None,
        increment_retries: bool = False,
    ) -> None:
        await self.pool.execute(
            """
            UPDATE frontier_items SET
                status = $2,
                error_code = $3,
                last_crawled_at = now(),
                next_crawl_at = COALESCE($4, next_crawl_at),
                priority = COALESCE($5, priority),
                retries = retries + CASE WHEN $6 THEN 1 ELSE 0 END,
                locked_by = NULL, locked_at = NULL
            WHERE url_id = $1
            """,
            url_id,
            status,
            error_code,
            next_crawl_at,
            priority,
            increment_retries,
        )

    async def reset_stale_items(self, older_than_minutes: int = 30) -> int:
        """Relance les items in_progress abandonnés (worker mort). Retourne le nombre."""
        result = await self.pool.execute(
            """
            UPDATE frontier_items SET status = 'pending', locked_by = NULL, locked_at = NULL
            WHERE status = 'in_progress' AND locked_at < now() - ($1 || ' minutes')::interval
            """,
            str(older_than_minutes),
        )
        return int(result.split()[-1])

    async def requeue_url(self, url: str) -> bool:
        result = await self.pool.execute(
            """
            UPDATE frontier_items SET status = 'pending', next_crawl_at = now(), retries = 0
            WHERE url_id = (SELECT id FROM urls WHERE url = $1)
            """,
            url,
        )
        return result.endswith("1")

    async def requeue_domain(self, domain: str) -> int:
        result = await self.pool.execute(
            """
            UPDATE frontier_items SET status = 'pending', next_crawl_at = now()
            WHERE status IN ('done', 'failed')
              AND url_id IN (SELECT u.id FROM urls u JOIN domains d ON d.id = u.domain_id
                             WHERE d.domain = $1)
            """,
            domain,
        )
        return int(result.split()[-1])

    # ------------------------------------------------------------------ crawl / documents

    async def record_attempt(
        self,
        url_id: int,
        *,
        status_code: int | None,
        error_code: str,
        outcome: str,
        content_hash: str | None = None,
        num_bytes: int | None = None,
        duration_ms: int | None = None,
        redirect_url: str | None = None,
        title: str | None = None,
        quebec_score: float | None = None,
    ) -> None:
        await self.pool.execute(
            """
            INSERT INTO crawl_attempts
                (url_id, status_code, error_code, outcome, content_hash, bytes,
                 duration_ms, redirect_url, title, quebec_score)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
            """,
            url_id,
            status_code,
            error_code,
            outcome,
            content_hash,
            num_bytes,
            duration_ms,
            redirect_url,
            title,
            quebec_score,
        )

    async def get_document(self, url_id: int) -> asyncpg.Record | None:
        return await self.pool.fetchrow("SELECT * FROM documents WHERE url_id = $1", url_id)

    async def find_duplicate(self, chash: str, url_id: int) -> asyncpg.Record | None:
        """Doublon exact : même hash de contenu sur une autre URL."""
        return await self.pool.fetchrow(
            "SELECT * FROM documents WHERE content_hash = $1 AND url_id != $2 LIMIT 1",
            chash,
            url_id,
        )

    async def upsert_document(
        self,
        url_id: int,
        *,
        content_hash: str,
        etag: str | None,
        last_modified: str | None,
        title: str,
        description: str,
        language: str | None,
        page_quebec_score: float,
        published_at: datetime | None,
        changed: bool,
    ) -> int:
        row = await self.pool.fetchrow(
            """
            INSERT INTO documents
                (url_id, content_hash, etag, last_modified, title, description, language,
                 page_quebec_score, published_at, last_changed_at)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, now())
            ON CONFLICT (url_id) DO UPDATE SET
                content_hash = EXCLUDED.content_hash,
                etag = EXCLUDED.etag,
                last_modified = EXCLUDED.last_modified,
                title = EXCLUDED.title,
                description = EXCLUDED.description,
                language = EXCLUDED.language,
                page_quebec_score = EXCLUDED.page_quebec_score,
                published_at = EXCLUDED.published_at,
                last_indexed_at = now(),
                last_changed_at = CASE WHEN $10 THEN now() ELSE documents.last_changed_at END
            RETURNING id
            """,
            url_id,
            content_hash,
            etag,
            last_modified,
            title,
            description,
            language,
            page_quebec_score,
            published_at,
            changed,
        )
        return row["id"]

    async def set_canonical(self, url_id: int, canonical_url: str) -> None:
        await self.pool.execute(
            "UPDATE urls SET canonical_url = $2 WHERE id = $1", url_id, canonical_url
        )

    # ------------------------------------------------------------------ analytics / statut

    async def record_search_query(
        self, query: str, language: str | None, total: int, took_ms: int
    ) -> None:
        await self.pool.execute(
            """
            INSERT INTO search_queries (query, language, results_total, took_ms, zero_result)
            VALUES ($1, $2, $3, $4, $3 = 0)
            """,
            query[:500],
            language,
            total,
            took_ms,
        )

    async def add_submission(self, url: str) -> None:
        await self.pool.execute("INSERT INTO submissions (url) VALUES ($1)", url)

    async def status_snapshot(self) -> dict[str, Any]:
        since = datetime.now(UTC) - timedelta(hours=1)
        frontier = {
            r["status"]: r["n"]
            for r in await self.pool.fetch(
                "SELECT status, count(*)::int AS n FROM frontier_items GROUP BY status"
            )
        }
        row = await self.pool.fetchrow(
            """
            SELECT
              (SELECT count(*)::int FROM documents) AS docs,
              (SELECT count(*)::int FROM domains WHERE page_count > 0) AS domains,
              (SELECT count(*)::int FROM documents WHERE last_indexed_at >= $1) AS indexed_1h,
              (SELECT count(*)::int FROM crawl_attempts WHERE fetched_at >= $1) AS fetched_1h,
              (SELECT count(*)::int FROM crawl_attempts
                 WHERE fetched_at >= $1 AND outcome = 'error') AS errors_1h
            """,
            since,
        )
        return {
            "pages_indexed": row["docs"],
            "domains_count": row["domains"],
            "indexed_last_hour": row["indexed_1h"],
            "fetched_last_hour": row["fetched_1h"],
            "errors_last_hour": row["errors_1h"],
            "frontier_pending": frontier.get("pending", 0),
            "frontier_in_progress": frontier.get("in_progress", 0),
            "frontier": frontier,
        }

    async def admin_overview(self) -> dict[str, Any]:
        since = datetime.now(UTC) - timedelta(hours=1)
        base = await self.status_snapshot()
        http_status = {
            str(r["status_code"]): r["n"]
            for r in await self.pool.fetch(
                """
                SELECT status_code, count(*)::int AS n FROM crawl_attempts
                WHERE fetched_at >= $1 AND status_code IS NOT NULL
                GROUP BY status_code ORDER BY n DESC
                """,
                since,
            )
        }
        outcomes = {
            r["outcome"]: r["n"]
            for r in await self.pool.fetch(
                "SELECT outcome, count(*)::int AS n FROM crawl_attempts WHERE fetched_at >= $1 GROUP BY outcome",
                since,
            )
        }
        top_domains = [
            dict(r)
            for r in await self.pool.fetch(
                """
                SELECT domain, page_count AS pages, round(quebec_score::numeric, 3)::float AS quebec_score
                FROM domains WHERE page_count > 0
                ORDER BY page_count DESC LIMIT 20
                """
            )
        ]
        recent_errors = [
            {"url": r["url"], "error_code": r["error_code"], "at": r["fetched_at"].isoformat()}
            for r in await self.pool.fetch(
                """
                SELECT u.url, a.error_code, a.fetched_at
                FROM crawl_attempts a JOIN urls u ON u.id = a.url_id
                WHERE a.outcome = 'error' ORDER BY a.fetched_at DESC LIMIT 20
                """
            )
        ]
        latency = await self.pool.fetchrow(
            """
            SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY duration_ms) AS p50,
                   percentile_cont(0.95) WITHIN GROUP (ORDER BY duration_ms) AS p95
            FROM crawl_attempts WHERE fetched_at >= $1 AND duration_ms IS NOT NULL
            """,
            since,
        )
        return {
            "frontier": {
                "pending": base["frontier"].get("pending", 0),
                "in_progress": base["frontier"].get("in_progress", 0),
                "done": base["frontier"].get("done", 0),
                "failed": base["frontier"].get("failed", 0),
                "blocked": base["frontier"].get("blocked", 0),
            },
            "rates": {
                "fetched_1h": base["fetched_last_hour"],
                "indexed_1h": outcomes.get("indexed", 0),
                "errors_1h": outcomes.get("error", 0),
                "robots_blocked_1h": outcomes.get("robots_blocked", 0),
                "duplicates_1h": outcomes.get("duplicate", 0),
                "parsed_1h": outcomes.get("indexed", 0) + outcomes.get("not_quebec", 0),
            },
            "http_status": http_status,
            "top_domains": top_domains,
            "recent_errors": recent_errors,
            "latency": {
                "fetch_p50_ms": round(latency["p50"]) if latency and latency["p50"] else None,
                "fetch_p95_ms": round(latency["p95"]) if latency and latency["p95"] else None,
            },
        }

    async def recent_events(self, limit: int = 50) -> list[dict[str, Any]]:
        rows = await self.pool.fetch(
            """
            SELECT a.fetched_at, u.url, a.status_code, a.outcome, a.quebec_score, a.title
            FROM crawl_attempts a JOIN urls u ON u.id = a.url_id
            ORDER BY a.fetched_at DESC LIMIT $1
            """,
            min(limit, 200),
        )
        return [
            {
                "at": r["fetched_at"].isoformat(),
                "url": r["url"],
                "status": r["status_code"],
                "outcome": r["outcome"],
                "quebec_score": r["quebec_score"],
                "title": r["title"],
            }
            for r in rows
        ]

    async def frontier_inspect(
        self, domain: str | None = None, status: str | None = None, limit: int = 100
    ) -> list[dict[str, Any]]:
        rows = await self.pool.fetch(
            """
            SELECT u.url, f.priority, f.depth, f.status, f.retries, f.next_crawl_at, f.error_code
            FROM frontier_items f JOIN urls u ON u.id = f.url_id
            JOIN domains d ON d.id = u.domain_id
            WHERE ($1::text IS NULL OR d.domain = $1)
              AND ($2::text IS NULL OR f.status = $2)
            ORDER BY f.priority DESC LIMIT $3
            """,
            domain,
            status,
            min(limit, 500),
        )
        return [dict(r) | {"next_crawl_at": r["next_crawl_at"].isoformat()} for r in rows]