from __future__ import annotations

import argparse
import base64
import hashlib
import json
import os
import platform
import queue
import subprocess
import sys
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import Any
from urllib.error import HTTPError
from urllib.error import URLError
from urllib.request import Request
from urllib.request import urlopen


ARTIFACT_NAMES = (
    "lstm_forecaster.json",
    "torch_retrain_guard.json",
    "torch_threshold_calibration.json",
)


def main() -> None:
    args = parse_args()
    repo_root = Path(args.repo_root).resolve()
    runtime_dir = repo_root / "runtime"
    runtime_dir.mkdir(parents=True, exist_ok=True)
    log_path = Path(args.log_file).resolve() if args.log_file else runtime_dir / "windows_training_agent.log"

    log(log_path, f"TradeBot Windows training agent started for {args.api_base_url}")
    while True:
        try:
            poll_once(args, repo_root, runtime_dir, log_path)
        except Exception as exc:  # noqa: BLE001 - agent must keep running.
            log(log_path, f"ERROR: {exc}")
        if args.once:
            break
        time.sleep(max(5, args.poll_seconds))


def poll_once(args: argparse.Namespace, repo_root: Path, runtime_dir: Path, log_path: Path) -> None:
    worker = worker_payload(args, repo_root)
    api_json(args, "/api/training/heartbeat", worker)
    claim = api_json(args, "/api/training/claim", worker)
    if not claim.get("claimed"):
        return
    job = claim.get("job") if isinstance(claim.get("job"), dict) else {}
    job_id = str(job.get("id") or "")
    if not job_id:
        return
    log(log_path, f"Claimed retrain job {job_id}")
    report_progress(args, job_id, "running", "claimed", 2, "Задание получено Windows-agent")
    success = False
    message = ""
    summary: dict[str, Any] = {}
    try:
        run_retrain(args, job_id, job, repo_root, log_path)
        summary = read_json(runtime_dir / "torch_retrain_guard.json")
        report_progress(args, job_id, "running", "uploading", 72, "Обучение завершено, загружаю артефакты")
        for name in ARTIFACT_NAMES:
            path = runtime_dir / name
            if path.is_file():
                upload_artifact(args, job_id, path, log_path)
        success = True
        message = "training completed"
        log(log_path, f"Completed retrain job {job_id}")
    except Exception as exc:  # noqa: BLE001 - report failure to the bot.
        message = str(exc)
        log(log_path, f"Job {job_id} failed: {message}")
    finally:
        payload = {"success": success, "message": message, "summary": summary}
        api_json(args, f"/api/training/jobs/{job_id}/complete", payload)


def run_retrain(args: argparse.Namespace, job_id: str, job: dict[str, Any], repo_root: Path, log_path: Path) -> None:
    script = repo_root / "tools" / "run_torch_retrain.ps1"
    if not script.is_file():
        raise RuntimeError(f"retrain script not found: {script}")
    cmd = [
        "powershell.exe",
        "-NoProfile",
        "-ExecutionPolicy",
        "Bypass",
        "-File",
        str(script),
    ]
    parameters = job.get("parameters") if isinstance(job.get("parameters"), dict) else {}
    arg_map = {
        "symbols": "-Symbols",
        "limit": "-Limit",
        "lookbacks": "-Lookbacks",
        "architectures": "-Architectures",
        "hidden_sizes": "-HiddenSizes",
        "layers": "-Layers",
        "dropouts": "-Dropouts",
        "epochs": "-Epochs",
    }
    for key, ps_arg in arg_map.items():
        value = parameters.get(key)
        if value not in (None, ""):
            cmd.extend([ps_arg, str(value)])
    log(log_path, "Running retrain: " + " ".join(quote_for_log(part) for part in cmd))
    report_progress(args, job_id, "running", "training", 8, "PyTorch retrain запущен")
    line_count = 0
    output_queue: queue.Queue[str] = queue.Queue()

    def read_output() -> None:
        assert process.stdout is not None
        for raw_line in process.stdout:
            output_queue.put(raw_line.rstrip())

    with subprocess.Popen(
        cmd,
        cwd=str(repo_root),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding="utf-8",
        errors="replace",
    ) as process:
        reader = threading.Thread(target=read_output, name="training-output-reader", daemon=True)
        reader.start()
        last_report_at = 0.0
        last_message = "PyTorch retrain выполняется"
        while True:
            got_line = False
            try:
                message = output_queue.get(timeout=5)
                got_line = True
                log(log_path, message)
                line_count += 1
                if message:
                    last_message = message[-220:]
            except queue.Empty:
                pass

            progress = min(70, 8 + line_count // 3)
            now = time.monotonic()
            if got_line or now - last_report_at >= 30:
                safe_report_progress(args, job_id, "running", "training", progress, last_message, log_path)
                last_report_at = now

            if process.poll() is not None and output_queue.empty():
                break

        reader.join(timeout=2)
        code = process.wait()
    if code != 0:
        raise RuntimeError(f"retrain failed with exit code {code}")
    report_progress(args, job_id, "running", "guard", 70, "Guard завершён, подготавливаю артефакты")


def upload_artifact(args: argparse.Namespace, job_id: str, path: Path, log_path: Path) -> None:
    digest = hashlib.sha256(path.read_bytes()).hexdigest()
    size = path.stat().st_size
    chunk_size = max(64 * 1024, args.chunk_size)
    total = max(1, (size + chunk_size - 1) // chunk_size)
    log(log_path, f"Uploading {path.name}: {size} bytes, {total} chunks")
    with path.open("rb") as source:
        for index in range(total):
            data = source.read(chunk_size)
            payload = {
                "name": path.name,
                "index": index,
                "total": total,
                "sha256": digest,
                "data_base64": base64.b64encode(data).decode("ascii"),
            }
            api_json(args, f"/api/training/jobs/{job_id}/artifacts/chunk", payload, timeout=120)
            if index == 0 or index == total - 1 or index % 10 == 0:
                progress = 72 + int(((index + 1) / total) * 23)
                report_progress(args, job_id, "running", "uploading", progress, f"Загружаю {path.name}: {index + 1}/{total}")


def report_progress(
    args: argparse.Namespace,
    job_id: str,
    status: str,
    phase: str,
    progress_percent: int,
    message: str,
) -> None:
    api_json(
        args,
        f"/api/training/jobs/{job_id}/progress",
        {
            "status": status,
            "phase": phase,
            "progress_percent": progress_percent,
            "message": message,
            "worker": worker_payload(args, Path(args.repo_root).resolve()),
        },
    )


def safe_report_progress(
    args: argparse.Namespace,
    job_id: str,
    status: str,
    phase: str,
    progress_percent: int,
    message: str,
    log_path: Path,
) -> None:
    try:
        report_progress(args, job_id, status, phase, progress_percent, message)
    except Exception as exc:  # noqa: BLE001 - keep the local training process alive.
        log(log_path, f"Progress report failed: {exc}")


def api_json(args: argparse.Namespace, path: str, payload: dict[str, Any], timeout: int = 30) -> dict[str, Any]:
    url = args.api_base_url.rstrip("/") + path
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    headers = {"Content-Type": "application/json", "Accept": "application/json"}
    token = args.api_auth or os.environ.get("TRADEBOT_API_AUTH", "")
    headers.update(auth_headers(token))
    request = Request(url, data=body, headers=headers, method="POST")
    try:
        with urlopen(request, timeout=timeout) as response:
            text = response.read().decode("utf-8")
    except HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {exc.code} {path}: {detail[:300]}") from exc
    except URLError as exc:
        raise RuntimeError(f"network error {path}: {exc.reason}") from exc
    return json.loads(text) if text.strip() else {}


def auth_headers(token: str) -> dict[str, str]:
    value = token.strip()
    if not value:
        return {}
    headers = {"X-TradeBot-Token": value}
    if value.lower().startswith(("basic ", "bearer ")):
        headers["Authorization"] = value
    elif ":" in value:
        encoded = base64.b64encode(value.encode("utf-8")).decode("ascii")
        headers["Authorization"] = f"Basic {encoded}"
    else:
        headers["Authorization"] = f"Bearer {value}"
    return headers


def worker_payload(args: argparse.Namespace, repo_root: Path) -> dict[str, Any]:
    name = args.worker_name or platform.node() or "Windows training host"
    return {
        "worker_id": args.worker_id or f"{name}:{repo_root}",
        "name": name,
        "path": str(repo_root),
        "version": "1",
    }


def log(path: Path, message: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now().astimezone().isoformat(timespec="seconds")
    line = f"[{stamp}] {message}"
    print(line, flush=True)
    with path.open("a", encoding="utf-8") as handle:
        handle.write(line + "\n")


def read_json(path: Path) -> dict[str, Any]:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return {}
    return data if isinstance(data, dict) else {}


def quote_for_log(value: str) -> str:
    return f'"{value}"' if " " in value else value


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Poll TradeBot for retrain jobs and execute them on Windows.")
    parser.add_argument("--api-base-url", default=os.environ.get("TRADEBOT_API_BASE_URL", "https://tb.kusoft.xyz"))
    parser.add_argument("--api-auth", default=os.environ.get("TRADEBOT_API_AUTH", ""))
    parser.add_argument("--repo-root", default=str(Path(__file__).resolve().parents[1]))
    parser.add_argument("--worker-id", default=os.environ.get("TRADEBOT_TRAINING_WORKER_ID", ""))
    parser.add_argument("--worker-name", default=os.environ.get("TRADEBOT_TRAINING_WORKER_NAME", ""))
    parser.add_argument("--poll-seconds", type=int, default=int(os.environ.get("TRADEBOT_TRAINING_POLL_SECONDS", "60")))
    parser.add_argument("--chunk-size", type=int, default=int(os.environ.get("TRADEBOT_TRAINING_CHUNK_SIZE", str(512 * 1024))))
    parser.add_argument("--log-file", default=os.environ.get("TRADEBOT_TRAINING_LOG", ""))
    parser.add_argument("--once", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    main()