# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import json
import logging
import shutil
import sqlite3
from pathlib import Path

from devplacepy.config import DBAPI_DIR
from devplacepy.database import get_int_setting
from devplacepy.services.base import ConfigField
from devplacepy.services.jobs.base import JobService

from .progress import hub
from .validate import _database_path, validate_select

logger = logging.getLogger(__name__)

BATCH_SIZE = 200
DEFAULT_MAX_ROWS = 5000


class DbApiJobService(JobService):
    kind = "dbquery"
    title = "Database API"
    description = (
        "Runs validated read-only SQL SELECT queries off the request path and streams the "
        "result rows over a websocket. Backs the asynchronous /dbapi/query/async endpoint."
    )

    def __init__(self):
        super().__init__(name="dbquery", interval_seconds=2)
        self.config_fields.extend(
            [
                ConfigField(
                    "dbapi_max_rows",
                    "Max result rows",
                    type="int",
                    default=DEFAULT_MAX_ROWS,
                    minimum=1,
                    help="Hard cap on rows returned by a query (sync and async).",
                    group="Database API",
                ),
                ConfigField(
                    "dbapi_nl_model",
                    "NL-to-SQL model",
                    type="str",
                    default="",
                    help="Model used to design SQL from natural language. Blank uses molodetz.",
                    group="Database API",
                ),
                ConfigField(
                    "dbapi_nl_system_preamble",
                    "NL-to-SQL preamble",
                    type="text",
                    default="",
                    help="Optional operator text prepended to the NL-to-SQL system prompt.",
                    group="Database API",
                ),
                ConfigField(
                    "dbapi_deny_tables",
                    "Denied tables",
                    type="str",
                    default="",
                    help="Comma separated extra tables to hide from the database API.",
                    group="Database API",
                ),
            ]
        )

    def result_dir(self, uid: str) -> Path:
        return DBAPI_DIR / uid

    async def process(self, job: dict) -> dict:
        uid = job["uid"]
        payload = job.get("payload", {})
        sql = payload.get("sql", "")
        verdict = validate_select(sql)
        if not verdict.valid:
            hub.publish(uid, {"type": "failed", "message": verdict.error or "invalid query"})
            hub.clear(uid)
            raise RuntimeError(verdict.error or "invalid query")

        max_rows = max(1, get_int_setting("dbapi_max_rows", DEFAULT_MAX_ROWS))
        rows, truncated = await self._stream(uid, verdict.sql, max_rows)

        output_dir = self.result_dir(uid)
        await asyncio.to_thread(output_dir.mkdir, parents=True, exist_ok=True)
        result = {
            "sql": verdict.sql,
            "row_count": len(rows),
            "truncated": truncated,
            "suspicious": verdict.suspicious,
            "rows": rows,
        }
        await asyncio.to_thread(
            (output_dir / "result.json").write_text,
            json.dumps(result, default=str),
            encoding="utf-8",
        )
        hub.publish(
            uid,
            {
                "type": "done",
                "row_count": len(rows),
                "truncated": truncated,
                "result_url": f"/dbapi/query/{uid}/result",
            },
        )
        hub.clear(uid)
        return {
            "sql": verdict.sql,
            "row_count": len(rows),
            "truncated": truncated,
            "suspicious": verdict.suspicious,
            "result_url": f"/dbapi/query/{uid}/result",
            "bytes_out": len(json.dumps(result, default=str)),
            "item_count": len(rows),
        }

    async def _stream(self, uid: str, sql: str, max_rows: int):
        connection = sqlite3.connect(
            f"file:{_database_path()}?mode=ro", uri=True, timeout=30
        )
        rows: list[dict] = []
        truncated = False
        try:
            connection.row_factory = sqlite3.Row
            connection.execute("PRAGMA query_only=ON")
            cursor = connection.execute(sql)
            while True:
                batch = cursor.fetchmany(BATCH_SIZE)
                if not batch:
                    break
                for raw in batch:
                    if len(rows) >= max_rows:
                        truncated = True
                        break
                    rows.append(dict(raw))
                hub.publish(
                    uid, {"type": "progress", "row_count": len(rows)}
                )
                if truncated:
                    break
                await asyncio.sleep(0)
        finally:
            connection.close()
        return rows, truncated

    def cleanup(self, job: dict) -> None:
        hub.clear(job["uid"])
        shutil.rmtree(self.result_dir(job["uid"]), ignore_errors=True)
