#!/usr/bin/env python3

import os
import time
import json
import threading
import http.server
from pathlib import Path
from urllib.parse import urlencode
from urllib.request import Request, urlopen
from urllib.error import URLError, HTTPError
from collections.abc import Mapping
from urllib.response import addinfourl
from typing import cast, Any
from datetime import datetime

class H(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == "/healthz":
            self.send_response(200)
            self.end_headers()
            self.wfile.write(b"ok")
        else:
            self.send_response(404)
            self.end_headers()
    def log_message(self, format, *args):
        pass

def int_env(name: str, default: int) -> int:
    try:
        return int(os.getenv(name, str(default)))
    except Exception:
        return default

TASKS_PATH = Path("/app/tasks")
API_URL = os.getenv("API_URL", "https://api.sodew.ai/api/tasks").rstrip("/")
WORKER_UUID = os.getenv("WORKER_UUID", "worker-uuid")
WORKER_NAME = os.getenv("WORKER_NAME", "worker-name")
WORKER_POOL = os.getenv("WORKER_POOL", "")
WORKER_VERSION = "6.3.445"
GETTING_PAUSE = int_env("GETTING_PAUSE", 30)
SENDING_PAUSE = int_env("SENDING_PAUSE", 15)

HTTP_TIMEOUT = 15
DEFAULT_HEADERS = {"Accept": "application/json", "Content-Type": "application/json", "User-Agent": "kube-worker/1.1"}

def start_health():
    t = threading.Thread(target=http.server.HTTPServer(('0.0.0.0', 8080), H).serve_forever, daemon=True)
    t.start()

def ensure_dir() -> None:
    TASKS_PATH.mkdir(parents=True, exist_ok=True)

def log_info(text: str) -> None:
    print(datetime.now().strftime("[%Y.%m.%d %H:%M:%S]") + f" {text}")

def format_error_body(body: Any) -> str:
    if body is None:
        return "<no body>"

    if isinstance(body, dict):
        msg = body.get("message")
        if isinstance(msg, str) and msg.strip():
            return msg

        try:
            return json.dumps(body, ensure_ascii=False, sort_keys=True)
        except Exception:
            return repr(body)

    if isinstance(body, list):
        try:
            return json.dumps(body, ensure_ascii=False, sort_keys=True)
        except Exception:
            return repr(body)

    try:
        return str(body)
    except Exception:
        return "<unprintable body>"

def task_read(path: Path) -> dict[str, str]:
    result: dict[str, str] = {}
    for line in path.read_text(encoding="utf-8", errors="ignore").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        result[k.strip()] = v.strip()
    return result

def task_save(path: Path, data: Mapping[str, object]) -> None:
    flat: dict[str, object] = dict(data)
    lines: list[str] = []
    for key, value in flat.items():
        if not isinstance(key, str):
            key = str(key)

        if value is None:
            value_str = ""
        elif isinstance(value, (dict, list)):
            try:
                value_str = json.dumps(value, ensure_ascii=False)
            except Exception:
                value_str = str(value)
        else:
            value_str = str(value)

        value_str = value_str.replace("\n", " ").replace("\r", " ")
        lines.append(f"{key}={value_str}")

    content = "\n".join(lines) + "\n"
    path.write_text(content, encoding="utf-8")

def http_request_json(
    method: str,
    url: str,
    *,
    params: Mapping[str, str] | None = None,
    json_body: Mapping[str, object] | None = None,
    headers: Mapping[str, str] | None = None,
    timeout: int = HTTP_TIMEOUT,
) -> tuple[int, object | None]:
    if params:
        qs = urlencode(params)
        url = f"{url}?{qs}"
    body_bytes = None
    req_headers = dict(DEFAULT_HEADERS)
    if headers:
        req_headers.update(headers)
    if json_body is not None:
        body_bytes = json.dumps(json_body).encode("utf-8")
    req = Request(url, data=body_bytes, headers=req_headers, method=method)
    try:
        with urlopen(req, timeout=timeout) as resp:
            resp_typed = cast(addinfourl, resp)
            code = resp_typed.getcode() or 0
            body = resp_typed.read()

    except HTTPError as e:
        try:
            err_bytes = e.read()
        except Exception:
            err_bytes = b""
        if not err_bytes:
            return e.code, None
        try:
            return e.code, json.loads(err_bytes.decode("utf-8", errors="replace"))
        except Exception:
            return e.code, err_bytes.decode("utf-8", errors="replace")
    except URLError as e:
        return 0, {"error": "network", "reason": str(e)}

    if not body:
        return code, None
    try:
        return code, json.loads(body.decode("utf-8", errors="replace"))
    except Exception:
        return code, None

def task_receive() -> dict[str, Any] | None:
    log_info(f"Receiving new tasks from API | Worker: {WORKER_NAME}")
    url = f"{API_URL}/receive"
    payload: dict[str, str] = {
        "worker_name": WORKER_NAME,
        "worker_uuid": WORKER_UUID,
        "worker_version": WORKER_VERSION
    }
    if WORKER_POOL and WORKER_POOL.strip():
        payload["worker_pool"] = WORKER_POOL.strip()

    code, body = http_request_json("POST", url, json_body=payload)

    if code == 204:
        log_info("Code 204 | No tasks from API")
        return None

    if code == 200:
        try:
            log_info("Code: 200 | Raw data: " + (json.dumps(body, ensure_ascii=False, sort_keys=True) if isinstance(body, (dict, list)) else repr(body)))
        except Exception:
            log_info("Code: 200 | Raw data: <unserializable>")

        if isinstance(body, dict):
            d = cast(dict[str, object], body)
            try:
                log_info("Task: received | " + json.dumps(d, ensure_ascii=False, sort_keys=True))
            except Exception:
                log_info("Task: received | <unserializable dict>")

            if "uuid" in d and "action" in d:
                return d

            log_info("Task: missing required fields 'uuid' or 'action'")
        else:
            log_info("Task: not recognized (body is not a dict)")
    else:
        message_error = format_error_body(body)
        log_info(f"Code: {code} | Error: {message_error}")

    return None

def main() -> None:
    start_health()
    ensure_dir()

    while True:
        task_files = list(TASKS_PATH.glob("*.task"))
        task_count = len(task_files)
        log_info(f"Task files found: {task_count}")

        if not task_files:
            task = task_receive()
            if task:
                uuid = str(task.get("uuid", "")).strip()
                action = str(task.get("action", "")).strip()
                task.pop("uuid", None)
                task["status"] = "waiting"
                if uuid and action:
                    log_info(f"Task from API: {uuid}")
                    path = TASKS_PATH / f"{uuid}.task"
                    if not path.exists():
                        task_save(path=path, data=task)
                else:
                    log_info(f"Task: {uuid} | Error: uuid or action is empty, skip")
            time.sleep(GETTING_PAUSE)
        else:
            for path in task_files:
                uuid = path.stem
                try:
                    data = task_read(path)
                    log_info(f"Task: {uuid} | Parse task success")
                except Exception:
                    log_info(f"Task: {uuid} | Task not recognized")
                    continue
                action = (data.get("action") or "unknown").strip().lower()
                if action == "pause":
                    continue
                status = (data.get("status") or "unknown").strip().lower()

                payload: dict[str, str] = dict(data)
                payload["worker_uuid"] = WORKER_UUID
                payload["status"] = status
                log_info(f"Task: {uuid} | Payload: {json.dumps(payload, ensure_ascii=False)}")
                url = f"{API_URL}/{uuid}/status"
                code, body = http_request_json("PATCH", url, json_body=payload)
                if not (200 <= code < 300):
                    message_error = format_error_body(body)
                    log_info(f"Task: {uuid} | Code: {code} | Error: {message_error}")
                    continue

                if status not in {"new", "waiting", "execution"}:
                    log_info(f"Task: {uuid} | Remove...")
                    try:
                        path.unlink(missing_ok=True)
                    except Exception:
                        pass

            time.sleep(SENDING_PAUSE)

if __name__ == "__main__":
    print(f"[ Worker Agent {WORKER_VERSION} | {datetime.now():%Y-%m-%d %H-%M-%S} ______________ ]")
    print(f"🔹Worker: {WORKER_NAME} | {WORKER_UUID}")
    print(f"🔹Tasks path: {TASKS_PATH}")
    print(f"🔹API URL: {API_URL}")
    try:
        main()
    except KeyboardInterrupt:
        print("🔥Interrupted by user.")