src/db.py

from __future__ import annotations

import sqlite3
from pathlib import Path
from typing import Any

from .models import EvaluationResult, ItineraryPlan, UserRequest


BASE_DIR = Path(__file__).parent.parent
DB_PATH = BASE_DIR / "qiddiya.db"


def get_connection() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_db() -> None:
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        """
        CREATE TABLE IF NOT EXISTS plans (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            visit_date TEXT NOT NULL,
            request_json TEXT NOT NULL,
            plan_json TEXT NOT NULL
        );
        """
    )
    cur.execute(
        """
        CREATE TABLE IF NOT EXISTS evaluations (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            visit_date TEXT NOT NULL,
            request_json TEXT NOT NULL,
            plan_json TEXT NOT NULL,
            evaluation_json TEXT NOT NULL
        );
        """
    )
    conn.commit()
    conn.close()


def _to_json(obj: Any) -> str:
    import json

    if hasattr(obj, "model_dump"):
        return json.dumps(obj.model_dump())
    if isinstance(obj, dict):
        return json.dumps(obj)
    raise TypeError("Object is not serializable")


def save_plan_result(request: UserRequest, plan: ItineraryPlan) -> None:
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        """
        INSERT INTO plans (visit_date, request_json, plan_json)
        VALUES (?, ?, ?);
        """,
        (
            plan.visit_date,
            _to_json(request),
            _to_json(plan),
        ),
    )
    conn.commit()
    conn.close()


def save_evaluation_result(
    request: UserRequest, plan: ItineraryPlan, evaluation: EvaluationResult
) -> None:
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        """
        INSERT INTO evaluations (visit_date, request_json, plan_json, evaluation_json)
        VALUES (?, ?, ?, ?);
        """,
        (
            plan.visit_date,
            _to_json(request),
            _to_json(plan),
            _to_json(evaluation),
        ),
    )
    conn.commit()
    conn.close()