from __future__ import annotations

import json
import os
import shutil
import threading
import time
import uuid
import zipfile
from collections.abc import Callable
from pathlib import Path
from typing import Any

from app.config import Settings
from app.crawler import (
    DEFAULT_ASSET_TYPES,
    CaptureCancelled,
    CaptureError,
    CaptureOptions,
    SiteCrawler,
)
from app.db import Database
from app.url_policy import NetworkRules, UrlPolicyError, resolve_and_validate


def _remove_tree(path: Path, allowed_root: Path) -> None:
    """Delete only a child path beneath the configured work root."""
    root = allowed_root.resolve()
    resolved = path.resolve()
    if resolved == root or not resolved.is_relative_to(root):
        raise RuntimeError("Refusing to remove a path outside SiteHarbor's work directory.")
    shutil.rmtree(resolved, ignore_errors=True)


def _create_archive(
    source_dir: Path,
    destination: Path,
    archive_root: str,
    cancelled: Callable[[], bool],
    deadline_exceeded: Callable[[], bool],
    heartbeat: Callable[[], None],
) -> int:
    partial = destination.with_suffix(".zip.partial")
    partial.unlink(missing_ok=True)
    last_heartbeat = time.monotonic()
    try:
        with zipfile.ZipFile(
            partial,
            mode="w",
            compression=zipfile.ZIP_DEFLATED,
            compresslevel=6,
            strict_timestamps=False,
        ) as archive:
            for source_path in sorted(source_dir.rglob("*")):
                if cancelled():
                    raise CaptureCancelled("Capture cancelled while packaging the archive")
                if deadline_exceeded():
                    raise CaptureError(
                        "The configured capture time limit was reached while packaging"
                    )
                if time.monotonic() - last_heartbeat >= 5:
                    heartbeat()
                    last_heartbeat = time.monotonic()
                if not source_path.is_file() or source_path.is_symlink():
                    continue
                relative = source_path.relative_to(source_dir).as_posix()
                with source_path.open("rb") as source_file:
                    with archive.open(f"{archive_root}/{relative}", mode="w") as archive_file:
                        while chunk := source_file.read(1_024 * 1_024):
                            if cancelled():
                                raise CaptureCancelled(
                                    "Capture cancelled while packaging the archive"
                                )
                            if deadline_exceeded():
                                raise CaptureError(
                                    "The configured capture time limit was reached while packaging"
                                )
                            if time.monotonic() - last_heartbeat >= 5:
                                heartbeat()
                                last_heartbeat = time.monotonic()
                            archive_file.write(chunk)
        os.replace(partial, destination)
        return destination.stat().st_size
    except Exception:
        partial.unlink(missing_ok=True)
        destination.unlink(missing_ok=True)
        raise


class CaptureService:
    """Runs one leased capture and persists a terminal SQLite job state."""

    def __init__(self, database: Database, settings: Settings, worker_id: str) -> None:
        self.database = database
        self.settings = settings
        self.worker_id = worker_id

    def run(self, job: dict[str, Any]) -> None:
        job_id = str(job["id"])
        config = dict(job["config"])
        attempt_id = uuid.uuid4().hex[:12]
        work_dir = self.settings.work_dir / f"{job_id}-{attempt_id}"
        artifact_path = self.settings.artifacts_dir / f"{job_id}-{attempt_id}.zip"
        report_path = self.settings.reports_dir / f"{job_id}-{attempt_id}.json"
        heartbeat_stop = threading.Event()
        lease_lost = threading.Event()
        heartbeat_thread: threading.Thread | None = None
        capture_deadline = time.monotonic() + int(config["max_duration_seconds"])

        def capture_cancelled() -> bool:
            return (
                lease_lost.is_set()
                or self.database.is_cancel_requested(job_id)
                or not self.database.owns_lease(job_id, self.worker_id)
            )

        def renew_lease() -> None:
            interval = min(20, max(3, self.settings.lease_seconds // 3))
            while not heartbeat_stop.wait(interval):
                try:
                    if not self.database.renew_lease(
                        job_id, self.worker_id, self.settings.lease_seconds
                    ):
                        lease_lost.set()
                        return
                except Exception:
                    # Ownership checks stop this attempt if the lease eventually expires.
                    continue

        try:
            _remove_tree_if_present(work_dir, self.settings.work_dir)
            work_dir.mkdir(parents=True, exist_ok=True)
            if capture_cancelled():
                raise CaptureCancelled("Capture cancelled before it started")

            heartbeat_thread = threading.Thread(
                target=renew_lease,
                name=f"siteharbor-lease-{job_id[-8:]}",
                daemon=True,
            )
            heartbeat_thread.start()

            rules = NetworkRules(
                allow_private_networks=self.settings.allow_private_networks,
                allow_nonstandard_ports=self.settings.allow_nonstandard_ports,
            )
            resolve_and_validate(str(job["normalized_url"]), rules)
            remaining_seconds = int(capture_deadline - time.monotonic())
            if remaining_seconds <= 0:
                raise CaptureError(
                    "The configured capture time limit was reached during validation"
                )
            options = CaptureOptions(
                source_url=str(job["normalized_url"]),
                include_external_assets=bool(config["include_external_assets"]),
                max_pages=int(config["max_pages"]),
                max_depth=int(config["max_depth"]),
                max_bytes=int(config["max_bytes"]),
                max_duration_seconds=remaining_seconds,
                parallel_connections=int(
                    config.get("parallel_connections", self.settings.fetch_concurrency)
                ),
                download_asset_types=tuple(
                    config.get("download_asset_types", DEFAULT_ASSET_TYPES)
                ),
                skip_file_types=tuple(config.get("skip_file_types", ())),
                use_proxy=bool(config.get("use_proxy", False)),
            )

            def update(
                phase: str, message: str, stats: dict[str, int], warnings: list[str]
            ) -> None:
                renewed = self.database.update_progress(
                    job_id,
                    self.worker_id,
                    self.settings.lease_seconds,
                    phase=phase,
                    message=message,
                    stats=stats,
                    warnings=warnings,
                )
                if not renewed:
                    raise CaptureCancelled("Capture lease was reassigned")

            if not self.database.update_progress(
                job_id,
                self.worker_id,
                self.settings.lease_seconds,
                phase="Validating",
                message="Checking capture policy and target reachability",
                stats=_job_stats(job),
                warnings=[],
            ):
                raise CaptureCancelled("Capture lease was reassigned")

            crawler = SiteCrawler(
                settings=self.settings,
                options=options,
                cancelled=capture_cancelled,
                on_progress=update,
            )
            result = crawler.run(work_dir)

            if capture_cancelled():
                raise CaptureCancelled("Capture cancelled before packaging")
            if time.monotonic() >= capture_deadline:
                raise CaptureError("The configured capture time limit was reached before packaging")

            update("Packaging", "Creating a portable ZIP archive", result.stats, result.warnings)
            archive_size = _create_archive(
                result.site_dir,
                artifact_path,
                archive_root=f"siteharbor-{job_id}",
                cancelled=capture_cancelled,
                deadline_exceeded=lambda: time.monotonic() >= capture_deadline,
                heartbeat=lambda: update(
                    "Packaging",
                    "Creating a portable ZIP archive",
                    result.stats,
                    result.warnings,
                ),
            )

            report = {
                **result.report,
                "job_id": job_id,
                "archive": {
                    "filename": artifact_path.name,
                    "size_bytes": archive_size,
                    "archive_root": f"siteharbor-{job_id}",
                },
            }
            if time.monotonic() >= capture_deadline:
                raise CaptureError(
                    "The configured capture time limit was reached before completion"
                )
            report_path.write_text(
                json.dumps(report, indent=2, ensure_ascii=True), encoding="utf-8"
            )
            marked_ready = self.database.mark_ready(
                job_id,
                self.worker_id,
                artifact_path=str(artifact_path),
                report_path=str(report_path),
                archive_size=archive_size,
                entry_point=result.entry_point,
                stats=result.stats,
                warnings=result.warnings,
                retention_days=int(config["retention_days"]),
            )
            if not marked_ready:
                artifact_path.unlink(missing_ok=True)
                report_path.unlink(missing_ok=True)
                if self.database.is_cancel_requested(job_id):
                    self.database.mark_cancelled(
                        job_id,
                        self.worker_id,
                        "Capture cancelled before the archive became available",
                    )
        except CaptureCancelled as error:
            artifact_path.unlink(missing_ok=True)
            report_path.unlink(missing_ok=True)
            if self.database.is_cancel_requested(job_id):
                self.database.mark_cancelled(job_id, self.worker_id, str(error))
        except UrlPolicyError as error:
            artifact_path.unlink(missing_ok=True)
            self.database.mark_failed(job_id, self.worker_id, "target_blocked", str(error))
        except CaptureError as error:
            artifact_path.unlink(missing_ok=True)
            self.database.mark_failed(job_id, self.worker_id, "capture_failed", str(error))
        except Exception as error:
            artifact_path.unlink(missing_ok=True)
            self.database.mark_failed(
                job_id,
                self.worker_id,
                "internal_error",
                f"Capture failed: {error}",
            )
        finally:
            heartbeat_stop.set()
            if heartbeat_thread:
                heartbeat_thread.join(timeout=1)
            try:
                _remove_tree_if_present(work_dir, self.settings.work_dir)
            except OSError:
                pass


def _job_stats(job: dict[str, Any]) -> dict[str, int]:
    return {
        "pages_found": int(job.get("pages_found", 0)),
        "pages_fetched": int(job.get("pages_fetched", 0)),
        "assets_found": int(job.get("assets_found", 0)),
        "assets_fetched": int(job.get("assets_fetched", 0)),
        "bytes_downloaded": int(job.get("bytes_downloaded", 0)),
        "files_written": int(job.get("files_written", 0)),
    }


def _remove_tree_if_present(path: Path, root: Path) -> None:
    if path.exists():
        _remove_tree(path, root)