# apps/aroflo_connector_app/ui_automation/core/runtime.py
from __future__ import annotations

import time
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Dict

from playwright.sync_api import sync_playwright, Playwright, Browser, BrowserContext, Page

from .browser import launch_browser_context
from .paths import DEFAULT_STATE_FILE
from ..auth.session import ensure_logged_in
from ..auth.detect import is_mfa_screen


class MFARequired(RuntimeError):
    pass


@dataclass
class Runtime:
    tenant_id: str
    browser: Browser
    context: BrowserContext
    page: Page
    state_file: Path
    _pw_cm: Optional[object] = None   # context manager de sync_playwright (solo one-shot)
    _pw: Optional[Playwright] = None  # playwright instance (solo managed)
    managed: bool = False
    last_active_ts: float = 0.0

    def touch(self) -> None:
        self.last_active_ts = time.time()

    def save_state(self) -> None:
        try:
            self.context.storage_state(path=str(self.state_file))
        except Exception:
            pass

    def close(self) -> None:
        # managed: no cerrar si lo controla SessionPool (se hace en pool.close)
        try:
            self.context.close()
        except Exception:
            pass
        try:
            self.browser.close()
        except Exception:
            pass
        # one-shot: cerrar playwright CM
        if self._pw_cm:
            try:
                self._pw_cm.__exit__(None, None, None)
            except Exception:
                pass


def _tenant_state_file(cfg, tenant_id: str) -> Path:
    base = Path(getattr(cfg, "state_file", DEFAULT_STATE_FILE))
    state_dir = base.parent
    tdir = state_dir / "tenants"
    tdir.mkdir(parents=True, exist_ok=True)
    safe = tenant_id.replace("/", "_").replace("\\", "_").replace("..", "_")
    return tdir / f"{safe}.storageState.json"


class SessionPool:
    """
    Pool sync en memoria: 1 runtime por tenant.
    - Lock por tenant para evitar colisiones de UI.
    - Mantiene Playwright vivo mientras exista al menos una sesión.
    """
    def __init__(self):
        self._lock = threading.Lock()
        self._pw = None
        self._started = False
        self._sessions: Dict[str, Runtime] = {}
        self._tenant_locks: Dict[str, threading.Lock] = {}

    def _ensure_playwright(self):
        if self._started:
            return
        self._pw = sync_playwright().start()
        self._started = True

    def tenant_lock(self, tenant_id: str) -> threading.Lock:
        with self._lock:
            if tenant_id not in self._tenant_locks:
                self._tenant_locks[tenant_id] = threading.Lock()
            return self._tenant_locks[tenant_id]

    def status(self, tenant_id: str) -> dict:
        with self._lock:
            rt = self._sessions.get(tenant_id)
            if not rt:
                return {"tenant_id": tenant_id, "state": "OFFLINE"}
            return {
                "tenant_id": tenant_id,
                "state": "ONLINE",
                "state_file": str(rt.state_file),
                "last_active_ts": rt.last_active_ts,
            }

    def close(self, tenant_id: str) -> bool:
        with self._lock:
            rt = self._sessions.pop(tenant_id, None)
        if not rt:
            return False
        try:
            rt.save_state()
        except Exception:
            pass
        try:
            rt.close()
        except Exception:
            pass
        return True

    def get(
        self,
        *,
        cfg,
        tenant_id: str,
        run_dir,
        mfa_code: str,
        allow_mfa: bool,
        force_new: bool = False,
    ) -> Runtime:
        self._ensure_playwright()

        with self._lock:
            rt = self._sessions.get(tenant_id)
            if rt and not force_new:
                return rt

        # Crear runtime nuevo
        state_file = _tenant_state_file(cfg, tenant_id)
        storage_state_arg = str(state_file) if state_file.exists() else None

        browser, context = launch_browser_context(self._pw, cfg, storage_state=storage_state_arg)
        page = context.new_page()

        try:
            ensure_logged_in(page, cfg, run_dir, mfa_code=mfa_code, allow_mfa=allow_mfa)

            # Si MFA screen y no pudimos resolverlo, señal controlada
            if is_mfa_screen(page):
                raise MFARequired("MFA_REQUIRED")

            # Guardar state por tenant
            try:
                context.storage_state(path=str(state_file))
            except Exception:
                pass

            rt = Runtime(
                tenant_id=tenant_id,
                browser=browser,
                context=context,
                page=page,
                state_file=state_file,
                managed=True,
                _pw=self._pw,
            )
            rt.touch()

            with self._lock:
                self._sessions[tenant_id] = rt

            return rt

        except Exception:
            try:
                context.close()
            except Exception:
                pass
            try:
                browser.close()
            except Exception:
                pass
            raise


SESSION_POOL = SessionPool()


def open_one_shot(cfg, *, state_file: Optional[Path] = None) -> Runtime:
    """
    Modo legacy: abre playwright/browser/context/page y lo cierra al final.
    """
    cm = sync_playwright()
    p = cm.__enter__()
    sf = Path(state_file) if state_file else Path(getattr(cfg, "state_file", DEFAULT_STATE_FILE))
    storage_state_arg = str(sf) if sf.exists() else None
    browser, context = launch_browser_context(p, cfg, storage_state=storage_state_arg)
    page = context.new_page()
    rt = Runtime(
        tenant_id="",
        browser=browser,
        context=context,
        page=page,
        state_file=sf,
        _pw_cm=cm,
        managed=False,
    )
    rt.touch()
    return rt