"""Job synchronization: poll upstream jobs and report completions."""

from __future__ import annotations

import asyncio
import json
import logging
import uuid
from datetime import datetime, timezone

from radix_edge.config import get_config
from radix_edge.db import get_db
from radix_edge.upstream.client import UpstreamClient
from radix_edge.upstream.translate import cloud_job_to_edge_fields

logger = logging.getLogger("radix_edge.upstream")


async def job_sync_loop(client: UpstreamClient) -> None:
    """Poll for upstream jobs and report completions."""
    config = get_config()
    logger.info("Upstream job sync started (interval=%.0fs)", config.job_poll_interval)

    while True:
        try:
            await _poll_and_insert_jobs(client)
            await _report_completions(client)
        except Exception:
            logger.exception("Job sync error")

        await asyncio.sleep(config.job_poll_interval)


async def _poll_and_insert_jobs(client: UpstreamClient) -> int:
    """Poll control plane for jobs and insert into local DB. Returns count inserted."""
    try:
        jobs = await client.poll_jobs(max_jobs=1)
    except Exception:
        logger.exception("Failed to poll upstream jobs")
        return 0

    inserted = 0
    for cloud_job in jobs:
        upstream_job_id = cloud_job.get("job_id")
        if not upstream_job_id:
            continue

        with get_db() as db:
            existing = db.execute(
                "SELECT job_id FROM jobs WHERE upstream_job_id = ?",
                (upstream_job_id,),
            ).fetchone()
            if existing:
                logger.debug("Upstream job %s already exists locally", upstream_job_id)
                continue

        fields = cloud_job_to_edge_fields(cloud_job)
        local_job_id = f"job-{uuid.uuid4().hex[:12]}"

        with get_db() as db:
            db.execute(
                """INSERT INTO jobs (
                    job_id, user_id, name, image, command, args, env,
                    job_type, priority, num_gpus, gpu_mem_required_mb,
                    timeout_seconds, origin, upstream_job_id, metadata
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    local_job_id,
                    fields["user_id"],
                    fields["name"],
                    fields["image"],
                    json.dumps(fields["command"]),
                    json.dumps(fields["args"]),
                    json.dumps(fields["env"]),
                    fields["job_type"],
                    fields["priority"],
                    fields["num_gpus"],
                    fields["gpu_mem_required_mb"],
                    fields["timeout_seconds"],
                    fields["origin"],
                    fields["upstream_job_id"],
                    json.dumps(fields["metadata"]) if fields["metadata"] else None,
                ),
            )

        logger.info(
            "Inserted upstream job %s as local %s (kind=%s, gpus=%d)",
            upstream_job_id, local_job_id, fields["job_type"], fields["num_gpus"],
        )
        inserted += 1

    return inserted


async def _report_completions(client: UpstreamClient) -> int:
    """Find completed upstream jobs and report status back. Returns count reported."""
    with get_db() as db:
        rows = db.execute(
            """SELECT job_id, upstream_job_id, status, exit_code, error_message,
                      stdout_tail, stderr_tail, runtime_seconds, metadata
            FROM jobs
            WHERE origin = 'upstream'
              AND upstream_job_id IS NOT NULL
              AND status IN ('succeeded', 'failed')
            ORDER BY completed_at
            LIMIT 10""",
        ).fetchall()

    reported = 0
    for row in rows:
        # Check if already reported
        meta = json.loads(row["metadata"]) if row["metadata"] else {}
        if meta.get("reported"):
            continue

        upstream_job_id = row["upstream_job_id"]

        output_payload: dict = {}
        if row["stdout_tail"]:
            output_payload["stdout"] = row["stdout_tail"][-5000:]
        if row["stderr_tail"]:
            output_payload["stderr"] = row["stderr_tail"][-2000:]

        metrics: dict = {}
        if row["runtime_seconds"] is not None:
            metrics["runtime_seconds"] = round(row["runtime_seconds"], 2)
        if row["exit_code"] is not None:
            metrics["exit_code"] = row["exit_code"]

        try:
            await client.complete_job(
                job_id=upstream_job_id,
                status=row["status"],
                output_payload=output_payload,
                metrics=metrics,
                error_message=row["error_message"],
            )

            meta["reported"] = True
            meta["reported_at"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
            with get_db() as db:
                db.execute(
                    "UPDATE jobs SET metadata = ? WHERE job_id = ?",
                    (json.dumps(meta), row["job_id"]),
                )

            logger.info(
                "Reported completion for upstream job %s (status=%s, local=%s)",
                upstream_job_id, row["status"], row["job_id"],
            )
            reported += 1

        except Exception:
            logger.exception("Failed to report completion for upstream job %s", upstream_job_id)

    return reported
