import socket, json, threading, os, time

--- CONFIGURATION & QUEUES ---

QI_FAST = []
QI_SLOW = []
QO = {}

New Mappings based on your ranges

INTERNAL_TASKS = ["1", "2", "3", "4"] # TM Pulls, Results, Status DEV_TASKS = ["101", "102"] # Admin/Dev FAST_TASKS = ["201"] # Data Fetching SLOW_TASKS = ["301"] # Email Generation

events = {} task_counter = 1000 lock = threading.Lock()

BASE_DIR = os.path.dirname(os.path.abspath(file)) STORAGE_DIR = os.path.join(BASE_DIR, "storage") if not os.path.exists(STORAGE_DIR): os.makedirs(STORAGE_DIR)

def handle_client(conn, addr): global task_counter try: data = conn.recv(65536).decode('utf-8') if not data: return raw_input = json.loads(data)

    t_type = str(raw_input.get("task_type", "0"))
    u_id = raw_input.get("user_id", "Unknown")
    prio = str(raw_input.get("priority", "0"))
    print(t_type)

    # --- 1. TRANSFORMATION TO NEW SCHEMA ---
    with lock:
        task_counter += 1
        tid = task_counter

    formatted_task = {
        "header": {
            "task_id": tid,
            "user_id": u_id,
            "task_type": t_type,
            "priority": prio,
            "timestamp": int(time.time())
        },
        "payload": {
            "ingredients": raw_input.get("ingredients", {}),
            "result": {"data": {}, "db_reference_id": None}
        },
        "meta": {
            "status": "Pending",
            "worker_id": None,
            "start_time": None,
            "end_time": None,
            "total_time": None,
            "progress": 0,
            "error_log": []
        }
    }

    # --- 2. ROUTING LOGIC ---

    # A. FAST QUEUE (Range 201-300)
    if 201 <= int(t_type) <= 300:
        event = threading.Event()
        events[tid] = event
        with lock:
            # Priority 1 = Premium
            if prio == "1":
                index = 0
                for t in QI_FAST:
                    if t["header"]["priority"] == "1": index += 1
                    else: break
                QI_FAST.insert(index, formatted_task)
            else:
                QI_FAST.append(formatted_task)

        conn.send(json.dumps({"task_id": tid, "status": "Queued"}).encode('utf-8'))

        # Wait for TM1 to finish
        if event.wait(timeout=15):
            result = QO.pop(tid, None)
            if result:
                conn.send(json.dumps(result).encode('utf-8'))
        else:
            conn.send(json.dumps({"task_id": tid, "status": "Processing_In_Background"}).encode('utf-8'))

    # B. SLOW/DEV QUEUE (Range 101-200 or 301-400)
    elif (101 <= int(t_type) <= 200) or (301 <= int(t_type) <= 400):
        with lock:
            # Dev (Priority 2) takes absolute top spot
            if prio == "2":
                QI_SLOW.insert(0, formatted_task)
            else:
                QI_SLOW.append(formatted_task)
        conn.send(json.dumps({"task_id": tid, "status": "Accepted_Async"}).encode('utf-8'))

    # --- 3. INTERNAL TM OPERATIONS (Range 1-100) ---

    elif t_type == "1": # TM1 PULL (Fast)
        with lock: task = QI_FAST.pop(0) if QI_FAST else None
        conn.send(json.dumps(task).encode('utf-8'))

    elif t_type == "2": # TM2 PULL (Slow)
        with lock: task = QI_SLOW.pop(0) if QI_SLOW else None
        conn.send(json.dumps(task).encode('utf-8'))

    elif t_type == "3": # TM POST RESULT
        # Expects the full formatted_task back
        res_tid = formatted_task["header"]["task_id"]
        formatted_task["meta"]["status"] = "Completed"

        if res_tid in events:
            with lock: QO[res_tid] = formatted_task
            events[res_tid].set()
            conn.send(b"ACK_PUSHED")
        else:
            # Save to disk if user disconnected
            filename = os.path.join(STORAGE_DIR, f"results_{u_id}.jsonl")
            with lock:
                with open(filename, "a") as f:
                    f.write(json.dumps(formatted_task) + "\n")
            conn.send(b"ACK_STORED")

    elif t_type == "4": # RS STATUS CHECK
        stats = {
            "qi_fast_count": len(QI_FAST), 
            "qi_slow_count": len(QI_SLOW), 
            "active_events": len(events),
            "qo_count": len(QO),
            "last_tid": task_counter
        }
        conn.send(json.dumps(stats).encode('utf-8'))

    elif t_type == "5": # NEW: FETCH USER RESULTS
        u_id = raw_input.get("user_id", "Unknown")
        filename = os.path.join(STORAGE_DIR, f"results_{u_id}.jsonl")

        results = []
        if os.path.exists(filename):
            with lock:
                with open(filename, "r") as f:
                    # Read last 20 lines to avoid memory lag
                    lines = f.readlines()
                    for line in lines[-20:]: 
                        results.append(json.loads(line))

        conn.send(json.dumps(results).encode('utf-8'))

except Exception as e:
    print(f"Error handling request: {e}")
finally:
    conn.close()

def start_rs(): server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) server.bind(('localhost', 7000)) server.listen(50) print("RS Layer Online [New Schema Active]") while True: conn, addr = server.accept() threading.Thread(target=handle_client, args=(conn, addr), daemon=True).start()

if name == "main": start_rs()