"""Runtime loops for commander and worker roles."""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from dataclasses import replace
from datetime import datetime
from datetime import timezone
import json
import os
from pathlib import Path
import queue
import signal
import threading
import time
import uuid

from a2a_agent.client import MqttClient
from a2a_agent.config import AgentConfig
from a2a_agent.config import BrokerConfig
from a2a_agent.config import DeviceRole
from a2a_agent.config import ShareConfig
from a2a_agent.config import default_config_path
from a2a_agent.config import load_config
from a2a_agent.config import save_config
from a2a_agent.config import share_config_from_env
from a2a_agent.executor import execute_text
from a2a_agent.executor import ExecutorResult
from a2a_agent.executor import selected_executor
from a2a_agent.invite import validate_pairing_code
from a2a_agent.pairing import consume_pending_pairing
from a2a_agent.protocol import AgentCard
from a2a_agent.protocol import FilePart
from a2a_agent.protocol import Message
from a2a_agent.protocol import PairingReply
from a2a_agent.protocol import PairingRequest
from a2a_agent.protocol import Task
from a2a_agent.protocol import build_file_part
from a2a_agent.protocol import parse_file_part
from a2a_agent.share import FileStationClient
from a2a_agent.share import ShareError
from a2a_agent.updater import start_runtime_auto_update_monitor
from a2a_agent import __version__
from a2a_agent import topics


PAIRING_REPLY_TIMEOUT_SECONDS = 30
SEND_TIMEOUT_SECONDS = int(os.environ.get("A2A_SEND_TIMEOUT", "120"))


@dataclass(frozen=True)
class SendResult:
    task_id: str
    status: str
    text: str
    downloaded_files: tuple[str, ...] = ()
    file_errors: tuple[str, ...] = ()


def join_pairing(code: str, *, timeout: int = PAIRING_REPLY_TIMEOUT_SECONDS) -> str:
    normalized = validate_pairing_code(code)
    broker = BrokerConfig()
    reply_event = threading.Event()
    reply_holder: dict[str, PairingReply] = {}
    client = MqttClient(client_id=f"a2a-join-{uuid.uuid4().hex}")
    reply_topic = topics.pairing_reply_topic(normalized)

    def handle_reply(topic: str, payload: bytes, message: object) -> None:
        try:
            reply_holder["reply"] = PairingReply.from_json(payload)
        except (ValueError, KeyError, json.JSONDecodeError) as exc:
            reply_holder["reply"] = PairingReply(ok=False, error=str(exc))
        reply_event.set()

    client.connect(broker.host, broker.port, tls=broker.tls)
    client.loop_start()
    try:
        if not client.wait_connected():
            raise RuntimeError(f"could not connect to broker {broker.host}:{broker.port}")
        client.subscribe(reply_topic, handle_reply)
        time.sleep(0.2)
        _publish_and_wait(
            client,
            topics.pairing_request_topic(),
            PairingRequest(code=normalized).to_json(),
        )
        if not reply_event.wait(timeout):
            raise TimeoutError(f"pairing reply timed out after {timeout} seconds")

        reply = reply_holder["reply"]
        if not reply.ok:
            raise RuntimeError(reply.error or "pairing failed")
        if not reply.room or not reply.device_name or not reply.token or not reply.broker:
            raise RuntimeError("pairing reply missed required worker settings")

        broker_payload = reply.broker
        saved_path = save_config(
            AgentConfig(
                room=reply.room,
                device_name=reply.device_name,
                broker=BrokerConfig.from_dict(broker_payload),
                role="worker",
                token=reply.token,
                share=ShareConfig.from_dict(reply.share) if reply.share else None,
            )
        )
        return str(saved_path)
    finally:
        client.disconnect()
        client.loop_stop()


def run_commander(stop_event: threading.Event | None = None) -> None:
    config = load_config()
    if config.role != "commander":
        raise RuntimeError("commander run requires commander config")

    stop = stop_event or threading.Event()
    _install_signal_handlers(stop)
    status_topic = topics.agent_status_topic(config.room, config.device_name)

    def announce_online(connected_client: MqttClient) -> None:
        if stop.is_set():
            return
        connected_client.publish(
            status_topic,
            _status_payload(config, "online"),
            qos=1,
            retain=True,
        )

    client = _connected_client(config, "commander", status_topic, announce_online)

    def handle_pairing_request(topic: str, payload: bytes, message: object) -> None:
        try:
            request = PairingRequest.from_json(payload)
        except (ValueError, KeyError, json.JSONDecodeError) as exc:
            print(f"pairing request ignored: {exc}", flush=True)
            return

        current = load_config()
        try:
            share_config = share_config_from_env()
        except ValueError as exc:
            print(f"pairing share config skipped: {exc}", flush=True)
            share_config = None
        reply, updated = consume_pending_pairing(
            current,
            request.code,
            share=share_config,
        )
        # 페어링에 성공하면 그 기기를 커맨더 config에 영구 등록한다.
        # (재시작 후에도 어떤 기기가 있었는지 기억하고, 개별 지시가 가능하게)
        if reply.ok and reply.device_name and reply.device_name not in updated.devices:
            new_devices = dict(updated.devices)
            new_devices[reply.device_name] = DeviceRole(
                role="",
                description=f"paired {datetime.now(timezone.utc).isoformat()}",
            )
            updated = replace(updated, devices=new_devices)
        if updated != current:
            save_config(updated)
        _publish_and_wait(client, topics.pairing_reply_topic(request.code), reply.to_json())
        device = reply.device_name or "-"
        print(
            f"pairing reply sent: code={request.code} ok={str(reply.ok).lower()} "
            f"device={device}",
            flush=True,
        )

    def handle_worker_status(topic: str, payload: bytes, message: object) -> None:
        try:
            decoded = payload.decode("utf-8")
        except UnicodeDecodeError:
            decoded = "<binary>"
        print(f"worker status: topic={topic} payload={decoded}", flush=True)

    try:
        client.subscribe(topics.pairing_request_topic(), handle_pairing_request)
        client.subscribe(topics.agent_status_filter(config.room), handle_worker_status)
        _publish_and_wait(client, status_topic, _status_payload(config, "online"), retain=True)
        print(
            f"commander run ready: broker={config.broker.host}:{config.broker.port} "
            f"room={config.room} device={config.device_name}",
            flush=True,
        )
        _wait_until_stopped(stop)
    finally:
        stop.set()
        _publish_and_wait(client, status_topic, _status_payload(config, "offline"), retain=True)
        client.disconnect()
        client.loop_stop()


def run_worker(stop_event: threading.Event | None = None) -> None:
    config = load_config()
    if config.role != "worker":
        raise RuntimeError("worker run requires worker config")
    if not config.token:
        raise RuntimeError("worker config has no token; run join first")

    stop = stop_event or threading.Event()
    _install_signal_handlers(stop)
    status_topic = topics.agent_status_topic(config.room, config.device_name)
    card_topic = topics.agent_card_topic(config.room, config.device_name)
    card_payload = AgentCard(name=config.device_name).to_json()

    def announce_online(connected_client: MqttClient) -> None:
        if stop.is_set():
            return
        connected_client.publish(
            status_topic,
            _status_payload(config, "online"),
            qos=1,
            retain=True,
        )
        connected_client.publish(card_topic, card_payload, qos=1, retain=True)

    client = _connected_client(config, "worker", status_topic, announce_online)
    request_topic = topics.cmd_request_topic(config.room, config.device_name)
    response_topic = topics.cmd_response_topic(config.room, config.device_name)
    task_status_topic = topics.cmd_task_status_topic(config.room, config.device_name)
    task_queue: queue.Queue[Task] = queue.Queue()

    def handle_request(topic: str, payload: bytes, message: object) -> None:
        if stop.is_set():
            return
        try:
            task = Task.from_json(payload)
        except (ValueError, KeyError, json.JSONDecodeError) as exc:
            print(f"task ignored: {exc}", flush=True)
            return
        if stop.is_set():
            return

        print(f"task received: id={task.id} from={task.from_}", flush=True)
        working = replace(
            task,
            from_=config.device_name,
            to=task.from_,
            status="working",
            share=None,
        )
        # paho 네트워크 루프 콜백 안에서는 QoS 완료를 기다리면 같은 루프를
        # 막을 수 있으므로 발행만 요청하고 즉시 큐로 넘긴다.
        client.publish(task_status_topic, working.to_json(), qos=1, retain=False)
        task_queue.put(task)

    def process_task(task: Task) -> None:
        nonlocal config
        task_config = _cache_task_share_config(config, task)
        if task_config != config:
            config = task_config

        downloaded_files, input_file_errors = _download_input_files(task_config, task)
        out_dir = _task_out_dir(task.id)
        out_dir.mkdir(parents=True, exist_ok=True)
        prompt = _worker_prompt(task, downloaded_files, out_dir, input_file_errors)
        try:
            result = execute_text(prompt, task_config)
        except Exception as exc:  # 워커가 어떤 실행기 오류에도 죽지 않게 한다.
            result = ExecutorResult(ok=False, text=f"executor crashed: {exc}")
        status = "completed" if result.ok else "failed"
        response_parts = [{"type": "text", "text": result.text}]
        response_parts.extend(_upload_output_files(task_config, task.id, out_dir))
        response = Task(
            id=task.id,
            from_=config.device_name,
            to=task.from_,
            status=status,
            messages=[
                Message(
                    role="agent",
                    parts=response_parts,
                )
            ],
        )
        _publish_and_wait(client, response_topic, response.to_json())
        done = replace(response, messages=[])
        _publish_and_wait(client, task_status_topic, done.to_json())
        print(f"task responded: id={task.id} status={status}", flush=True)

    def publish_task_failure(task: Task, exc: Exception) -> None:
        failed = Task(
            id=task.id,
            from_=config.device_name,
            to=task.from_,
            status="failed",
            messages=[
                Message(
                    role="agent",
                    parts=[{"type": "text", "text": f"task processing crashed: {exc}"}],
                )
            ],
        )
        publications = (
            (response_topic, failed.to_json(), "response"),
            (task_status_topic, replace(failed, messages=[]).to_json(), "task status"),
        )
        for publish_topic, publish_payload, label in publications:
            try:
                _publish_and_wait(client, publish_topic, publish_payload)
            except Exception as publish_exc:
                print(
                    f"task failure {label} publish skipped: id={task.id}: {publish_exc}",
                    flush=True,
                )

    def process_tasks() -> None:
        while not stop.is_set():
            try:
                task = task_queue.get(timeout=0.2)
            except queue.Empty:
                continue
            try:
                process_task(task)
            except Exception as exc:
                # 개별 태스크 실패가 처리 스레드와 이후 태스크까지 죽이지 않게 한다.
                print(f"task processing crashed: id={task.id}: {exc}", flush=True)
                publish_task_failure(task, exc)
            finally:
                task_queue.task_done()

    task_processor = threading.Thread(
        target=process_tasks,
        name="a2a-worker-task-processor",
        daemon=True,
    )
    task_processor.start()

    try:
        client.subscribe(request_topic, handle_request)
        _publish_and_wait(
            client,
            card_topic,
            card_payload,
            retain=True,
        )
        _publish_and_wait(client, status_topic, _status_payload(config, "online"), retain=True)
        start_runtime_auto_update_monitor(stop)
        print(
            f"worker run ready: broker={config.broker.host}:{config.broker.port} "
            f"room={config.room} device={config.device_name} "
            f"executor={selected_executor(config)}",
            flush=True,
        )
        _wait_until_stopped(stop)
    finally:
        stop.set()
        task_processor.join(timeout=1)
        _publish_and_wait(client, status_topic, _status_payload(config, "offline"), retain=True)
        client.disconnect()
        client.loop_stop()


def send_task(
    device_name: str,
    text: str,
    *,
    files: list[Path] | None = None,
    download_dir: Path | None = None,
    timeout: int = SEND_TIMEOUT_SECONDS,
) -> SendResult:
    config = load_config()
    if config.role != "commander":
        raise RuntimeError("send requires commander config")
    resolved_device_name = resolve_device_name(config, device_name)

    task_id = f"task-{uuid.uuid4().hex}"
    local_files = files or []
    try:
        share_config = _commander_share_config(config)
    except ValueError as exc:
        if local_files:
            raise RuntimeError(str(exc)) from exc
        share_config = None
    request_parts = [{"type": "text", "text": text}]
    request_parts.extend(
        _upload_request_files(
            config,
            task_id,
            local_files,
            share_config=share_config,
        )
    )
    response_event = threading.Event()
    response_holder: dict[str, Task] = {}
    client = MqttClient(client_id=f"a2a-send-{uuid.uuid4().hex}")
    response_topic = topics.cmd_response_topic(config.room, resolved_device_name)

    def handle_response(topic: str, payload: bytes, message: object) -> None:
        try:
            task = Task.from_json(payload)
        except (ValueError, KeyError, json.JSONDecodeError):
            return
        if task.id != task_id:
            return
        response_holder["task"] = task
        response_event.set()

    client.connect(config.broker.host, config.broker.port, tls=config.broker.tls)
    client.loop_start()
    try:
        if not client.wait_connected():
            raise RuntimeError(
                f"could not connect to broker {config.broker.host}:{config.broker.port}"
            )
        client.subscribe(response_topic, handle_response)
        time.sleep(0.2)
        request = Task(
            id=task_id,
            from_=config.device_name,
            to=resolved_device_name,
            status="pending",
            messages=[
                Message(role="user", parts=request_parts),
            ],
            # 보안: share 자격은 페어링 때처럼 Task에도 평문 MQTT를 탄다.
            # 범위가 좁은 계정 전제이며 이후 TLS 강화 대상이다.
            share=share_config.to_dict() if share_config is not None else None,
        )
        _publish_and_wait(
            client,
            topics.cmd_request_topic(config.room, resolved_device_name),
            request.to_json(),
        )
        if not response_event.wait(timeout):
            raise TimeoutError(
                f"response from {resolved_device_name} timed out after {timeout} seconds"
            )

        response = response_holder["task"]
        downloaded_files, file_errors = _download_response_files(
            config,
            response,
            download_dir or _a2a_home() / "downloads" / task_id,
        )
        return SendResult(
            task_id=response.id,
            status=response.status,
            text=_task_text(response),
            downloaded_files=tuple(str(path) for path in downloaded_files),
            file_errors=tuple(file_errors),
        )
    finally:
        client.disconnect()
        client.loop_stop()


def _upload_request_files(
    config: AgentConfig,
    task_id: str,
    files: list[Path],
    *,
    share_config: ShareConfig | None = None,
) -> list[dict[str, object]]:
    if not files:
        return []
    if share_config is None:
        try:
            share_config = _commander_share_config(config)
        except ValueError as exc:
            raise RuntimeError(str(exc)) from exc
    if share_config is None:
        raise RuntimeError("file upload requires A2A_SHARE_USER and A2A_SHARE_PASSWORD")

    client = FileStationClient(share_config)
    parts: list[dict[str, object]] = []
    try:
        client.login()
        remote_dir = f"{task_id}"
        for local_path in files:
            source = Path(local_path)
            remote_path = client.upload(source, remote_dir)
            parts.append(
                build_file_part(
                    name=source.name,
                    path=client.relative_path(remote_path),
                    size=source.stat().st_size,
                )
            )
    except ShareError as exc:
        raise RuntimeError(f"file upload failed: {exc}") from exc
    finally:
        client.logout()
    return parts


def _download_response_files(
    config: AgentConfig,
    task: Task,
    download_dir: Path,
) -> tuple[list[Path], list[str]]:
    file_parts = _task_file_parts(task)
    if not file_parts:
        return [], []
    try:
        share_config = _commander_share_config(config)
    except ValueError as exc:
        return [], [str(exc)]
    if share_config is None:
        return [], ["file response received but share config is not available"]
    return _download_file_parts(share_config, file_parts, download_dir)


def _cache_task_share_config(config: AgentConfig, task: Task) -> AgentConfig:
    if config.share is not None or task.share is None:
        return config
    try:
        share_config = ShareConfig.from_dict(task.share)
    except ValueError as exc:
        print(f"task share ignored: {exc}", flush=True)
        return config

    updated = replace(config, share=share_config)
    try:
        save_config(updated)
        print("task share config cached", flush=True)
    except OSError as exc:
        print(f"task share config cache skipped: {exc}", flush=True)
    return updated


def _download_input_files(
    config: AgentConfig,
    task: Task,
) -> tuple[list[Path], list[str]]:
    file_parts = _task_file_parts(task)
    if not file_parts:
        return [], []
    if config.share is None:
        return [], ["input files were not downloaded because share config is missing"]
    return _download_file_parts(config.share, file_parts, _task_in_dir(task.id))


def _download_file_parts(
    share_config: ShareConfig,
    file_parts: list[FilePart],
    target_dir: Path,
) -> tuple[list[Path], list[str]]:
    target_dir.mkdir(parents=True, exist_ok=True)
    client = FileStationClient(share_config)
    downloaded: list[Path] = []
    errors: list[str] = []
    try:
        client.login()
        for part in file_parts:
            local_path = target_dir / _safe_local_name(part.name, part.path)
            try:
                downloaded.append(client.download(part.path, local_path))
            except ShareError as exc:
                errors.append(f"{part.name}: {exc}")
    except ShareError as exc:
        errors.append(str(exc))
    finally:
        client.logout()
    return downloaded, errors


def _upload_output_files(
    config: AgentConfig,
    task_id: str,
    out_dir: Path,
) -> list[dict[str, object]]:
    if config.share is None or not out_dir.exists():
        return []

    files = sorted(path for path in out_dir.rglob("*") if path.is_file())
    if not files:
        return []

    client = FileStationClient(config.share)
    parts: list[dict[str, object]] = []
    try:
        client.login()
        for local_path in files:
            relative = local_path.relative_to(out_dir)
            remote_dir = f"{task_id}/out"
            if relative.parent != Path("."):
                remote_dir = f"{remote_dir}/{relative.parent.as_posix()}"
            try:
                remote_path = client.upload(local_path, remote_dir)
            except ShareError as exc:
                print(f"output file upload skipped: {relative.as_posix()}: {exc}", flush=True)
                continue
            parts.append(
                build_file_part(
                    name=relative.as_posix(),
                    path=client.relative_path(remote_path),
                    size=local_path.stat().st_size,
                )
            )
    except ShareError as exc:
        print(f"output file upload skipped: {exc}", flush=True)
    finally:
        client.logout()
    return parts


def _task_file_parts(task: Task) -> list[FilePart]:
    parts: list[FilePart] = []
    for message in task.messages:
        for part in message.parts:
            if part.get("type") != "file":
                continue
            try:
                parts.append(parse_file_part(part))
            except (TypeError, ValueError) as exc:
                print(f"file part ignored: {exc}", flush=True)
    return parts


def _worker_prompt(
    task: Task,
    downloaded_files: list[Path],
    out_dir: Path,
    file_errors: list[str],
) -> str:
    lines = [_task_text(task).strip()]
    if downloaded_files:
        paths = " ".join(str(path) for path in downloaded_files)
        lines.append(f"첨부 파일: {paths}")
    elif file_errors:
        lines.append("첨부 파일 다운로드 실패: " + "; ".join(file_errors))
    lines.append(f"결과 파일 저장 폴더: {out_dir}")
    lines.append("결과 파일이 있으면 위 폴더에 저장하세요.")
    return "\n".join(line for line in lines if line)


def _commander_share_config(config: AgentConfig) -> ShareConfig | None:
    env_config = share_config_from_env()
    return env_config or config.share


def _safe_local_name(name: str, remote_path: str) -> str:
    candidate = name.replace("\\", "/").split("/")[-1]
    if candidate:
        return candidate
    fallback = remote_path.replace("\\", "/").split("/")[-1]
    return fallback or "downloaded-file"


def _a2a_home() -> Path:
    return default_config_path().parent


def _task_in_dir(task_id: str) -> Path:
    return _a2a_home() / "tasks" / task_id / "in"


def _task_out_dir(task_id: str) -> Path:
    return _a2a_home() / "tasks" / task_id / "out"


def resolve_device_name(config: AgentConfig, requested_name: str) -> str:
    if requested_name in config.aliases:
        return config.aliases[requested_name]
    if requested_name in config.devices or not config.devices:
        return requested_name
    aliases = ", ".join(sorted(config.aliases)) or "none"
    devices = ", ".join(sorted(config.devices)) or "none"
    raise RuntimeError(
        f"unknown device or alias: {requested_name} "
        f"(aliases: {aliases}; devices: {devices})"
    )


def _connected_client(
    config: AgentConfig,
    role: str,
    status_topic: str,
    on_connect_hook: Callable[[MqttClient], None] | None = None,
) -> MqttClient:
    client = MqttClient(client_id=f"a2a-{role}-{uuid.uuid4().hex}")
    client.set_last_will(
        status_topic,
        _status_payload(config, "offline"),
        retain=True,
    )
    if on_connect_hook is not None:
        client.add_on_connect_hook(lambda: on_connect_hook(client))
    client.connect(config.broker.host, config.broker.port, tls=config.broker.tls)
    client.loop_start()
    if not client.wait_connected():
        raise RuntimeError(f"could not connect to broker {config.broker.host}:{config.broker.port}")
    return client


def _publish_and_wait(
    client: MqttClient,
    topic: str,
    payload: str,
    *,
    retain: bool = False,
) -> None:
    info = client.publish(topic, payload, qos=1, retain=retain)
    info.wait_for_publish(timeout=10)


def _task_text(task: Task) -> str:
    parts: list[str] = []
    for message in task.messages:
        for part in message.parts:
            if part.get("type") == "text" and part.get("text") is not None:
                parts.append(str(part["text"]))
    return "\n".join(parts)


def _status_payload(config: AgentConfig, status: str) -> str:
    payload = {
        "device_name": config.device_name,
        "role": config.role,
        "version": __version__,
        "executor": selected_executor(config),
        "status": status,
        "ts": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
    }
    return json.dumps(payload, ensure_ascii=False, sort_keys=True)


def _wait_until_stopped(stop_event: threading.Event) -> None:
    try:
        while not stop_event.wait(1):
            pass
    except KeyboardInterrupt:
        stop_event.set()


def _install_signal_handlers(stop_event: threading.Event) -> None:
    def handle_signal(signum: int, frame: object) -> None:
        stop_event.set()

    try:
        signal.signal(signal.SIGTERM, handle_signal)
        signal.signal(signal.SIGINT, handle_signal)
    except ValueError:
        return
