from __future__ import annotations

import shutil
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from .client import ERPClient
from .config import AgentConfig
from .db import LocalDB
from .storage import resolve_under_root, safe_relative_path, sha256_file, sanitize_segment, write_stream_to_file
from .storage_identity import ensure_storage_identity
from . import __version__


class StorageAgent:
    def __init__(self, config: AgentConfig, client: ERPClient, db: LocalDB, logger):
        self.config = config
        self.client = client
        self.db = db
        self.logger = logger
        self.last_heartbeat = 0.0
        self.identity_file = ensure_storage_identity(config)
        self.logger.info("Storage root identity verified: %s", self.identity_file)

    def run_forever(self) -> None:
        self.logger.info("Starting ERP Local Agent for node=%s", self.config.node_code)
        while True:
            self.run_once()
            time.sleep(max(5, self.config.poll_interval_seconds))

    def run_once(self) -> None:
        self._maybe_heartbeat()
        self.process_storage_jobs()
        self.process_download_requests()

    def _maybe_heartbeat(self, force: bool = False) -> None:
        now = time.time()
        if not force and now - self.last_heartbeat < self.config.heartbeat_interval_seconds:
            return
        try:
            total, used, free = shutil.disk_usage(self.config.storage_root)
            payload = {
                "status": "ONLINE",
                "storage_root": str(self.config.storage_root),
                "total_bytes": total,
                "used_bytes": used,
                "free_bytes": free,
                "agent_version": __version__,
                "agent_name": "ERP Local Agent",
                "capabilities": ["storage"],
                "agent_time_utc": datetime.now(timezone.utc).isoformat(),
            }
            self.client.heartbeat(payload)
            self.last_heartbeat = now
            self.logger.info("Heartbeat sent. free_gb=%.2f", free / (1024 ** 3))
        except Exception as exc:
            self.logger.exception("Heartbeat failed: %s", exc)

    def process_storage_jobs(self) -> None:
        try:
            jobs = self.client.pending_storage_jobs()
        except Exception as exc:
            self.logger.exception("Could not fetch pending storage jobs: %s", exc)
            return

        self.process_storage_jobs_from_payload(jobs)

    def process_storage_jobs_from_payload(self, jobs: list[dict[str, Any]]) -> None:
        if not jobs:
            return
        self.logger.info("Found %s pending storage job(s)", len(jobs))
        for job in jobs:
            try:
                self._process_storage_job(job)
            except Exception as exc:
                self.logger.exception("Storage job failed job=%s error=%s", job.get("id") or job.get("job_id"), exc)

    def _process_storage_job(self, job: dict[str, Any]) -> None:
        job_id = job.get("id") or job.get("job_id")
        if job_id is None:
            raise RuntimeError(f"Storage job missing id: {job}")

        original_filename = job.get("original_filename") or job.get("filename") or f"job_{job_id}.bin"
        relative_path_raw = job.get("local_relative_path") or job.get("relative_path")
        relative_path = safe_relative_path(relative_path_raw, original_filename)
        destination = resolve_under_root(self.config.storage_root, relative_path)

        self.logger.info("Downloading storage job=%s to %s", job_id, destination)
        response = self.client.download_storage_job(job_id)
        file_size, digest = write_stream_to_file(response.raw, destination)

        payload = {
            "status": "COMPLETED",
            "sha256_hash": digest,
            "file_hash": digest,
            "file_size": file_size,
            "local_relative_path": str(relative_path).replace("\\", "/"),
            "stored_at_utc": datetime.now(timezone.utc).isoformat(),
        }
        self.client.acknowledge_storage_job(job_id, payload)
        self.db.record_storage_job(str(job_id), str(relative_path).replace("\\", "/"), digest, file_size)
        self.logger.info("Storage job completed job=%s size=%s sha256=%s", job_id, file_size, digest)

    def process_download_requests(self) -> None:
        try:
            requests_ = self.client.pending_download_requests()
        except Exception as exc:
            self.logger.exception("Could not fetch pending download requests: %s", exc)
            return

        self.process_download_requests_from_payload(requests_)

    def process_download_requests_from_payload(self, requests_: list[dict[str, Any]]) -> None:
        if not requests_:
            return
        self.logger.info("Found %s pending download request(s)", len(requests_))
        for request in requests_:
            try:
                self._process_download_request(request)
            except Exception as exc:
                self.logger.exception("Download request failed request=%s error=%s", request.get("id") or request.get("request_id"), exc)

    def _process_download_request(self, request: dict[str, Any]) -> None:
        request_id = request.get("id") or request.get("request_id")
        if request_id is None:
            raise RuntimeError(f"Download request missing id: {request}")

        relative_path_raw = request.get("local_relative_path") or request.get("relative_path")
        filename = request.get("original_filename") or request.get("filename") or f"download_request_{request_id}.bin"
        relative_path = safe_relative_path(relative_path_raw, filename)
        local_file = resolve_under_root(self.config.storage_root, relative_path)
        if not local_file.exists() or not local_file.is_file():
            raise FileNotFoundError(f"Requested file is not found locally: {local_file}")

        file_size = local_file.stat().st_size
        digest = sha256_file(local_file)
        expected_hash = request.get("sha256_hash") or request.get("file_hash")
        if expected_hash and str(expected_hash).lower() != digest.lower():
            raise RuntimeError(f"SHA256 mismatch before upload. expected={expected_hash} actual={digest}")

        self.logger.info("Uploading file for download request=%s path=%s", request_id, local_file)
        extra = {
            "sha256_hash": digest,
            "file_hash": digest,
            "file_size": file_size,
            "local_relative_path": str(relative_path).replace("\\", "/"),
        }
        self.client.upload_download_request_file(request_id, local_file, extra)
        self.db.record_download_request(str(request_id), str(relative_path).replace("\\", "/"), digest, file_size)
        self.logger.info("Download request fulfilled request=%s size=%s sha256=%s", request_id, file_size, digest)
