import json
import os
import re
import sqlite3
import threading
from queue import Empty, Queue
from threading import Event, Thread

POOL = os.getenv("LLMS_POOL", "0") == "1"


def create_reader_connection(db_path):
    # isolation_level=None leaves the connection in autocommit mode
    conn = sqlite3.connect(
        db_path, timeout=1.0, check_same_thread=False, isolation_level=None
    )  # Lower - reads should be fast
    conn.execute("PRAGMA query_only=1")  # Read-only optimization
    return conn


def create_writer_connection(db_path):
    conn = sqlite3.connect(db_path)
    conn.execute("PRAGMA busy_timeout=5000")  # Reasonable timeout for busy connections
    conn.execute("PRAGMA journal_mode=WAL")  # Enable WAL mode for better concurrency
    conn.execute("PRAGMA cache_size=-128000")  # Increase cache size for better performance
    conn.execute("PRAGMA synchronous=NORMAL")  # Reasonable durability/performance balance
    return conn


def writer_thread(ctx, db_path, task_queue, stop_event):
    conn = create_writer_connection(db_path)
    try:
        while not stop_event.is_set():
            try:
                # Use timeout to check stop_event periodically
                task = task_queue.get(timeout=0.1)

                if task is None:  # Poison pill for clean shutdown
                    break

                sql, args, callback = task  # Optional callback for results

                try:
                    ctx.dbg("SQL>" + ("\n" if "\n" in sql else " ") + sql + ("\n" if args else " ") + str(args))
                    cursor = conn.execute(sql, args)
                    conn.commit()
                    ctx.dbg(f"lastrowid {cursor.lastrowid}, rowcount {cursor.rowcount}")
                    if callback:
                        callback(cursor.lastrowid, cursor.rowcount)
                except sqlite3.Error as e:
                    ctx.err("writer_thread", e)
                    if callback:
                        callback(None, None, error=e)
                finally:
                    task_queue.task_done()

            except Empty:
                continue

    finally:
        conn.close()


def to_dto(ctx, row, json_columns):
    # as=column -> [0,1,2]
    if not isinstance(row, dict):
        return row

    to = {}
    for k, v in row.items():
        if k in json_columns and v is not None and isinstance(v, str):
            try:
                to[k] = json.loads(v)
            except Exception as e:
                print(f"Failed to parse JSON for {k}: {v} ({type(v)})", e)
                to[k] = v
        else:
            to[k] = v
    return to


def valid_columns(all_columns, fields):
    if fields:
        if not isinstance(fields, list):
            fields = fields.split(",")
        cols = []
        for k in fields:
            k = k.strip()
            if k in all_columns:
                cols.append(k)
        return cols
    return []


def table_columns(all_columns, fields):
    cols = valid_columns(all_columns, fields)
    return ", ".join(cols) if len(cols) > 0 else ", ".join(all_columns)


def select_columns(all_columns, fields, select=None):
    columns = table_columns(all_columns, fields)
    if select == "distinct":
        return f"SELECT DISTINCT {columns}"
    return f"SELECT {columns}"


def order_by(all_columns, sort):
    cols = []
    for k in sort.split(","):
        k = k.strip()
        by = ""
        if k[0] == "-":
            by = " DESC"
            k = k[1:]
        if k in all_columns:
            cols.append(f"{k}{by}")
    return f"ORDER BY {', '.join(cols)} " if len(cols) > 0 else ""


def count_tokens_approx(messages: list[dict]) -> int:
    """
    Approximate token count for chat completion messages without external libraries.

    Handles various message formats:
    - Simple string content: {"role": "user", "content": "hello"}
    - Content arrays: {"role": "user", "content": [{"type": "text", "text": "hello"}]}
    - Tool calls, images, etc.
    """

    def count_text_tokens(text: str) -> int:
        if not text:
            return 0

        tokens = 0
        chunks = re.findall(r"\d+|[a-zA-Z]+|[^\s\w]|\s+", text)
        for chunk in chunks:
            if not chunk.strip():
                tokens += len(chunk) // 4
            elif chunk.isdigit():
                tokens += (len(chunk) + 2) // 3
            elif chunk.isalpha():
                tokens += 1 if len(chunk) <= 4 else (len(chunk) + 3) // 4
            else:
                tokens += len(chunk)
        return tokens

    def extract_text_content(value) -> str:
        """Recursively extract text from various content structures."""
        if value is None:
            return ""
        if isinstance(value, str):
            return value
        if isinstance(value, list):
            texts = []
            for item in value:
                if isinstance(item, dict):
                    # Handle content blocks: {"type": "text", "text": "..."}
                    if item.get("type") == "text" and "text" in item:
                        texts.append(item["text"])
                    # Handle tool use, tool results, etc.
                    elif "content" in item:
                        texts.append(extract_text_content(item["content"]))
                    elif "text" in item:
                        texts.append(item["text"])
                elif isinstance(item, str):
                    texts.append(item)
            return " ".join(texts)
        if isinstance(value, dict):
            if "text" in value:
                return value["text"]
            if "content" in value:
                return extract_text_content(value["content"])
        return ""

    total = 0
    for message in messages:
        # Message overhead
        total += 4

        role = message.get("role", "")
        total += count_text_tokens(role)

        content = message.get("content")
        text = extract_text_content(content)
        total += count_text_tokens(text)

        # Handle thinking/reasoning content
        for key in ["thinking", "reasoning", "reasoning_content"]:
            if key in message:
                text = extract_text_content(message[key])
                total += count_text_tokens(text)

        # Handle tool calls if present
        if "tool_calls" in message:
            for tool_call in message.get("tool_calls") or []:
                if isinstance(tool_call, dict):
                    fn = tool_call.get("function", {})
                    total += count_text_tokens(fn.get("name", ""))
                    total += count_text_tokens(fn.get("arguments", ""))
    # Reply priming
    total += 3

    return total


class DbManager:
    def __init__(self, ctx, db_path, clone=None):
        if db_path is None:
            raise ValueError("db_path is required")
        self.ctx = ctx
        self.db_path = db_path
        self.read_only_pool = Queue()
        if not clone:
            self.task_queue = Queue()
            self.stop_event = Event()
            self.writer_thread = Thread(target=writer_thread, args=(ctx, db_path, self.task_queue, self.stop_event))
            self.writer_thread.start()
        else:
            # share singleton writer thread in clones
            self.task_queue = clone.task_queue
            self.stop_event = clone.stop_event
            self.writer_thread = clone.writer_thread

    def create_reader_connection(self):
        return create_reader_connection(self.db_path)

    def create_writer_connection(self):
        return create_writer_connection(self.db_path)

    def resolve_connection(self):
        if POOL:
            try:
                return self.read_only_pool.get_nowait()
            except Empty:
                return self.create_reader_connection()
        else:
            return self.create_reader_connection()

    def release_connection(self, conn):
        if POOL:
            conn.rollback()
            self.read_only_pool.put(conn)
        else:
            conn.close()

    def write(self, query, args=None, callback=None):
        """
        Execute a write operation asynchronously.

        Args:
            query (str): The SQL query to execute.
            args (tuple, optional): Arguments for the query.
            callback (callable, optional): A function called after execution with signature:
                callback(lastrowid, rowcount, error=None)
                - lastrowid (int): output of cursor.lastrowid
                - rowcount (int): output of cursor.rowcount
                - error (Exception): exception if operation failed, else None
        """
        self.task_queue.put((query, args, callback))

    def log_sql(self, sql, parameters=None):
        if self.ctx.debug:
            self.ctx.dbg(
                "SQL>" + ("\n" if "\n" in sql else " ") + sql + ("\n" if parameters else " ") + str(parameters)
            )

    def exec(self, connection, sql, parameters=None):
        self.log_sql(sql, parameters)
        return connection.execute(sql, parameters or ())

    def all(self, sql, parameters=None, connection=None):
        """
        Execute a query and return all rows as a list of dictionaries.
        """
        conn = self.resolve_connection() if connection is None else connection

        try:
            self.log_sql(sql, parameters)
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(sql, parameters or ())
            rows = [dict(row) for row in cursor.fetchall()]
            return rows
        finally:
            if connection is None:
                conn.row_factory = None
                self.release_connection(conn)

    def one(self, sql, parameters=None, connection=None):
        """
        Execute a query and return the first row as a dictionary.
        """
        conn = self.resolve_connection() if connection is None else connection

        try:
            self.log_sql(sql, parameters)
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(sql, parameters or ())
            row = cursor.fetchone()
            return dict(row) if row else None
        finally:
            if connection is None:
                conn.row_factory = None
                self.release_connection(conn)

    def scalar(self, sql, parameters=None, connection=None):
        """
        Execute a scalar query and return the first column of the first row.
        """
        conn = self.resolve_connection() if connection is None else connection

        try:
            self.log_sql(sql, parameters)
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(sql, parameters or ())
            row = cursor.fetchone()
            return row[0] if row else None
        finally:
            if connection is None:
                conn.row_factory = None
                self.release_connection(conn)

    def column(self, sql, parameters=None, connection=None):
        """
        Execute a 1 column query and return the values as a list.
        """
        conn = self.resolve_connection() if connection is None else connection

        try:
            self.log_sql(sql, parameters)
            cursor = conn.execute(sql, parameters or ())
            return [row[0] for row in cursor.fetchall()]
        finally:
            if connection is None:
                self.release_connection(conn)

    def dict(self, sql, parameters=None, connection=None):
        """
        Execute a 2 column query and return the keys as the first column and the values as the second column.
        """
        conn = self.resolve_connection() if connection is None else connection

        try:
            self.log_sql(sql, parameters)
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(sql, parameters or ())
            rows = cursor.fetchall()
            return {row[0]: row[1] for row in rows}
        finally:
            if connection is None:
                conn.row_factory = None
                self.release_connection(conn)

    # Helper to safely dump JSON if value exists
    def value(self, val):
        if val is None or val == "":
            return None
        if isinstance(val, (dict, list)):
            return json.dumps(val)
        return val

    def insert(self, table, columns, info, callback=None):
        if not info:
            raise Exception("info is required")

        args = {}
        known_columns = columns.keys()
        for k, val in info.items():
            if k in known_columns and k != "id":
                args[k] = self.value(val)

        insert_keys = list(args.keys())
        insert_body = ", ".join(insert_keys)
        insert_values = ", ".join(["?" for _ in insert_keys])

        sql = f"INSERT INTO {table} ({insert_body}) VALUES ({insert_values})"

        self.write(sql, tuple(args[k] for k in insert_keys), callback)

    async def insert_async(self, table, columns, info):
        event = threading.Event()

        ret = [None, None]

        def cb(lastrowid, rowcount, error=None):
            nonlocal ret
            if error:
                ret[1] = error
            else:
                ret[0] = lastrowid
            event.set()

        self.insert(table, columns, info, cb)
        event.wait()
        if ret[1]:
            raise ret[1]
        return ret[0]

    def update(self, table, columns, info, callback=None):
        if not info:
            raise Exception("info is required")

        args = {}
        known_columns = columns.keys()
        for k, val in info.items():
            if k in known_columns and k != "id":
                args[k] = self.value(val)

        update_keys = list(args.keys())
        update_body = ", ".join([f"{k} = :{k}" for k in update_keys])

        args["id"] = info["id"]
        sql = f"UPDATE {table} SET {update_body} WHERE id = :id"

        self.write(sql, args, callback)

    async def update_async(self, table, columns, info):
        event = threading.Event()

        ret = [None, None]

        def cb(lastrowid, rowcount, error=None):
            nonlocal ret
            if error:
                ret[1] = error
            else:
                ret[0] = rowcount
            event.set()

        self.update(table, columns, info, cb)
        event.wait()
        if ret[1]:
            raise ret[1]
        return ret[0]

    def close(self):
        self.ctx.dbg("Closing database")
        self.stop_event.set()
        self.task_queue.put(None)  # Poison pill to signal shutdown
        self.writer_thread.join()

        while not self.read_only_pool.empty():
            try:
                conn = self.read_only_pool.get_nowait()
                conn.close()
            except Empty:
                break
