from flask import Flask, request, jsonify
from tinydb import TinyDB, Query, where
from tinydb.operations import set as ops_set
import json

from flask_cors import CORS
import os
import time
import hmac
import hashlib
import base64
import re
from typing import Dict, Optional, Tuple

app = Flask(__name__)

UI_ORIGIN = os.getenv("UI_ORIGIN", "https://db.northdanubesoft.eu")
CORS(app, resources={r"/*": {"origins": [UI_ORIGIN]}})

# --- Auth & DB config ---------------------------------------------------------
SECRET_KEY = os.getenv("SECRET_KEY", "dev-secret-change-me")
TOKEN_TTL = int(os.getenv("TOKEN_TTL", "3600"))  # seconds
# Default to a writable local ./data dir in dev; override to /data in Docker via env
DB_DIR = os.path.abspath(os.getenv("DB_DIR", os.path.join(os.path.dirname(__file__), "..", "data")))

# Cache TinyDB instances per application token file
_DB_CACHE: Dict[str, TinyDB] = {}

# --- Helpers -----------------------------------------------------------------

def json_body(required: bool = True):
    data = request.get_json(silent=True)
    if data is None:
        raw = request.get_data(as_text=True)
        if raw:
            try:
                data = json.loads(raw)
            except Exception:
                if required:
                    return None, (jsonify({"error": "Expected JSON body"}), 400)
                else:
                    data = None
    elif isinstance(data, str):
        try:
            data = json.loads(data)
        except Exception:
            if required:
                return None, (jsonify({"error": "Expected JSON body"}), 400)
            else:
                data = None
    if required and data is None:
        return None, (jsonify({"error": "Expected JSON body"}), 400)
    return data, None

# --- Token helpers ------------------------------------------------------------

def _b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")

def _b64url_json(obj: dict) -> str:
    return _b64url(json.dumps(obj, separators=(",", ":")).encode("utf-8"))

def sign_token(app_id: str, ttl: int = TOKEN_TTL) -> str:
    now = int(time.time())
    payload = {"app_id": app_id, "iat": now, "exp": now + ttl}
    header = {"alg": "HS256", "typ": "JWT"}
    h = _b64url_json(header)
    p = _b64url_json(payload)
    to_sign = f"{h}.{p}".encode("utf-8")
    sig = hmac.new(SECRET_KEY.encode("utf-8"), to_sign, hashlib.sha256).digest()
    return f"{h}.{p}.{_b64url(sig)}"

def verify_token(token: str) -> Tuple[bool, Optional[dict], Optional[str]]:
    try:
        parts = token.split(".")
        if len(parts) != 3:
            return False, None, "Malformed token"
        h_b64, p_b64, s_b64 = parts
        to_sign = f"{h_b64}.{p_b64}".encode("utf-8")
        sig = base64.urlsafe_b64decode(s_b64 + "==")
        expected = hmac.new(SECRET_KEY.encode("utf-8"), to_sign, hashlib.sha256).digest()
        if not hmac.compare_digest(sig, expected):
            return False, None, "Bad signature"
        payload = json.loads(base64.urlsafe_b64decode(p_b64 + "==").decode("utf-8"))
        now = int(time.time())
        if payload.get("exp", 0) < now:
            return False, None, "Expired"
        return True, payload, None
    except Exception as e:
        return False, None, str(e)

# --- Access token / DB helpers -----------------------------------------------

def _sanitize_filename(name: str) -> str:
    # allow only safe chars; replace others with '_'
    return re.sub(r"[^A-Za-z0-9._-]", "_", name)[:128]

def get_application_token_from(body: Optional[dict]) -> Optional[str]:
    # Prefer header, then JSON body
    hdr = request.headers.get("X-Application-Token")
    if hdr:
        return hdr.strip()
    if isinstance(body, dict):
        at = body.get("application_token")
        if isinstance(at, str) and at.strip():
            return at.strip()
    return None

def get_db_for_application_token(application_token: str) -> TinyDB:
    safe = _sanitize_filename(application_token)
    path = os.path.join(DB_DIR, f"{safe}.json")
    # return cached TinyDB if exists
    db = _DB_CACHE.get(path)
    if db is None:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        db = TinyDB(path)
        _DB_CACHE[path] = db
    return db

def authenticate_and_select_db(body: Optional[dict]) -> Tuple[Optional[TinyDB], Optional[tuple]]:
    # Verify Bearer token
    authz = request.headers.get("Authorization", "")
    if not authz.startswith("Bearer "):
        return None, (jsonify({"error": "Missing Authorization Bearer token"}), 401)
    ok, payload, err = verify_token(authz[len("Bearer "):].strip())
    if not ok:
        return None, (jsonify({"error": f"Unauthorized: {err}"}), 401)
    # Get application token to select DB
    application_token = get_application_token_from(body)
    if not application_token:
        return None, (jsonify({"error": "Missing application_token (header X-Application-Token or JSON body)"}), 400)
    db = get_db_for_application_token(application_token)
    return db, None

def build_query(field: str, op: str, value):
    """Translate a simple JSON filter into a TinyDB Query.
    Supported ops: ==, !=, >, >=, <, <=, in, contains
    """
    f = where(field)
    if op == "==":
        return f == value
    if op == "!=":
        return f != value
    if op == ">":
        return f > value
    if op == ">=":
        return f >= value
    if op == "<":
        return f < value
    if op == "<=":
        return f <= value
    if op == "in":
        # value should be a list
        if not isinstance(value, list):
            value = [value]
        return f.one_of(value)
    if op == "contains":
        # substring for strings; membership for lists
        return f.test(lambda v: (isinstance(v, str) and isinstance(value, str) and value in v)
                                or (isinstance(v, (list, tuple, set)) and value in v))
    raise ValueError(f"Unsupported op: {op}")

def parse_assignment(expr: str):
    """Parse a string expression like 'field = \"value\"' into (field, value)."""
    if '=' not in expr:
        raise ValueError("Expression must contain '='")
    field_part, value_part = expr.split('=', 1)
    field = field_part.strip()
    value_raw = value_part.strip()
    # Remove trailing comma if present
    if value_raw.endswith(','):
        value_raw = value_raw[:-1].rstrip()
    # Strip quotes if value starts and ends with same quote
    if (len(value_raw) >= 2) and ((value_raw[0] == value_raw[-1]) and value_raw[0] in ("'", '"')):
        value = value_raw[1:-1]
    else:
        # Try to parse JSON for numbers, booleans, null, objects, arrays, or quoted strings
        # Check if looks like JSON
        json_like_start = ('{', '[', '"', '-', 't', 'f', 'n') + tuple(str(i) for i in range(10))
        if value_raw and (value_raw[0] in json_like_start):
            try:
                value = json.loads(value_raw)
            except Exception:
                value = value_raw
        else:
            value = value_raw
    return field, value

@app.route("/connect", methods=["POST"]) 
def connect():
    body, err = json_body()
    if err:
        return err
    app_id = body.get("application_id") if isinstance(body, dict) else None
    application_token = body.get("application_token") if isinstance(body, dict) else None
    if not isinstance(app_id, str) or not app_id.strip():
        return jsonify({"error": "application_id required"}), 400
    if not isinstance(application_token, str) or not application_token.strip():
        return jsonify({"error": "application_token required"}), 400
    # Ensure DB file exists for this application token
    _ = get_db_for_application_token(application_token)
    tok = sign_token(app_id.strip())
    return jsonify({
        "token": tok,
        "expires_in": TOKEN_TTL,
        "token_type": "Bearer",
        "application_token": application_token
    }), 200

# --- Routes ------------------------------------------------------------------

@app.route("/healthz", methods=["GET"]) 
def healthz():
    return jsonify({"status": "ok"}), 200


@app.route("/insert", methods=["POST"]) 
def insert():
    body, err = json_body()
    if err:
        return err
    db, auth_err = authenticate_and_select_db(body)
    if auth_err:
        return auth_err
    # Accept either {"doc": {...}} or raw JSON object as the document
    if isinstance(body, dict) and "doc" in body:
        doc = body["doc"]
    else:
        doc = body
    if not isinstance(doc, dict):
        return jsonify({"error": "Body must be an object or {doc: {...}}"}), 400

    doc_id = db.insert(doc)
    return jsonify({"message": "inserted", "doc_id": doc_id}), 200


@app.route("/insert_many", methods=["POST"]) 
def insert_many():
    body, err = json_body()
    if err:
        return err
    db, auth_err = authenticate_and_select_db(body)
    if auth_err:
        return auth_err
    docs = body.get("docs") if isinstance(body, dict) else None
    if not isinstance(docs, list) or not all(isinstance(d, dict) for d in docs):
        return jsonify({"error": "Expected {docs: [ {...}, {...} ]}"}), 400
    ids = db.insert_multiple(docs)
    return jsonify({"message": "inserted", "doc_ids": ids}), 200


@app.route("/get_all", methods=["POST", "GET"]) 
def get_all():
    body, _ = json_body(required=False)
    db, auth_err = authenticate_and_select_db(body)
    if auth_err:
        return auth_err
    data = db.all()
    _items = []
    for item in data:
        d = dict(item)
        d["doc_id"] = item.doc_id
        _items.append(d)
    return jsonify(_items), 200


@app.route("/get", methods=["POST"]) 
def get_one():
    body, err = json_body()
    if err:
        return err
    db, auth_err = authenticate_and_select_db(body)
    if auth_err:
        return auth_err
    doc_id = body.get("doc_id") if isinstance(body, dict) else None
    if not isinstance(doc_id, int):
        return jsonify({"error": "Expected {doc_id: <int>}"}), 400
    doc = db.get(doc_id=doc_id)
    if doc is None:
        return jsonify({"error": "not found"}), 404
    # Include doc_id so clients can reference it later
    doc_with_id = dict(doc)
    doc_with_id["doc_id"] = doc_id
    return jsonify(doc_with_id), 200


@app.route("/search", methods=["POST"]) 
def search():
    body, err = json_body()
    print (body)
    if err:
        return err
    db, auth_err = authenticate_and_select_db(body)
    if auth_err:
        return auth_err
    if not isinstance(body, dict):
        return jsonify({"error": "Expected JSON object"}), 400

    # Accept either a single filter or a list of filters (ANDed)
    filters = body.get("where")
    if isinstance(filters, dict):
        filters = [filters]
    if not isinstance(filters, list) or not filters:
        return jsonify({"error": "Expected {where: {field, op, value}} or a list of them"}), 400

    try:
        q = None
        for f in filters:
            field = f.get("field")
            op = f.get("op")
            value = f.get("value")
            if not isinstance(field, str) or not isinstance(op, str):
                return jsonify({"error": "Each filter needs 'field' and 'op'"}), 400
            clause = build_query(field, op, value)
            q = clause if q is None else (q & clause)
    except ValueError as e:
        return jsonify({"error": str(e)}), 400
    if not callable(q):
        return jsonify({"error":"Invalid query built"}), 400
    results = db.search(q)
    _items = []
    for item in results:
        d = dict(item)
        d["doc_id"] = item.doc_id
        _items.append(d)
    return jsonify(_items), 200


@app.route("/update", methods=["POST"]) 
def update():
    body, err = json_body()
    if err:
        return err
    db, auth_err = authenticate_and_select_db(body)
    if auth_err:
        return auth_err
    if not isinstance(body, dict):
        return jsonify({"error": "Expected JSON object"}), 400
    print(body)
    # Option A: update by doc_id
    if "doc_id" in body and "fields" in body:
        doc_id = body.get("doc_id")
        fields = body.get("fields")
        if isinstance(fields, str):
            try:
                field, value = parse_assignment(fields)
            except Exception as e:
                return jsonify({"error": f"Failed to parse fields string: {str(e)}"}), 400
            updated = db.update(ops_set(field, value), doc_ids=[doc_id])
            return jsonify({"updated": len(updated)}), 200
        if not isinstance(doc_id, int) or not isinstance(fields, dict):
            return jsonify({"error": "Expected {doc_id: int, fields: {...}}"}), 400
        updated = db.update(fields, doc_ids=[doc_id])
        return jsonify({"updated": len(updated)}), 200

    # Option B: update by query
    if "where" in body and "fields" in body:
        filters = body.get("where")
        fields = body.get("fields")
        if isinstance(filters, dict):
            filters = [filters]
        if not isinstance(filters, list):
            return jsonify({"error": "Expected {where: [...], fields: {...}}"}), 400
        try:
            q = None
            for f in filters:
                clause = build_query(f.get("field"), f.get("op"), f.get("value"))
                q = clause if q is None else (q & clause)
        except ValueError as e:
            return jsonify({"error": str(e)}), 400
        if isinstance(fields, str):
            try:
                field, value = parse_assignment(fields)
            except Exception as e:
                return jsonify({"error": f"Failed to parse fields string: {str(e)}"}), 400
            updated = db.update(ops_set(field, value), q)
            return jsonify({"updated": len(updated)}), 200
        if not isinstance(fields, dict):
            return jsonify({"error": "Expected {where: [...], fields: {...}}"}), 400
        updated = db.update(fields, q)
        return jsonify({"updated": len(updated)}), 200

    return jsonify({"error": "Provide either {doc_id, fields} or {where, fields}"}), 400


@app.route("/remove", methods=["POST"]) 
def remove():
    body, err = json_body()
    if err:
        return err
    db, auth_err = authenticate_and_select_db(body)
    if auth_err:
        return auth_err
    if not isinstance(body, dict):
        return jsonify({"error": "Expected JSON object"}), 400

    # Option A: by doc_id
    if "doc_id" in body:
        doc_id = body.get("doc_id")
        if not isinstance(doc_id, int):
            return jsonify({"error": "Expected {doc_id: int}"}), 400
        removed = db.remove(doc_ids=[doc_id])
        return jsonify({"removed": len(removed)}), 200

    # Option B: by query
    if "where" in body:
        filters = body.get("where")
        if isinstance(filters, dict):
            filters = [filters]
        if not isinstance(filters, list):
            return jsonify({"error": "Expected {where: [...]}"}), 400
        try:
            q = None
            for f in filters:
                clause = build_query(f.get("field"), f.get("op"), f.get("value"))
                q = clause if q is None else (q & clause)
        except ValueError as e:
            return jsonify({"error": str(e)}), 400
        removed = db.remove(q)
        return jsonify({"removed": len(removed)}), 200

    return jsonify({"error": "Provide {doc_id} or {where}"}), 400


@app.route("/truncate", methods=["POST"]) 
def truncate():
    body, _ = json_body(required=False)
    db, auth_err = authenticate_and_select_db(body)
    if auth_err:
        return auth_err
    db.truncate()
    return jsonify({"message": "truncated"}), 200


if __name__ == "__main__":
    # Note: Flask dev server is single-threaded by default; good for local testing.
    # In Docker, gunicorn is used to run this app.
    app.run(debug=True, port=5001)