# websocket_server/task_api.py

from flask import request, jsonify
from datetime import datetime
import json
import logging

from websocket_server.models import Task
from websocket_server.config import get_db_session, get_redis_client

logger = logging.getLogger("websocket_server")
redis_client = get_redis_client()

# Import dispatch logic
from websocket_server.task_assigner import assign_task_to_worker
from websocket_server.shared_state import connected_workers, worker_lock


def register_routes(app):
    """
    Register all Flask routes to the given app.
    """

    @app.route("/api/create_task", methods=["POST"])
    def create_task():
        """
        Create a new task and insert it into the database.
        """
        data = request.get_json()
        logger.debug(f"Received create_task request: {data}")
        worker_id = data.get("worker_id")
        job_details_dict = data.get("job_details")
        job_details = json.dumps(job_details_dict) if job_details_dict else None
        task_type = data.get("task_type", "default")
        depends_on = data.get("depends_on")
        not_before = data.get("not_before")

        logger.info(f"Creating new task for worker {worker_id}, type {task_type}")
        if job_details_dict and "containers" in job_details_dict:
            for container in job_details_dict["containers"]:
                if "injected_files" in container:
                    logger.info(f"Task contains injected_files for container {container.get('container_name', 'unknown')}")
                else:
                    logger.debug(f"No injected_files for container {container.get('container_name', 'unknown')}")

        if not worker_id or not job_details:
            return jsonify({"error": "worker_id and job_details are required"}), 400

        try:
            not_before_dt = None
            if not_before:
                not_before_dt = datetime.fromisoformat(not_before)

            with get_db_session() as session:
                task = Task(
                    worker_id=worker_id,
                    job_details=job_details,
                    task_type=task_type,
                    status="pending",
                    success=None,
                    creation_time=datetime.utcnow(),
                    depends_on=depends_on,
                    not_before=not_before_dt
                )
                session.add(task)
                session.commit()

                logger.info(f"Created task {task.id} for worker {worker_id}")
                redis_client.set(f"worker_queue_{worker_id}", "True")

                return jsonify({"message": "Task created", "task_id": task.id}), 201

        except Exception as e:
            logger.error(f"Error creating task: {e}")
            return jsonify({"error": "Internal error creating task"}), 500

    @app.route("/api/assign_task", methods=["POST"])
    def manually_assign_task():
        """
        Force task assignment for a specific worker via HTTP trigger.
        """
        data = request.get_json()
        worker_id = data.get("worker_id")

        if not worker_id:
            return jsonify({"error": "Missing worker_id"}), 400

        logger.info(f"Manual trigger for assign_task to worker {worker_id}")
        assign_task_to_worker(worker_id)
        return jsonify({"message": f"Dispatched task assignment to worker {worker_id}."})

    @app.route("/api/connected_clients", methods=["GET"])
    def get_connected_clients():
        """
        Returns currently connected worker IDs and socket session IDs.
        """
        logger.debug("Listing connected clients")
        with worker_lock:
            result = [
                {"worker_id": wid, "socket_id": sid, "ip": "unknown"}
                for wid, sid in connected_workers.items()
            ]
        return jsonify({"connected_clients": result})