"""Command-line interface for chunked screenshot batch runs."""

from __future__ import annotations

import argparse
import json
import logging
import sys
from collections.abc import Callable, Iterable, Sequence
from pathlib import Path
from typing import Any

from screenshot import ScreenshotOptions

from .coordinator import (
    ScreenshotJobSpec,
    ScreenshotMultiRunCoordinator,
    ScreenshotRunSummary,
)
from .utils import ensure_capture_enabled, normalize_options_payload


def _parse_args(argv: Sequence[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Submit screenshot jobs in chunks via the Azure service")
    parser.add_argument("input", type=Path, help="Path to JSONL file with screenshot jobs")
    parser.add_argument("batch_id", help="Batch identifier to use when calling the service")
    parser.add_argument(
        "--base-url",
        required=False,
        help="Screenshot Azure Function base URL (defaults to SCREENSHOT_FUNC_BASE_URL)",
    )
    parser.add_argument(
        "--store-dir",
        type=Path,
        help="Directory where run artifacts should be downloaded (optional)",
    )
    parser.add_argument(
        "--state-path",
        type=Path,
        help="Optional JSON state file for tracking in-flight runs (enables resuming)",
    )
    parser.add_argument("--chunk-size", type=int, default=4, help="Jobs per API call (default: 4)")
    parser.add_argument(
        "--max-parallel-runs",
        type=int,
        default=32,
        help="Number of concurrent API calls to run (default: 32)",
    )
    parser.add_argument(
        "--per-run-concurrency",
        type=int,
        default=4,
        help="Concurrency value to send with each API payload (default: 4)",
    )
    parser.add_argument(
        "--download-results",
        action="store_true",
        help="Download run artifacts under --store-dir",
    )
    parser.add_argument(
        "--no-download",
        dest="download_results",
        action="store_false",
        help="Skip artifact download",
    )
    parser.set_defaults(download_results=False)
    parser.add_argument(
        "--upload-results",
        action="store_true",
        help="Upload downloaded artifacts via infra_core.azure.storage (requires config)",
    )
    parser.add_argument(
        "--metadata",
        type=str,
        help="JSON object merged into metadata for every job",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="Optional path to write job results (JSONL with status/blob_url)",
    )
    parser.add_argument(
        "--sample",
        type=int,
        metavar="N",
        help="Limit the run to the first N jobs (for smoke testing)",
    )
    parser.add_argument(
        "--log-level",
        default="INFO",
        choices=["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"],
        help="Log verbosity (default: INFO)",
    )
    return parser.parse_args(argv)


def _iter_jsonl(path: Path) -> Iterable[dict[str, Any]]:
    with path.open("r", encoding="utf-8") as handle:
        for line_no, line in enumerate(handle, start=1):
            text = line.strip()
            if not text:
                continue
            try:
                yield json.loads(text)
            except json.JSONDecodeError as exc:
                raise RuntimeError(f"Failed to parse {path} line {line_no}: {exc}") from exc


def _build_spec(record: dict[str, Any], default_metadata: dict[str, Any]) -> ScreenshotJobSpec:
    url = (record.get("url") or record.get("website") or "").strip()
    if not url:
        raise ValueError("Record missing 'url'")

    options_payload = record.get("options")
    if isinstance(options_payload, dict):
        normalized_payload = normalize_options_payload(options_payload)
        try:
            options = ScreenshotOptions.from_dict(normalized_payload)
        except ValueError as exc:
            raise ValueError("Job options must provide nested 'capture', 'browser', or 'runner' sections.") from exc
    elif options_payload is None:
        options = ScreenshotOptions()
    else:
        raise ValueError("options must be a JSON object when provided")
    options = ensure_capture_enabled(options)

    metadata = dict(default_metadata)
    if isinstance(record.get("metadata"), dict):
        metadata.update(record["metadata"])

    snapshot_path: Path | None = None
    snapshot = record.get("html_snapshot_path")
    if isinstance(snapshot, str) and snapshot.strip():
        snapshot_path = Path(snapshot)

    return ScreenshotJobSpec.from_url(
        url,
        job_id=record.get("job_id") or record.get("site_id"),
        options=options,
        partition_date=record.get("partition_date"),
        metadata=metadata,
        html_snapshot_path=snapshot_path,
    )


def _make_progress_callback(
    total_chunks: int,
) -> Callable[[int, int, ScreenshotRunSummary], None] | None:
    if total_chunks <= 0:
        return None

    bar_width = 20

    def _callback(completed: int, total: int, summary: ScreenshotRunSummary) -> None:
        filled = int(bar_width * completed / max(1, total))
        bar = "#" * filled + "-" * (bar_width - filled)
        percent = (completed / max(1, total)) * 100.0
        logging.info(
            "Progress [%s] %.1f%% (%d/%d) chunk=%d run=%s status=%s failed=%d",
            bar,
            percent,
            completed,
            total,
            summary.chunk_index,
            summary.result.job_id,
            summary.result.status,
            summary.result.job_failed,
        )

    return _callback


def main(argv: Sequence[str] | None = None) -> int:
    args = _parse_args(list(argv) if argv is not None else sys.argv[1:])

    logging.basicConfig(
        level=getattr(logging, args.log_level.upper()),
        format="%(asctime)s %(levelname)s %(message)s",
    )

    base_url = (args.base_url or "").strip()
    if not base_url:
        from os import getenv

        base_url = getenv("SCREENSHOT_FUNC_BASE_URL", "").strip()
    if not base_url:
        logging.error("Screenshot base URL missing (set --base-url or SCREENSHOT_FUNC_BASE_URL)")
        return 1

    if not args.input.exists():
        logging.error("Input file %s not found", args.input)
        return 1

    if args.download_results and not args.store_dir:
        logging.error("--store-dir is required when --download-results is enabled")
        return 1

    if args.upload_results and not args.download_results:
        logging.error("--upload-results requires --download-results")
        return 1

    metadata: dict[str, Any] = {}
    if args.metadata:
        try:
            parsed = json.loads(args.metadata)
        except json.JSONDecodeError as exc:
            logging.error("--metadata must be valid JSON: %s", exc)
            return 1
        if not isinstance(parsed, dict):
            logging.error("--metadata must be a JSON object")
            return 1
        metadata = dict(parsed)

    specs: list[ScreenshotJobSpec] = []
    skipped = 0
    records_data: list[dict[str, Any]] = []
    try:
        for record in _iter_jsonl(args.input):
            entry: dict[str, Any] = {"record": dict(record)}
            try:
                spec = _build_spec(record, metadata)
            except ValueError as exc:
                skipped += 1
                entry["status"] = "skipped_missing_url"
                entry["error"] = str(exc)
                records_data.append(entry)
                logging.warning("Skipping record: %s", exc)
                continue
            specs.append(spec)
            entry["spec"] = spec
            entry["job_id"] = spec.job_id
            records_data.append(entry)
            if args.sample and len(specs) >= int(args.sample):
                break
    except Exception as exc:  # pragma: no cover - user input errors
        logging.error("Failed to load input file: %s", exc)
        return 1

    if not specs:
        logging.error("No valid jobs found in %s", args.input)
        return 1

    if skipped:
        logging.info("Skipped %d record(s) without URL/website", skipped)

    runner = ScreenshotMultiRunCoordinator(
        base_url=base_url,
        batch_id=args.batch_id,
        store_dir=args.store_dir,
        per_run_concurrency=args.per_run_concurrency,
        download_results=args.download_results,
        upload_results=args.upload_results,
        state_path=args.state_path,
        default_metadata=metadata,
    )

    logging.info(
        "Submitting %d jobs in chunks of %d (parallel runs=%d)",
        len(specs),
        args.chunk_size,
        args.max_parallel_runs,
    )

    total_chunks = (len(specs) + args.chunk_size - 1) // args.chunk_size
    progress_callback = _make_progress_callback(total_chunks)

    summaries = runner.run_jobs(
        specs,
        chunk_size=args.chunk_size,
        max_parallel_runs=args.max_parallel_runs,
        progress_callback=progress_callback,
    )

    total_jobs = sum(summary.result.job_total for summary in summaries)
    total_failed = sum(summary.result.job_failed for summary in summaries)

    logging.info(
        "Runs completed: %d, jobs=%d failed=%d",
        len(summaries),
        total_jobs,
        total_failed,
    )

    for summary in summaries:
        logging.info(
            "Chunk %d → run=%s status=%s completed=%d failed=%d",
            summary.chunk_index,
            summary.result.job_id,
            summary.result.status,
            summary.result.job_completed,
            summary.result.job_failed,
        )

    output_path = args.output or args.input.with_suffix(".results.jsonl")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    status_map: dict[str, str] = {}
    blob_map: dict[str, str | None] = {}
    error_map: dict[str, Sequence[str]] = {}

    for summary in summaries:
        for job_id, job_summary in summary.job_results.items():
            status_map[job_id] = job_summary.status
            if job_summary.blob_url:
                blob_map[job_id] = job_summary.blob_url
            if job_summary.errors:
                error_map[job_id] = list(job_summary.errors)

    with output_path.open("w", encoding="utf-8") as handle:
        for entry in records_data:
            record = dict(entry["record"])
            spec_entry = entry.get("spec")
            if isinstance(spec_entry, ScreenshotJobSpec):
                job_id = spec_entry.job_id
                status = status_map.get(job_id, "unknown")
                blob_url = blob_map.get(job_id)
                if status != "succeeded":
                    blob_url = None
                if job_id:
                    record.setdefault("job_id", job_id)
                record["status"] = status
                record["blob_url"] = blob_url
                errors = error_map.get(job_id)
                if errors:
                    record["errors"] = errors
            else:
                record["status"] = entry.get("status", "skipped")
                record["blob_url"] = None
                if entry.get("error"):
                    record["errors"] = [entry["error"]]
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")

    logging.info("Wrote results to %s", output_path)

    return 0 if total_failed == 0 else 1


if __name__ == "__main__":  # pragma: no cover
    sys.exit(main())
