from __future__ import annotations

import argparse
import base64
import hashlib
import json
import os
import platform
import queue
import re
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

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from crypto_spot_bot.orderbook_features import load_orderbook_feature_map
from tools.sync_orderbook_observations import sync_orderbook_observations


ARTIFACT_NAMES = (
    "lstm_forecaster.json",
    "torch_retrain_guard.json",
    "torch_threshold_calibration.json",
)
SHADOW_ARTIFACT_NAMES = (
    "lstm_forecaster.shadow.json",
    "torch_shadow_guard.json",
    "torch_shadow_calibration.json",
)
_LAST_ORDERBOOK_AUTO_CHECK = 0.0


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"):
        maybe_auto_queue_orderbook(args, repo_root, runtime_dir, log_path)
        return
    job = claim.get("job") if isinstance(claim.get("job"), dict) else {}
    job_id = str(job.get("id") or "")
    lease_token = str(claim.get("lease_token") or "")
    if not job_id:
        return
    if not lease_token:
        raise RuntimeError("training server did not issue a job lease")
    log(log_path, f"Claimed retrain job {job_id}")
    report_progress(args, job_id, lease_token, "running", "claimed", 2, "Задание получено Windows-agent")
    success = False
    message = ""
    summary: dict[str, Any] = {}
    try:
        parameters = job.get("parameters") if isinstance(job.get("parameters"), dict) else {}
        use_orderbook = parameters.get("use_orderbook", True) is not False
        orderbook_status: dict[str, Any] = {}
        if use_orderbook:
            report_progress(
                args,
                job_id,
                lease_token,
                "running",
                "orderbook_sync",
                4,
                "Синхронизирую forward-наблюдения стакана",
            )
            orderbook_status = prepare_orderbook_data(args, repo_root, parameters, log_path)
            if orderbook_status["state"] != "ready":
                summary = orderbook_status
                message = "forward orderbook coverage is still accumulating"
                success = True
                log(log_path, f"Job {job_id} remains in collecting state: {orderbook_status}")
                return
        run_retrain(
            args,
            job_id,
            lease_token,
            job,
            repo_root,
            log_path,
            orderbook_db=(repo_root / "runtime" / "orderbook_observations.sqlite3") if use_orderbook else None,
        )
        summary = read_json(runtime_dir / "torch_retrain_guard.json")
        accepted = summary.get("accepted") is True
        if accepted:
            report_progress(
                args,
                job_id,
                lease_token,
                "running",
                "uploading",
                72,
                "Обучение завершено, загружаю артефакты",
            )
            artifact_names = SHADOW_ARTIFACT_NAMES if use_orderbook else ARTIFACT_NAMES
            if use_orderbook:
                summary["deployment"] = "shadow"
                summary["orderbook"] = orderbook_status
            for name in artifact_names:
                path = runtime_dir / name
                if path.is_file():
                    upload_artifact(args, job_id, lease_token, path, log_path)
            message = (
                "training completed; candidate staged in shadow"
                if use_orderbook
                else "training completed; candidate accepted"
            )
            log(log_path, f"Completed retrain job {job_id}; {message}")
        else:
            reason = str(summary.get("reason") or "validation failed")
            message = f"training completed; candidate rejected by quality gate: {reason}"
            log(log_path, f"Completed retrain job {job_id}; candidate rejected: {reason}")
        success = True
    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,
            "lease_token": lease_token,
        }
        api_json(args, f"/api/training/jobs/{job_id}/complete", payload)


def run_retrain(
    args: argparse.Namespace,
    job_id: str,
    lease_token: str,
    job: dict[str, Any],
    repo_root: Path,
    log_path: Path,
    orderbook_db: Path | None = None,
) -> 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",
        "validation_window": "-ValidationWindow",
        "holdout_window": "-HoldoutWindow",
        "ensemble_seeds": "-EnsembleSeeds",
        "selection_folds": "-SelectionFolds",
        "learning_rate": "-LearningRate",
        "weight_decay": "-WeightDecay",
        "horizon": "-Horizon",
        "horizons": "-Horizons",
        "patience": "-Patience",
        "context_symbols": "-ContextSymbols",
        "features": "-Features",
        "seed": "-Seed",
        "interval": "-Interval",
    }
    for key, ps_arg in arg_map.items():
        value = parameters.get(key)
        if value not in (None, ""):
            cmd.extend([ps_arg, str(value)])
    if parameters.get("pooled", True) is True:
        cmd.append("-Pooled")
    if parameters.get("resume_candidate") is True:
        cmd.append("-ResumeCandidate")
    if orderbook_db is not None:
        cmd.extend(["-OrderbookDb", str(orderbook_db)])
        for key, ps_arg, default in (
            ("orderbook_min_samples_per_bucket", "-OrderbookMinSamplesPerBucket", 20),
            ("orderbook_min_covered_buckets", "-OrderbookMinCoveredBuckets", 240),
            ("orderbook_min_symbols", "-OrderbookMinSymbols", 2),
        ):
            cmd.extend([ps_arg, str(int(parameters.get(key, default) or default))])
    log(log_path, "Running retrain: " + " ".join(quote_for_log(part) for part in cmd))
    report_progress(
        args,
        job_id,
        lease_token,
        "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",
        **hidden_subprocess_kwargs(),
    ) as process:
        reader = threading.Thread(target=read_output, name="training-output-reader", daemon=True)
        reader.start()
        last_report_at = 0.0
        started_at = time.monotonic()
        last_output_at = started_at
        last_message = "PyTorch retrain выполняется"
        while True:
            got_line = False
            try:
                message = output_queue.get(timeout=5)
                got_line = True
                last_output_at = time.monotonic()
                log(log_path, message)
                line_count += 1
                if message:
                    last_message = friendly_training_message(message)
            except queue.Empty:
                pass

            progress = min(70, 8 + line_count // 3)
            now = time.monotonic()
            if got_line or now - last_report_at >= 30:
                report_message = last_message
                if not got_line:
                    report_message = training_heartbeat_message(now, started_at, last_output_at, last_message)
                safe_report_progress(
                    args,
                    job_id,
                    lease_token,
                    "running",
                    "training",
                    progress,
                    report_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,
        lease_token,
        "running",
        "guard",
        70,
        "Guard завершён, подготавливаю артефакты",
    )


def prepare_orderbook_data(
    args: argparse.Namespace,
    repo_root: Path,
    parameters: dict[str, Any],
    log_path: Path,
) -> dict[str, Any]:
    database_path = repo_root / "runtime" / "orderbook_observations.sqlite3"
    token = args.api_auth or os.environ.get("TRADEBOT_API_AUTH", "")
    sync_result = sync_orderbook_observations(
        api_base_url=args.api_base_url,
        token=token,
        database_path=database_path,
    )
    interval = str(parameters.get("interval") or os.environ.get("TORCH_RETRAIN_INTERVAL") or "60")
    minimum_samples = int(parameters.get("orderbook_min_samples_per_bucket", 20) or 20)
    minimum_buckets = int(parameters.get("orderbook_min_covered_buckets", 240) or 240)
    minimum_symbols = int(parameters.get("orderbook_min_symbols", 2) or 2)
    requested_symbols = {
        item.strip().upper()
        for item in str(parameters.get("symbols") or "").split(",")
        if item.strip()
    }
    _features, manifest = load_orderbook_feature_map(
        database_path,
        interval=interval,
        symbols=sorted(requested_symbols) if requested_symbols else None,
        min_samples_per_bucket=minimum_samples,
    )
    eligible = sorted(
        symbol
        for symbol, row in manifest.items()
        if int(row.get("covered_buckets", 0) or 0) >= minimum_buckets
    )
    state = "ready" if len(eligible) >= minimum_symbols else "collecting_orderbook"
    coverage = {
        symbol: int(row.get("covered_buckets", 0) or 0)
        for symbol, row in sorted(manifest.items())
    }
    result = {
        "accepted": False,
        "state": state,
        "reason": (
            "orderbook coverage ready for training"
            if state == "ready"
            else "forward orderbook coverage is below the configured minimum"
        ),
        "eligible_symbols": eligible,
        "eligible_symbol_count": len(eligible),
        "minimum_symbols": minimum_symbols,
        "minimum_covered_buckets": minimum_buckets,
        "minimum_samples_per_bucket": minimum_samples,
        "covered_buckets_by_symbol": coverage,
        "local_samples": int(sync_result.get("local_samples", 0) or 0),
        "downloaded_samples": int(sync_result.get("downloaded", 0) or 0),
    }
    log(log_path, "Orderbook preparation: " + json.dumps(result, ensure_ascii=False, sort_keys=True))
    return result


def maybe_auto_queue_orderbook(
    args: argparse.Namespace,
    repo_root: Path,
    runtime_dir: Path,
    log_path: Path,
) -> None:
    global _LAST_ORDERBOOK_AUTO_CHECK
    try:
        interval_seconds = max(
            300,
            int(os.environ.get("TORCH_ORDERBOOK_AUTO_CHECK_SECONDS", "3600") or 3600),
        )
    except ValueError:
        interval_seconds = 3600
    now = time.monotonic()
    if _LAST_ORDERBOOK_AUTO_CHECK and now - _LAST_ORDERBOOK_AUTO_CHECK < interval_seconds:
        return
    _LAST_ORDERBOOK_AUTO_CHECK = now
    marker_path = runtime_dir / "orderbook_auto_queue.json"
    if marker_path.is_file() or (runtime_dir / "lstm_forecaster.shadow.json").is_file():
        return
    status = prepare_orderbook_data(args, repo_root, {}, log_path)
    if status.get("state") != "ready":
        return
    response = api_json(args, "/api/training/retrain/auto", {})
    if not response.get("queued"):
        log(log_path, f"Automatic orderbook retrain was not queued: {response.get('reason', 'unknown')}")
        return
    marker = {
        "queued_at": datetime.now().astimezone().isoformat(timespec="seconds"),
        "job_id": (response.get("job") or {}).get("id"),
        "coverage": status,
    }
    marker_tmp = marker_path.with_suffix(".tmp")
    marker_tmp.write_text(json.dumps(marker, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    marker_tmp.replace(marker_path)
    log(log_path, f"Automatically queued orderbook retrain job {marker['job_id']}")


def friendly_training_message(message: str) -> str:
    cleaned = message.strip()
    if not cleaned:
        return "PyTorch обучает модель"

    if "Starting PyTorch recurrent retrain:" in cleaned:
        return "PyTorch LSTM/GRU запущен: готовлю данные и варианты модели"

    started = re.search(
        r"training started: symbols=(?P<symbols>\d+) interval=(?P<interval>\d+) "
        r"limit=(?P<limit>\d+) epochs=(?P<epochs>\d+)",
        cleaned,
    )
    if started:
        interval = started.group("interval")
        timeframe = "1h" if interval == "60" else f"{interval}m"
        return (
            f"Старт обучения: {started.group('symbols')} пар, таймфрейм {timeframe}, "
            f"история {started.group('limit')} свечей, до {started.group('epochs')} эпох"
        )

    pair_started = re.search(r"^(?P<symbol>[A-Z0-9]+): training started \((?P<index>\d+)/(?P<total>\d+)\)", cleaned)
    if pair_started:
        return (
            f"{pair_started.group('symbol')}: обучение пары "
            f"{pair_started.group('index')}/{pair_started.group('total')}"
        )

    preparing = re.search(r"^(?P<symbol>[A-Z0-9]+): preparing lookback=(?P<lookback>\d+)", cleaned)
    if preparing:
        return f"{preparing.group('symbol')}: готовлю окно {preparing.group('lookback')} свечей"

    fitting = re.search(
        r"^(?P<symbol>[A-Z0-9]+): fitting (?P<arch>lstm|gru) "
        r"lookback=(?P<lookback>\d+) hidden=(?P<hidden>\d+) "
        r"layers=(?P<layers>\d+) dropout=(?P<dropout>[0-9.]+)",
        cleaned,
    )
    if fitting:
        return (
            f"{fitting.group('symbol')}: обучаю {fitting.group('arch').upper()}, "
            f"окно {fitting.group('lookback')}, нейронов {fitting.group('hidden')}, "
            f"слоёв {fitting.group('layers')}, dropout {fitting.group('dropout')}"
        )

    model = re.search(
        r"^(?P<symbol>[A-Z0-9]+): model=torch_(?P<arch>lstm|gru).*?"
        r"mae=(?P<mae>[0-9.]+)%.*?skill=(?P<skill>-?[0-9.]+).*?dir=(?P<direction>[0-9.]+)",
        cleaned,
    )
    if model:
        direction = float(model.group("direction")) * 100
        skill = float(model.group("skill")) * 100
        return (
            f"{model.group('symbol')}: выбран {model.group('arch').upper()}, "
            f"ошибка {model.group('mae')}%, skill {skill:.1f}%, направление {direction:.1f}%"
        )

    if "Calibrating current artifact" in cleaned:
        return "Проверяю текущую модель на replay"
    if "Calibrating candidate artifact" in cleaned:
        return "Проверяю новую модель на replay"
    if "Running retrain guard" in cleaned:
        return "Gate сравнивает новую модель с текущей"
    if "Candidate rejected by guard" in cleaned:
        return "Новая модель обучилась, но gate не дал ей ходу"
    if "Candidate accepted by guard" in cleaned:
        return "Новая модель прошла gate и стала активной"

    return cleaned[-220:]


def training_heartbeat_message(now: float, started_at: float, last_output_at: float, last_message: str) -> str:
    elapsed = format_duration(now - started_at)
    idle_seconds = max(0.0, now - last_output_at)
    if idle_seconds >= 45:
        return (
            f"PyTorch обучает модель: процесс активен {elapsed}; "
            f"последний лог {format_duration(idle_seconds)} назад: {last_message[:140]}"
        )
    return last_message or f"PyTorch обучает модель: процесс активен {elapsed}"


def format_duration(seconds: float) -> str:
    total_seconds = max(0, int(seconds))
    minutes, seconds_part = divmod(total_seconds, 60)
    hours, minutes_part = divmod(minutes, 60)
    if hours:
        return f"{hours}ч {minutes_part}м"
    if minutes:
        return f"{minutes}м {seconds_part}с"
    return f"{seconds_part}с"


def upload_artifact(
    args: argparse.Namespace,
    job_id: str,
    lease_token: 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"),
                "lease_token": lease_token,
            }
            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,
                    lease_token,
                    "running",
                    "uploading",
                    progress,
                    f"Загружаю {path.name}: {index + 1}/{total}",
                )


def report_progress(
    args: argparse.Namespace,
    job_id: str,
    lease_token: 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()),
            "lease_token": lease_token,
        },
    )


def safe_report_progress(
    args: argparse.Namespace,
    job_id: str,
    lease_token: str,
    status: str,
    phase: str,
    progress_percent: int,
    message: str,
    log_path: Path,
) -> None:
    last_error: Exception | None = None
    for attempt in range(1, 4):
        try:
            report_progress(
                args,
                job_id,
                lease_token,
                status,
                phase,
                progress_percent,
                message,
            )
            return
        except Exception as exc:  # noqa: BLE001 - keep the local training process alive.
            last_error = exc
            if attempt < 3:
                time.sleep(attempt * 2)
    log(log_path, f"Temporary progress upload error; training continues: {last_error}")


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": "3",
    }


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}"
    if sys.stdout is not None:
        try:
            print(line, flush=True)
        except OSError:
            pass
    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 hidden_subprocess_kwargs() -> dict[str, Any]:
    if os.name != "nt":
        return {}
    startupinfo = subprocess.STARTUPINFO()
    startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
    startupinfo.wShowWindow = 0
    return {
        "creationflags": getattr(subprocess, "CREATE_NO_WINDOW", 0),
        "startupinfo": startupinfo,
    }


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()