import os
import time
import re
import traceback
import threading
import logging
import subprocess
import argparse
import ipaddress
import select

from flask import Flask, request, jsonify
from handlers import KasmLogHandler
from log_entry_pb2 import LogEntry
from datetime import datetime
from google.protobuf.message import DecodeError
from logging.handlers import RotatingFileHandler


log_path = "/var/log/kasm-logger"
os.makedirs(log_path, exist_ok=True)
os.makedirs("/run/docker/plugins", exist_ok=True)


app = Flask(__name__)
app_logger = logging.getLogger('kasmlogger')

sessions = {}
sessions_lock = threading.Lock()


def parse_container_data(data):
    container_info = data.get("Info", {})

    file = data.get("File")
    container_id = container_info.get("ContainerID")
    container_image = container_info.get("ContainerImageName")

    container_env = {k: v for env in container_info.get("ContainerEnv", []) if "=" in env for k, v in [env.split("=", 1)]}

    return container_id, container_image, container_env, file


def resolve_container_hostname(container_id, hostname):
    app_logger.info(f"resolve_container_hostname: resolving hostname({hostname}) for container({container_id})")

    # check if the hostname is already an IP address
    try:
        api_address = str(ipaddress.ip_address(hostname))
        app_logger.info(f"using IP address({api_address}) for container({container_id})")
        return api_address
    except ValueError:
        pass

    # resolve the hostname using the docker exec
    try:
        api_address = None
        cmd = f"docker exec {container_id} getent hosts {hostname} | awk '{{print $1}}'"
        api_address = subprocess.check_output(cmd, shell=True, stderr=subprocess.STDOUT, timeout=1).decode("utf-8").strip()
        app_logger.info(f"resolved {hostname} to {api_address} for container({container_id})")
    except subprocess.TimeoutExpired:
        app_logger.error(f"failed to resolve hostname({hostname}) for container({container_id}) due to timeout, the container may not be running")
    except subprocess.CalledProcessError as exception:
        app_logger.error(f"failed to resolve hostname({hostname}) for container({container_id}): {exception}")

    return api_address

def consume_container_logs(container_id, container_image, container_env, fd, stop_event):
    try:
        container_logger = logging.getLogger(f"container-{container_id}")
        container_logger.setLevel(logging.INFO)

        file_handler = RotatingFileHandler(
            filename=f"{log_path}/container-{container_id}.log",
            maxBytes=10 * 1024 * 1024,
            backupCount=5
        )
        file_handler.setFormatter(logging.Formatter('%(asctime)s - %(message)s'))
        container_logger.addHandler(file_handler)

        api_hostname = container_env.get("KASM_API_HOST")
        if not api_hostname:
            raise ValueError("Missing KASM_API_HOST in container environment")

        api_address = resolve_container_hostname(container_id, api_hostname)

        if api_address:
            config = {
                'log_path': "/api/kasm_session_log",
                'manager': {
                    'hostnames': [api_address],
                    'public_port': container_env["KASM_API_PORT"]
                }
            }

            kasm_handler = KasmLogHandler.create_cached_kasmloghandler(
                application="Session",
                kasm_api_jwt=container_env["KASM_API_JWT"],
                config=config,
                hostname=container_env["KASM_ID"]
            )
            container_logger.addHandler(kasm_handler)
        else:
            app_logger.warning(f"KASM logging disabled for container({container_id}) - could not resolve API hostname, logging to file only")

        app_logger.info(f"starting logging for container({container_id})")

        buffer = b""
        stop_requested = False

        while not stop_requested:
            if stop_event.is_set():
                stop_requested = True

            r = select.select([fd], [], [fd], 0.1)
            if not r[0]:
                continue # no data available

            chunk = os.read(fd, 1024)
            if chunk is None or len(chunk) == 0:
                continue # not enough data to read

            buffer += chunk

            while len(buffer) >= 4:
                message_size = int.from_bytes(buffer[:4], byteorder="big")
                if len(buffer) < message_size + 4:
                    break # not enough data to read

                message_bytes = buffer[4:message_size + 4]
                buffer = buffer[message_size + 4:]

                entry = LogEntry()
                try:
                    entry.ParseFromString(message_bytes)
                except Exception:
                    container_logger.exception(f"unexpected error while parsing log entry for container({container_id})")
                    continue

                channel = None
                line = entry.line.decode("utf-8", errors="replace")
                channel_match = re.match(r'^<KASM\|(\d+)>\s*(.*)', line)
                if channel_match:
                    channel = channel_match.group(1)
                    line = channel_match.group(2)

                container_logger.info(line, extra={
                    'channel': channel,
                    'timestamp': entry.time_nano,
                    'container_id': container_id,
                    'container_image': container_image
                })

        app_logger.info(f"stopping logging for container({container_id})")
    except Exception:
        app_logger.exception(f"error processing logs for container({container_id})")
    finally:
        os.close(fd)

@app.route("/Plugin.Activate", methods=["POST"])
def plugin_activate():
    app_logger.info("Plugin.Activate endpoint called")
    return jsonify({"Implements": ["LogDriver"]})


@app.route("/LogDriver.Capabilities", methods=["POST"])
def capabilities():
    return jsonify({"Cap": {"ReadLogs": True}})


@app.route("/LogDriver.ReadLogs", methods=["POST"])
def read_logs():
    data = request.get_json(force=True)
    follow = data.get("Config", {}).get("Follow", False)

    container_id, container_image, container_env, file = parse_container_data(data)

    log_file = f"{log_path}/container-{container_id}.log"
    if not os.path.exists(log_file):
        return jsonify({"Error": f"No logs found for container {container_id}"}), 404

    try:
        def generate_logs():
            with open(log_file, "r") as f:
                while True:
                    line = f.readline()
                    if line:
                        entry = LogEntry()
                        entry.time_nano = int(datetime.fromisoformat(line.split(" ")[0]).timestamp() * 1e9)
                        entry.line = line.encode("utf-8")
                        entry.source = "stdout"
                        entry.partial = False
                        payload = entry.SerializeToString()
                        payload_size = len(payload).to_bytes(4, byteorder="big")
                        yield payload_size + payload
                    elif follow:
                        position = f.tell()
                        f.seek(position)
                        time.sleep(0.1)
                        yield b'\x00\x00\x00\x00'
                    else:
                        break

        return generate_logs(), 200, {"Content-Type": "application/x-protobuf"}

    except Exception as exception:
        app_logger.exception("error reading logs.")
        return jsonify({"Error": f"Error reading logs: {str(exception)}"}), 500


@app.route("/LogDriver.StartLogging", methods=["POST"])
def start_logging():
    data = request.get_json(force=True)
    app_logger.info(f"container data:\n{data}")

    container_id, container_image, container_env, file_path = parse_container_data(data)

    kasm_id = container_env.get("KASM_ID", None)
    kasm_api_host = container_env.get("KASM_API_HOST", None)
    kasm_api_jwt = container_env.get("KASM_API_JWT", None)
    kasm_api_port = container_env.get("KASM_API_PORT", None)

    if not kasm_id or not kasm_api_host or not kasm_api_jwt or not kasm_api_port:
        app_logger.debug(f"container({container_id}) is not a KASM session container, skipping")
        return jsonify({}), 200

    try:
        fd = os.open(file_path, os.O_RDONLY | os.O_NONBLOCK)
    except Exception:
        return jsonify({"Error": f"Failed to open container log file: {file_path}"}), 500

    stop_event = threading.Event()
    thread = threading.Thread(target=consume_container_logs, args=(container_id, container_image, container_env, fd, stop_event), daemon=True)

    with sessions_lock:
        sessions[file_path] = (thread, stop_event)

    thread.start()

    return jsonify({}), 200


@app.route("/LogDriver.StopLogging", methods=["POST"])
def stop_logging():
    data = request.get_json(force=True)
    _, _, _, file_path = parse_container_data(data)

    with sessions_lock:
        session = sessions.pop(file_path, None)

    if session:
        thread, stop_event = session
        stop_event.set()
        thread.join()

    return jsonify({}), 200


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--debug", type=int, default=0)
    args = parser.parse_args()

    driver_log_file = os.path.join(log_path, "logging-driver.log")
    rotating_handler = RotatingFileHandler(
        filename=driver_log_file,
        maxBytes=10 * 1024 * 1024,
        backupCount=5
    )
    rotating_handler.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    rotating_handler.setFormatter(formatter)

    root_logger = logging.getLogger()
    root_logger.setLevel(logging.DEBUG)
    root_logger.addHandler(rotating_handler)

    app.run(host="unix:///run/docker/plugins/kasmlogger.sock", threaded=True, debug=bool(args.debug))