# screen_capturer.py
"""
Capture d'écran pour agent_v0.

v0+ :
- capture de l'écran principal en PNG
- mode "full" : écran complet
- mode "crop" : zone centrée autour d'une position (curseur)
- stockage dans sessions/<session_id>/shots/shot_XXXX.png
"""

from __future__ import annotations

import os
import time
from typing import Tuple, Optional

from mss import mss
from mss.tools import to_png

from config import SESSIONS_ROOT
from raw_session import RawSession


class ScreenCapturer:
    """
    Gère la capture d'écran pour une session donnée.
    """

    def __init__(
        self,
        session: RawSession,
        base_dir: str = SESSIONS_ROOT,
        screenshot_mode: str = "full",  # "full" | "crop"
        crop_width: int = 800,
        crop_height: int = 600,
    ) -> None:
        self.session = session
        self.base_dir = base_dir
        self._counter = 0

        self.screenshot_mode = screenshot_mode
        self.crop_width = crop_width
        self.crop_height = crop_height

        # Résolution écran pour clamp la zone de crop
        self.screen_width, self.screen_height = self.detect_primary_resolution()

    def _get_session_shots_dir(self) -> str:
        session_dir = os.path.join(self.base_dir, self.session.session_id)
        shots_dir = os.path.join(session_dir, "shots")
        os.makedirs(shots_dir, exist_ok=True)
        return shots_dir

    def capture(self, focus_pos: Optional[Tuple[int, int]] = None) -> Tuple[str, str]:
        """
        Capture l'écran :
        - si screenshot_mode == "crop" et focus_pos fourni -> zone centrée autour de focus_pos
        - sinon -> plein écran

        Retourne (screenshot_id, relative_path)
        """
        self._counter += 1
        screenshot_id = f"shot_{self._counter:04d}"
        shots_dir = self._get_session_shots_dir()
        filename = f"{screenshot_id}.png"
        file_path = os.path.join(shots_dir, filename)
        relative_path = os.path.join("shots", filename)

        with mss() as sct:
            if self.screenshot_mode == "crop" and focus_pos is not None:
                cx, cy = focus_pos
                w = min(self.crop_width, self.screen_width)
                h = min(self.crop_height, self.screen_height)

                left = max(0, cx - w // 2)
                top = max(0, cy - h // 2)

                if left + w > self.screen_width:
                    left = self.screen_width - w
                if top + h > self.screen_height:
                    top = self.screen_height - h

                region = {"left": int(left), "top": int(top), "width": int(w), "height": int(h)}
                img = sct.grab(region)
            else:
                monitor = sct.monitors[0]  # écran principal complet
                img = sct.grab(monitor)

            with open(file_path, "wb") as f:
                f.write(to_png(img.rgb, img.size))

        captured_at = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
        self.session.add_screenshot(
            screenshot_id=screenshot_id,
            relative_path=relative_path,
            captured_at=captured_at,
        )

        return screenshot_id, relative_path

    @staticmethod
    def detect_primary_resolution() -> Tuple[int, int]:
        """
        Renvoie (width, height) de l'écran principal.
        Utile pour remplir RawSession.environment.screen.primary_resolution.
        """
        with mss() as sct:
            mon = sct.monitors[0]
            return mon["width"], mon["height"]