"""
╔══════════════════════════════════════════════════════════════════════════════╗
║  ALASK v2.0 Project                                                          ║
║  Task Manager - Multiprocess/Thread Management System                        ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  Version : 2.0.0                                                             ║
║  Date    : 2026-02-06                                                        ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  Classes:                                                                    ║
║    - TaskInfo    : Task metadata and IPC resources                           ║
║    - TaskManager : Task lifecycle management                                 ║
║    - TaskRuntime : Task execution wrapper with RMI handler                   ║
║    - RmiClient   : Remote method invocation client                           ║
║    - _RmiProxy   : Lazy proxy for remote variable access                     ║
╚══════════════════════════════════════════════════════════════════════════════╝
"""

from __future__ import annotations
from typing import Any, Dict, Optional
from dataclasses import dataclass, field
import threading
import multiprocessing
import time
import socket
import traceback
import queue
import atexit
import signal

from .task_performance import TimingRecorder, PerformanceCollector
from .task_signal import SignalBroker, SignalClient
from .sm_block import SmBlockHandler
from .task_decoration import rmi_task, has_run_flag, get_signal_handlers, TaskValidator

_MISSING = object()


# ─────────────────────────────────────────────────────────────────────────────
# Utility Functions
# ─────────────────────────────────────────────────────────────────────────────

def _check_port_in_use(port: int) -> bool:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        try:
            s.bind(('0.0.0.0', port))
            return False
        except OSError:
            return True


def _check_singleton(port: int) -> bool:
    if _check_port_in_use(port):
        print(f"\n{'='*60}\n[ERROR] Port {port} in use\n{'='*60}\n")
        return False
    return True


# ─────────────────────────────────────────────────────────────────────────────
# TaskInfo
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class TaskInfo:
    task_class_name: str
    task_id: str = ""
    mode: str = "thread"
    job_class: Optional[type] = None
    injection: Dict[str, Any] = field(default_factory=dict)
    auto_restart: bool = True
    restart_delay: float = 3.0
    request_q: Any = None
    response_q: Any = None
    stop_flag: Any = None
    rmi_rx: Any = None
    rmi_tx: Any = None
    rmi_fail: Any = None
    rmi_time: Any = None
    rmi_time_min: Any = None
    rmi_time_max: Any = None
    restart_cnt: Any = None
    rmi_methods: Any = None
    rmi_timing: Any = None

    def __post_init__(self):
        self.task_id = self.task_id or self.task_class_name
        if not self.job_class:
            n = self.task_class_name
            self.job_class = rmi_task.get(n)
            self.mode = rmi_task.attr(n, '_tc_mode', 'thread')
            self.auto_restart = rmi_task.attr(n, '_tc_restart', True)
            self.restart_delay = rmi_task.attr(n, '_tc_delay', 3.0)


# ─────────────────────────────────────────────────────────────────────────────
# _RmiProxy - Lazy proxy for remote variable access
# ─────────────────────────────────────────────────────────────────────────────

class _RmiProxy:
    __slots__ = ('_client', '_name', '_value', '_resolved')

    def __init__(self, client: 'RmiClient', name: str):
        sa = object.__setattr__
        sa(self, '_client', client)
        sa(self, '_name', name)
        sa(self, '_value', None)
        sa(self, '_resolved', False)

    def _resolve(self):
        if not self._resolved:
            sa = object.__setattr__
            sa(self, '_value', self._client._call('__getvar__', self._name))
            sa(self, '_resolved', True)
        return self._value

    def __call__(self, *a, **kw): return self._client._call(self._name, *a, **kw)

    def __eq__(self, o): return self._resolve() == o
    def __ne__(self, o): return self._resolve() != o
    def __lt__(self, o): return self._resolve() < o
    def __le__(self, o): return self._resolve() <= o
    def __gt__(self, o): return self._resolve() > o
    def __ge__(self, o): return self._resolve() >= o

    def __add__(self, o): return self._resolve() + o
    def __radd__(self, o): return o + self._resolve()
    def __sub__(self, o): return self._resolve() - o
    def __rsub__(self, o): return o - self._resolve()
    def __mul__(self, o): return self._resolve() * o
    def __rmul__(self, o): return o * self._resolve()
    def __truediv__(self, o): return self._resolve() / o
    def __floordiv__(self, o): return self._resolve() // o
    def __mod__(self, o): return self._resolve() % o

    def __neg__(self): return -self._resolve()
    def __pos__(self): return +self._resolve()
    def __abs__(self): return abs(self._resolve())
    def __int__(self): return int(self._resolve())
    def __float__(self): return float(self._resolve())
    def __str__(self): return str(self._resolve())
    def __bool__(self): return bool(self._resolve())
    def __repr__(self): return f"_RmiProxy({self._name})"

    def __len__(self): return len(self._resolve())
    def __getitem__(self, k): return self._resolve()[k]
    def __contains__(self, i): return i in self._resolve()
    def __iter__(self): return iter(self._resolve())


# ─────────────────────────────────────────────────────────────────────────────
# RmiClient - Remote Method Invocation Client
# ─────────────────────────────────────────────────────────────────────────────

class RmiClient:
    __slots__ = ('_id', '_req_q', '_res_q', '_timeout', '_tx_cnt', '_time_cnt', '_time_min', '_time_max', '_caller')

    def __init__(self, target_id: str, request_q, response_q, timeout: float = 2.0,
                 tx_counter=None, time_counter=None, time_min=None, time_max=None, caller: str = None):
        sa = object.__setattr__
        sa(self, '_id', target_id)
        sa(self, '_req_q', request_q)
        sa(self, '_res_q', response_q)
        sa(self, '_timeout', timeout)
        sa(self, '_tx_cnt', tx_counter)
        sa(self, '_time_cnt', time_counter)
        sa(self, '_time_min', time_min)
        sa(self, '_time_max', time_max)
        sa(self, '_caller', caller)

    def _call(self, m: str, *a, _chain: list = None, **kw):
        tx = self._tx_cnt
        if tx: tx.value += 1
        t0 = time.time()
        chain = list(_chain) if _chain else []
        if self._caller:
            chain.append(self._caller)
        self._req_q.put({'m': m, 'a': a, 'kw': kw, 'rq': self._res_q, 't0': t0, 'chain': chain})
        try:
            r = self._res_q.get(timeout=self._timeout)
            elapsed = (time.time() - t0) * 1000
            tc, tmin, tmax = self._time_cnt, self._time_min, self._time_max
            if tc: tc.value += elapsed
            if tmin and elapsed < tmin.value: tmin.value = elapsed
            if tmax and elapsed > tmax.value: tmax.value = elapsed
            if r.get('e'):
                err_msg = r['e']
                if r.get('trace'):
                    err_msg += f"\n--- Remote Traceback ---\n{r['trace']}"
                raise Exception(err_msg)
            return r.get('r')
        except queue.Empty:
            raise Exception(f"RmiClient({self._id}).{m}: Timeout after {self._timeout}s")
        except Exception as e:
            raise Exception(f"RmiClient({self._id}).{m}: {e}") from e

    def __getattr__(self, m: str):
        if m[0] == '_': raise AttributeError(m)
        return _RmiProxy(self, m)

    def __setattr__(self, m: str, value):
        if m[0] == '_':
            object.__setattr__(self, m, value)
        else:
            self._call('__setvar__', m, value)

    def __repr__(self): return f"RmiClient({self._id})"


# ─────────────────────────────────────────────────────────────────────────────
# TaskManager
# ─────────────────────────────────────────────────────────────────────────────

class TaskManager:
    __slots__ = ('_tasks', '_workers', '_mgr', '_monitor_port', '_monitor', '_gconfig',
                 '_exit_hook', '_measurement', '_performance', '_signal_broker', '_rmi_timeout', '__weakref__')

    def __init__(self, gconfig):
        missing = []
        if not gconfig.data_get("app_info"): missing.append("app_info")
        if not gconfig.data_get("platform_config"): missing.append("platform_config")
        if not gconfig.data_get("task_config"): missing.append("task_config")
        if missing:
            raise ValueError(f"[TaskManager] Missing config: {', '.join(missing)}")

        self._tasks: Dict[str, TaskInfo] = {}
        self._workers: Dict[str, Any] = {}
        self._mgr = multiprocessing.Manager()
        self._gconfig = gconfig
        self._monitor = None
        self._exit_hook = False

        platform_config = gconfig.data_get("platform_config", {})
        config = gconfig.data_get("task_config", {})
        monitor_cfg = platform_config.get("_monitor", {}) or config.get("_monitor", {})
        self._monitor_port = monitor_cfg.get("port") if monitor_cfg else None
        self._exit_hook = monitor_cfg.get("exit_hook", False) if monitor_cfg else False
        measurement_default = monitor_cfg.get("measurement", False) if monitor_cfg else False
        self._measurement = self._mgr.Value('b', measurement_default)
        rmi_cfg = platform_config.get("rmi", {})
        self._rmi_timeout = rmi_cfg.get("timeout", 2.0)
        mgr = self._mgr

        # LogTask 자동 등록
        log_cfg = platform_config.get("_log/log_task", {})
        if log_cfg:
            from . import task_log  # noqa: F401
            inj = {k: v for k, v in log_cfg.items() if k[0] != "@"}
            ti = TaskInfo(task_class_name="log_task", task_id="log_task", injection=inj)
            self._init_task_info(ti, mgr)
            self._tasks["log_task"] = ti

        # 사용자 Task 등록
        C, U, R = "\033[36m", "\033[4m", "\033[0m"  # Cyan, Underline, Reset
        config_path = getattr(gconfig, '_filepath', None)
        for key, cfg in config.items():
            if key.startswith("_"): continue
            # @import 처리
            if "@import" in cfg:
                try:
                    __import__(cfg["@import"])
                except ImportError as e:
                    print(f"[TaskManager] @import failed: {cfg['@import']} ({e})")
                    if config_path:
                        print(f"  {C}{U}{config_path}{R}  task: {key}")
            tid, tcn = key.split("/", 1) if "/" in key else (key, cfg.get("@task") or cfg.get("@job") or key)
            inj = {k: v for k, v in cfg.items() if k[0] != "@"}
            ti = TaskInfo(task_class_name=tcn, task_id=tid, injection=inj)
            self._init_task_info(ti, mgr)
            self._tasks[tid] = ti

        self._performance = PerformanceCollector(self._tasks, self._workers, self._measurement)
        self._signal_broker = SignalBroker(self._mgr)

    def _init_task_info(self, ti: TaskInfo, mgr):
        ti.request_q, ti.response_q = mgr.Queue(), mgr.Queue()
        ti.stop_flag = mgr.Value('b', False)
        ti.rmi_rx, ti.rmi_tx, ti.rmi_fail = mgr.Value('i', 0), mgr.Value('i', 0), mgr.Value('i', 0)
        ti.rmi_time, ti.rmi_time_min, ti.rmi_time_max = mgr.Value('d', 0.0), mgr.Value('d', float('inf')), mgr.Value('d', 0.0)
        ti.restart_cnt = mgr.Value('i', 0)
        ti.rmi_methods = mgr.dict()
        ti.rmi_timing = mgr.dict()

    @property
    def performance(self) -> PerformanceCollector:
        return self._performance

    def start_all(self, exit_on_duplicate: bool = True):
        if self._monitor_port and not _check_singleton(self._monitor_port):
            raise SystemExit(1) if exit_on_duplicate else RuntimeError(f"Port {self._monitor_port} in use")

        validator = TaskValidator()
        for tid, ti in self._tasks.items():
            validator.add(tid, ti.task_class_name)
        if not validator.validate():
            validator.print_errors()
            raise SystemExit(1)

        self._print_task_table()
        SmBlockHandler.creation(self._gconfig, self._mgr, self._tasks)

        if self._signal_broker:
            for tid in self._tasks.keys():
                self._signal_broker.register(tid)

        for tid, ti in self._tasks.items():
            if tid not in self._workers:
                w = TaskRuntime(ti, self._tasks, measurement=self._measurement, signal_broker=self._signal_broker, rmi_timeout=self._rmi_timeout)
                h = multiprocessing.Process(target=w.run, daemon=True) if ti.mode == 'process' else threading.Thread(target=w.run, daemon=True)
                h.start()
                self._workers[tid] = h

        if self._monitor_port and not self._monitor:
            from .task_monitor import TaskMonitor
            self._monitor = TaskMonitor(self, self._monitor_port, gconfig=self._gconfig)
            self._monitor.start()

        if self._exit_hook:
            atexit.register(self.stop_all)
            if threading.current_thread() is threading.main_thread():
                signal.signal(signal.SIGINT, lambda s, f: self.stop_all())

    def stop_all(self, skip_monitor: bool = False):
        if self._monitor and not skip_monitor:
            self._monitor.stop()
            self._monitor = None
        for ti in self._tasks.values():
            ti.stop_flag.value = True
        for tid, w in self._workers.items():
            w.join(timeout=4.0)
            if w.is_alive() and hasattr(w, 'terminate'):
                w.terminate()
        self._workers.clear()
        SmBlockHandler.close_all()

    def clear_stats(self):
        inf = float('inf')
        for ti in self._tasks.values():
            ti.rmi_rx.value = ti.rmi_tx.value = ti.rmi_fail.value = 0
            ti.rmi_time.value = ti.rmi_time_max.value = 0.0
            ti.rmi_time_min.value = inf
            ti.restart_cnt.value = 0
            ti.rmi_methods.clear()
            ti.rmi_timing.clear()

    def get_status(self) -> Dict[str, Dict[str, Any]]:
        result = {}
        inf = float('inf')
        for tid, ti in self._tasks.items():
            w = self._workers.get(tid)
            tx, tt, tmin = ti.rmi_tx.value, ti.rmi_time.value, ti.rmi_time_min.value
            result[tid] = {
                'class': ti.task_class_name, 'mode': ti.mode,
                'alive': w.is_alive() if w else False,
                'rmi_rx': ti.rmi_rx.value, 'rmi_tx': tx, 'rmi_fail': ti.rmi_fail.value,
                'rmi_avg': tt / tx if tx else 0.0,
                'rmi_min': 0.0 if tmin == inf else tmin,
                'rmi_max': ti.rmi_time_max.value,
                'rxq_size': ti.request_q.qsize(), 'txq_size': ti.response_q.qsize(),
                'restart': ti.restart_cnt.value,
                'rmi_methods': dict(ti.rmi_methods),
            }
        return result

    def get_client(self, task_id: str) -> RmiClient:
        """외부에서 Task에 RMI 호출하기 위한 클라이언트 생성"""
        ti = self._tasks.get(task_id)
        if not ti:
            raise ValueError(f"Task '{task_id}' not found")
        response_q = self._mgr.Queue()
        return RmiClient(
            target_id=task_id,
            request_q=ti.request_q,
            response_q=response_q,
            timeout=self._rmi_timeout
        )

    def get_signal_client(self, client_id: str = "_external") -> 'SignalClient':
        """외부에서 Signal 수신하기 위한 SignalClient 생성

        Usage:
            signal_client = manager.get_signal_client("my_ui")
            signal_client.on("camera_connected", lambda sig: print(sig.data))
            signal_client.start()  # 수신 시작
            # ...
            signal_client.stop()   # 종료 시
        """
        if not self._signal_broker:
            raise RuntimeError("SignalBroker not initialized")
        return SignalClient(self._signal_broker, client_id)

    def _print_task_table(self):
        headers = ['ID', 'Mode', 'Class', 'rmi_name', 'rmi_run']
        rows = []
        for tid, ti in self._tasks.items():
            cls = ti.job_class
            rmi_name = getattr(cls, '_tc_name', '-') if cls else '-'
            rmi_run = self._find_rmi_run_method(cls) if cls else '-'
            rows.append([tid, ti.mode, ti.task_class_name, rmi_name, rmi_run])
        widths = [max(len(h), max((len(str(r[i])) for r in rows), default=0)) for i, h in enumerate(headers)]
        top    = '┌' + '┬'.join('─' * (w + 2) for w in widths) + '┐'
        mid    = '├' + '┼'.join('─' * (w + 2) for w in widths) + '┤'
        bottom = '└' + '┴'.join('─' * (w + 2) for w in widths) + '┘'
        print('\n' + top)
        print('│' + '│'.join(f' {h:<{w}} ' for h, w in zip(headers, widths)) + '│')
        print(mid)
        for row in rows:
            print('│' + '│'.join(f' {str(c):<{w}} ' for c, w in zip(row, widths)) + '│')
        print(bottom + '\n')

    def _find_rmi_run_method(self, cls: type) -> str:
        for name in dir(cls):
            if name.startswith('__'): continue
            method = getattr(cls, name, None)
            if callable(method) and has_run_flag(method):
                return name
        if hasattr(cls, 'run') and callable(getattr(cls, 'run', None)):
            return 'run'
        if hasattr(cls, 'rmi_loop') and callable(getattr(cls, 'rmi_loop', None)):
            return 'rmi_loop'
        return '-'

    def __enter__(self): self.start_all(); return self
    def __exit__(self, *_): self.stop_all(); return False


# ─────────────────────────────────────────────────────────────────────────────
# TaskRuntime
# ─────────────────────────────────────────────────────────────────────────────

class TaskRuntime:
    __slots__ = ('ti', 'tasks', 'job', '_run', '_th', '_measurement', '_timing_recorder',
                 '_signal_broker', '_signal_client', '_rmi_timeout', 'log')

    def __init__(self, ti: TaskInfo, tasks: Dict[str, TaskInfo], measurement=None, signal_broker=None, rmi_timeout: float = 2.0):
        self.ti, self.tasks, self.job, self._run, self._th = ti, tasks, None, True, None
        self._measurement = measurement
        self._rmi_timeout = rmi_timeout
        self._timing_recorder = TimingRecorder(ti.rmi_methods, ti.rmi_timing)
        self._signal_broker = signal_broker
        self._signal_client = None
        from .task_log import RmiLogger
        self.log = RmiLogger(self)

    def run(self):
        ti = self.ti
        if not ti.job_class:
            print(f"[W:{ti.task_id}] no job_class"); return
        self.job = ti.job_class()  # __init__(self) 호출
        self.job.runtime = self   # runtime 속성 주입
        self.job._input_queue = queue.Queue()  # 기본 입력 큐 주입
        # on_input 메서드 자동 생성 (미정의 시)
        if not hasattr(self.job, 'on_input'):
            self.job.on_input = lambda data: self.job._input_queue.put_nowait(data)
        if self._signal_broker:
            self._signal_client = SignalClient(self._signal_broker, ti.task_id)
            self._signal_client.start()
            self._auto_subscribe_signals()
        self._th = threading.Thread(target=self._rmi_handler, daemon=True)
        self._th.start()
        self._inject_dependencies()
        if not self._validate_injection():
            if self._signal_client: self._signal_client.stop()
            return
        print(f"[W:{ti.task_id}] started")
        self._run_main_loop()
        if self._signal_client:
            self._signal_client.off_all()
            self._signal_client.stop()

    def _validate_injection(self) -> bool:
        ti, job = self.ti, self.job
        errors = []
        if 'smblock' in ti.injection and (not hasattr(job, 'smblock') or job.smblock is None):
            errors.append("smblock not injected")
        if 'next' in ti.injection and (not hasattr(job, 'next') or job.next is None):
            errors.append("next not injected")
        if errors:
            print(f"[W:{ti.task_id}] ERROR: {', '.join(errors)}")
            time.sleep(3)
            return False
        return True

    def _auto_subscribe_signals(self):
        if not self._signal_client or not self.job: return
        for signal_name, method in get_signal_handlers(self.job):
            self._signal_client.on(signal_name, method)

    def _rmi_handler(self):
        ti, q, job = self.ti, self.ti.request_q, self.job
        stop, rx, fail = ti.stop_flag, ti.rmi_rx, ti.rmi_fail
        recorder = self._timing_recorder

        while self._run and not stop.value:
            try:
                req = q.get(timeout=0.5)
            except queue.Empty:
                if self._measurement and self._measurement.value: recorder.flush()
                continue
            except Exception:
                continue
            if stop.value: break

            t_recv = time.time()
            rx.value += 1
            m, a, kw, rq = req['m'], req.get('a', ()), req.get('kw', {}), req.get('rq')
            t0, chain = req.get('t0'), req.get('chain', [])

            if ti.task_id in chain:
                chain_str = ' → '.join(chain + [ti.task_id])
                if rq: rq.put({'e': f"Circular RMI: {chain_str}"})
                fail.value += 1
                continue

            measure = m[0] != '_' and self._measurement and self._measurement.value
            t_func_start = time.time()

            try:
                if m[0] != '_':
                    fn = getattr(job, m, _MISSING)
                    if fn is _MISSING:
                        res = {'e': f"'{m}' not found"}; fail.value += 1
                    else:
                        res = {'r': fn(*a, **kw) if callable(fn) else fn}
                elif m == '__getvar__':
                    val = getattr(job, a[0], _MISSING)
                    res = {'e': f"'{a[0]}' not found"} if val is _MISSING else {'r': val() if callable(val) else val}
                    if val is _MISSING: fail.value += 1
                elif m == '__setvar__':
                    setattr(job, a[0], a[1]); res = {'r': True}
                elif m == '__set_config__':
                    setattr(job, a[0], a[1]); res = {'r': True}
                elif m == '__get_config__':
                    res = {'r': {k: getattr(job, k, None) for k in (a[0] if a else [])}}
                else:
                    res = {'e': f"Unknown '{m}'"}; fail.value += 1
            except Exception as e:
                res = {'e': str(e), 'trace': traceback.format_exc()}; fail.value += 1

            if measure:
                t_func_end = time.time()
                if t0:
                    recorder.record(m, (t_recv - t0) * 1000, (t_func_end - t_func_start) * 1000)
                else:
                    recorder.record_count_only(m)

            if rq:
                try: rq.put(res)
                except: pass

        recorder.flush()

    def _inject_dependencies(self):
        job, tasks, ti = self.job, self.tasks, self.ti
        SmBlockHandler.injection(self, tasks)
        for k, v in ti.injection.items():
            if isinstance(v, dict) and "_smblock" in v: continue
            if isinstance(v, str) and v[:5] == "task:":
                t = tasks.get(v[5:])
                if t:
                    setattr(job, k, RmiClient(
                        v[5:], t.request_q, ti.response_q, timeout=self._rmi_timeout,
                        tx_counter=ti.rmi_tx, time_counter=ti.rmi_time,
                        time_min=ti.rmi_time_min, time_max=ti.rmi_time_max,
                        caller=ti.task_id
                    ))
            else:
                setattr(job, k, v)

    def _run_main_loop(self):
        ti, job = self.ti, self.job

        # @rmi_run 데코레이터 탐지
        loop_method = None
        for name in dir(job):
            if name.startswith('__'): continue
            method = getattr(job, name, None)
            if callable(method) and has_run_flag(method):
                loop_method = method
                print(f"[W:{ti.task_id}] @rmi_run detected: {name}")
                break

        # Fallback: run() 또는 rmi_loop()
        if loop_method is None:
            loop_method = getattr(job, 'run', None) or getattr(job, 'rmi_loop', None)
            if loop_method:
                print(f"[W:{ti.task_id}] fallback: {loop_method.__name__}")

        if not loop_method:
            print(f"[W:{ti.task_id}] no run method found!")
            return

        print(f"[W:{ti.task_id}] executing: {loop_method.__name__}")
        while self._run:
            try:
                loop_method()  # run(self) - self.runtime으로 접근
                break
            except Exception as e:
                print(f"[W:{ti.task_id}] {loop_method.__name__}: {e}")
                traceback.print_exc()
                if not ti.auto_restart: break
                ti.restart_cnt.value += 1
                time.sleep(ti.restart_delay)

    def should_stop(self) -> bool:
        return not self._run or self.ti.stop_flag.value

    @property
    def name(self) -> str:
        return self.ti.task_id

    def get_client(self, task_id: str) -> RmiClient:
        t = self.tasks.get(task_id)
        if not t: raise ValueError(f"Task '{task_id}' not found")
        return RmiClient(
            task_id, t.request_q, self.ti.response_q, timeout=self._rmi_timeout,
            tx_counter=self.ti.rmi_tx, time_counter=self.ti.rmi_time,
            time_min=self.ti.rmi_time_min, time_max=self.ti.rmi_time_max,
            caller=self.ti.task_id
        )

    @property
    def signal(self) -> SignalClient:
        return self._signal_client
