ai-flow-python/app/service/template/template_service.py

import hashlib
import io
import logging
import os
import shutil
import time
import zipfile
from concurrent.futures import ThreadPoolExecutor
from typing import Optional

import httpx

from app.config import settings

logger = logging.getLogger(__name__)


class _TemplateFileInfo:
    __slots__ = ("size", "etag", "last_modified")

    def __init__(self, size: int, etag: str | None, last_modified: str | None):
        self.size = size
        self.etag = etag
        self.last_modified = last_modified


class TemplateService:
    def __init__(self):
        self._cache: dict[str, str] = {}

    def get_template_path(self, template_url: str) -> str:
        if not template_url or not template_url.strip():
            raise ValueError("template_url is empty")

        if settings.template_cache_enabled and template_url in self._cache:
            cached = self._cache[template_url]
            if os.path.isdir(cached):
                logger.info("Using cached template from: %s", cached)
                return cached
            self._cache.pop(template_url, None)

        existing = self._find_in_filesystem(template_url)
        if existing:
            logger.info("Found existing template in file system: %s", existing)
            if settings.template_cache_enabled:
                self._cache[template_url] = existing
            return existing

        path = self._download_and_extract(template_url)
        if settings.template_cache_enabled:
            self._cache[template_url] = path
        return path

    def _download_and_extract(self, template_url: str) -> str:
        logger.info("Processing template from: %s", template_url)
        cache_dir = settings.template_cache_dir
        os.makedirs(cache_dir, exist_ok=True)

        file_info = self._get_file_info(template_url)
        url_hash = str(abs(hash(template_url)))

        if file_info:
            existing = self._find_by_file_info(cache_dir, url_hash, file_info)
            if existing:
                logger.info("Template with same content already exists: %s", existing)
                return existing

        zip_path = self._download_zip(template_url, cache_dir)
        try:
            md5 = self._calc_md5(zip_path)
            template_dir = os.path.join(cache_dir, f"template_{url_hash}_{md5}")
            self._cleanup_old_versions(cache_dir, url_hash)
            self._extract_zip(zip_path, template_dir)
            if file_info:
                self._save_metadata(template_dir, file_info)
            logger.info("Template extracted to: %s", template_dir)
            return template_dir
        finally:
            if os.path.isfile(zip_path):
                os.remove(zip_path)

    def _download_zip(self, url: str, cache_dir: str) -> str:
        zip_path = os.path.join(cache_dir, f"temp_{int(time.time() * 1000)}.zip")
        timeout_sec = settings.template_download_timeout / 1000
        max_size = settings.template_download_max_size

        with httpx.Client(timeout=timeout_sec, follow_redirects=True) as client:
            with client.stream("GET", url) as resp:
                resp.raise_for_status()
                total = 0
                with open(zip_path, "wb") as f:
                    for chunk in resp.iter_bytes(8192):
                        f.write(chunk)
                        total += len(chunk)
                        if total > max_size:
                            raise ValueError("Downloaded size exceeds maximum allowed")

        logger.info("Downloaded zip file: %d bytes", os.path.getsize(zip_path))
        return zip_path

    @staticmethod
    def _extract_zip(zip_path: str, target_dir: str):
        os.makedirs(target_dir, exist_ok=True)
        with zipfile.ZipFile(zip_path, "r") as zf:
            zf.extractall(target_dir)
        logger.info("Extracted zip to: %s", target_dir)

    @staticmethod
    def _calc_md5(filepath: str) -> str:
        md5 = hashlib.md5()
        with open(filepath, "rb") as f:
            for chunk in iter(lambda: f.read(8192), b""):
                md5.update(chunk)
        return md5.hexdigest()

    @staticmethod
    def _cleanup_old_versions(cache_dir: str, url_hash: str):
        prefix = f"template_{url_hash}_"
        try:
            for name in os.listdir(cache_dir):
                full = os.path.join(cache_dir, name)
                if os.path.isdir(full) and name.startswith(prefix):
                    shutil.rmtree(full, ignore_errors=True)
                    logger.info("Cleaned up old template version: %s", full)
        except Exception as e:
            logger.warning("Error during cleanup: %s", e)

    def _find_in_filesystem(self, template_url: str) -> str | None:
        cache_dir = settings.template_cache_dir
        if not os.path.isdir(cache_dir):
            return None
        prefix = f"template_{hash(template_url)}_"
        try:
            for name in os.listdir(cache_dir):
                if os.path.isdir(os.path.join(cache_dir, name)) and name.startswith(prefix):
                    return os.path.join(cache_dir, name)
        except Exception as e:
            logger.warning("Error finding template: %s", e)
        return None

    def _get_file_info(self, url: str) -> _TemplateFileInfo | None:
        timeout_sec = settings.template_download_timeout / 1000
        try:
            with httpx.Client(timeout=timeout_sec) as client:
                resp = client.head(url, follow_redirects=True)
                size = int(resp.headers.get("content-length", -1))
                etag = resp.headers.get("etag")
                last_modified = resp.headers.get("last-modified")
                return _TemplateFileInfo(size, etag, last_modified)
        except Exception as e:
            logger.warning("Failed to get file info: %s", e)
            return None

    def _find_by_file_info(self, cache_dir: str, url_hash: str, info: _TemplateFileInfo) -> str | None:
        prefix = f"template_{url_hash}_"
        try:
            for name in os.listdir(cache_dir):
                full = os.path.join(cache_dir, name)
                if not os.path.isdir(full) or not name.startswith(prefix):
                    continue
                meta_file = os.path.join(full, ".metadata")
                if os.path.isfile(meta_file) and self._compare_metadata(meta_file, info):
                    return full
        except Exception as e:
            logger.warning("Error finding by file info: %s", e)
        return None

    @staticmethod
    def _compare_metadata(meta_path: str, info: _TemplateFileInfo) -> bool:
        try:
            with open(meta_path, "r") as f:
                for line in f:
                    if line.startswith("ETag:") and info.etag:
                        if line[5:].strip() == info.etag:
                            return True
                    elif line.startswith("Size:") and info.size > 0:
                        if int(line[5:].strip()) == info.size:
                            return True
        except Exception:
            pass
        return False

    @staticmethod
    def _save_metadata(template_dir: str, info: _TemplateFileInfo):
        try:
            meta = os.path.join(template_dir, ".metadata")
            with open(meta, "w") as f:
                f.write(f"ETag: {info.etag or 'unknown'}\n")
                f.write(f"Size: {info.size}\n")
                f.write(f"Last-Modified: {info.last_modified or 'unknown'}\n")
                f.write(f"Created: {time.strftime('%Y-%m-%dT%H:%M:%SZ')}\n")
        except Exception as e:
            logger.warning("Failed to save metadata: %s", e)

    def cleanup_expired_templates(self, max_age_hours: int):
        if not settings.template_cache_enabled:
            return
        cache_dir = settings.template_cache_dir
        if not os.path.isdir(cache_dir):
            return
        cutoff = time.time() - max_age_hours * 3600
        try:
            for name in os.listdir(cache_dir):
                full = os.path.join(cache_dir, name)
                if os.path.isdir(full) and os.path.getmtime(full) < cutoff:
                    shutil.rmtree(full, ignore_errors=True)
                    logger.info("Cleaned up expired template: %s", full)
        except Exception as e:
            logger.error("Error during template cleanup: %s", e)

    def get_cache_stats(self) -> dict:
        stats: dict = {
            "cacheEnabled": settings.template_cache_enabled,
            "cacheDir": settings.template_cache_dir,
            "cachedTemplates": len(self._cache),
        }
        cache_dir = settings.template_cache_dir
        if os.path.isdir(cache_dir):
            total_size = 0
            dir_count = 0
            for root, dirs, files in os.walk(cache_dir):
                for f in files:
                    total_size += os.path.getsize(os.path.join(root, f))
            for name in os.listdir(cache_dir):
                if os.path.isdir(os.path.join(cache_dir, name)) and name.startswith("template_"):
                    dir_count += 1
            stats["cacheSizeBytes"] = total_size
            stats["cacheSizeMB"] = total_size // (1024 * 1024)
            stats["templateDirectories"] = dir_count
        return stats