import json
import time
import threading
import sqlite3
from flask import Flask, jsonify, request
import paho.mqtt.client as mqtt

# ========= CONFIGURATION =========
MQTT_BROKER = 'localhost'
MQTT_PORT = 1883
MQTT_TOPICS = ['lte_message', 'nr_message', 'wifi_message']
MQTT_CONTROL_TOPIC = 'test_control_command'

DB_FILE = 'roadtest.db'
WEB_PORT = 5000

status_context = {
    'testing': False,
    'path_id': None,
    'start_time': None
}

# ========= DATABASE SETUP =========
def init_db():
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute('''CREATE TABLE IF NOT EXISTS test_paths (
        path_id TEXT PRIMARY KEY,
        start_time TEXT,
        end_time TEXT,
        remark TEXT
    )''')
    cursor.execute('''CREATE TABLE IF NOT EXISTS network_measurements (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        path_id TEXT,
        message_type TEXT,
        data_json TEXT,
        ts TEXT
    )''')
    conn.commit()
    conn.close()

# ========= MQTT CALLBACKS =========
def on_connect(client, userdata, flags, rc):
    print(f"[MQTT] Connected with result code {rc}")
    client.subscribe(MQTT_CONTROL_TOPIC)
    for topic in MQTT_TOPICS:
        client.subscribe(topic)


def on_message(client, userdata, msg):
    topic = msg.topic
    payload = msg.payload.decode('utf-8')

    if topic == MQTT_CONTROL_TOPIC:
        handle_control_message(payload)
        return

    if not status_context['testing']:
        return

    # Save measurement
    try:
        data = json.loads(payload)
        insert_measurement(
            status_context['path_id'],
            topic,
            payload,
            time.strftime("%Y-%m-%d %H:%M:%S")
        )
        print(f"[DATA] Saved data on {topic} for path {status_context['path_id']}")
    except Exception as e:
        print(f"[ERROR] Failed to parse or save message: {e}")


def handle_control_message(payload):
    try:
        cmd = json.loads(payload)
        action = cmd.get('action')
        if action == 'start':
            status_context['testing'] = True
            status_context['path_id'] = cmd['path_id']
            status_context['start_time'] = time.strftime("%Y-%m-%d %H:%M:%S")
            insert_path(status_context['path_id'], status_context['start_time'])
            print(f"[CONTROL] Test started with path_id={status_context['path_id']}")
        elif action == 'stop':
            status_context['testing'] = False
            update_path_endtime(status_context['path_id'], time.strftime("%Y-%m-%d %H:%M:%S"))
            print(f"[CONTROL] Test ended for path_id={status_context['path_id']}")
            status_context['path_id'] = None
    except Exception as e:
        print(f"[ERROR] Control message error: {e}")

# ========= DATABASE FUNCTIONS =========
def insert_path(path_id, start_time):
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute('INSERT OR REPLACE INTO test_paths (path_id, start_time) VALUES (?, ?)',
                   (path_id, start_time))
    conn.commit()
    conn.close()

def update_path_endtime(path_id, end_time):
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute('UPDATE test_paths SET end_time = ? WHERE path_id = ?', (end_time, path_id))
    conn.commit()
    conn.close()

def insert_measurement(path_id, message_type, data_json, ts):
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute('''INSERT INTO network_measurements (path_id, message_type, data_json, ts)
                      VALUES (?, ?, ?, ?)''', (path_id, message_type, data_json, ts))
    conn.commit()
    conn.close()

# ========= WEB SERVER =========
app = Flask(__name__)

@app.route('/status', methods=['GET'])
def get_status():
    return jsonify({
        'testing': status_context['testing'],
        'path_id': status_context['path_id'],
        'start_time': status_context['start_time']
    })

@app.route('/paths', methods=['GET'])
def list_paths():
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute('SELECT * FROM test_paths ORDER BY start_time DESC')
    rows = cursor.fetchall()
    conn.close()
    return jsonify(rows)

# ========= MAIN =========
def main():
    init_db()
    mqtt_client = mqtt.Client()
    mqtt_client.on_connect = on_connect
    mqtt_client.on_message = on_message
    mqtt_client.connect(MQTT_BROKER, MQTT_PORT, 60)

    threading.Thread(target=app.run, kwargs={'port': WEB_PORT, 'debug': False, 'use_reloader': False}, daemon=True).start()
    print("[SYSTEM] Web UI running on http://localhost:5000")
    mqtt_client.loop_forever()

if __name__ == '__main__':
    main()