from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
import uuid
from pydantic import BaseModel, Field


class BrowserInstance(BaseModel):
    id: str = Field(default_factory=lambda: f"browser-{uuid.uuid4().hex[:8]}")
    status: str = "WARM"
    active_mission: Optional[str] = None
    tabs_count: int = 1
    tabs: List[Dict[str, Any]] = Field(default_factory=lambda: [{"tab_id": "tab-1", "url": "about:blank", "title": "Blank"}])
    cookies: List[Dict[str, Any]] = Field(default_factory=list)
    history: List[str] = Field(default_factory=list)
    snapshots: List[Dict[str, Any]] = Field(default_factory=list)
    last_screenshot_base64: Optional[str] = None
    downloads: List[Dict[str, Any]] = Field(default_factory=list)
    health_score: float = 100.0
    memory_mb: float = 128.5
    created_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat())


class BrowserPoolManager:
    """Manages active browser instances, pool scaling, health monitoring, and session recycling."""

    def __init__(self, max_pool_size: int = 5):
        self.max_pool_size = max_pool_size
        self.pool: Dict[str, BrowserInstance] = {}
        self._initialize_pool()

    def _initialize_pool(self):
        for _ in range(min(2, self.max_pool_size)):
            inst = BrowserInstance()
            self.pool[inst.id] = inst

    def get_status(self) -> Dict[str, Any]:
        return {
            "total_browsers": len(self.pool),
            "max_pool_size": self.max_pool_size,
            "active_browsers": len([b for b in self.pool.values() if b.active_mission]),
            "browsers": list(self.pool.values()),
        }

    def acquire_instance(self, mission_id: str) -> BrowserInstance:
        for inst in self.pool.values():
            if not inst.active_mission:
                inst.active_mission = mission_id
                inst.status = "BUSY"
                return inst
        if len(self.pool) < self.max_pool_size:
            new_inst = BrowserInstance(status="BUSY", active_mission=mission_id)
            self.pool[new_inst.id] = new_inst
            return new_inst
        first_inst = list(self.pool.values())[0]
        return first_inst

    def release_instance(self, browser_id: str):
        if browser_id in self.pool:
            self.pool[browser_id].active_mission = None
            self.pool[browser_id].status = "WARM"

    def update_screenshot(self, browser_id: str, screenshot_base64: str, url: str):
        if browser_id in self.pool:
            inst = self.pool[browser_id]
            inst.last_screenshot_base64 = screenshot_base64
            if url not in inst.history:
                inst.history.append(url)

    def get_mission_screenshot(self, mission_id: str) -> Optional[str]:
        for inst in self.pool.values():
            if inst.active_mission == mission_id and inst.last_screenshot_base64:
                return inst.last_screenshot_base64
        return None

    def get_all_thumbnails(self) -> List[Dict[str, Any]]:
        return [
            {
                "browser_id": b.id,
                "mission_id": b.active_mission,
                "status": b.status,
                "url": b.tabs[0]["url"] if b.tabs else "about:blank",
                "has_screenshot": b.last_screenshot_base64 is not None,
            }
            for b in self.pool.values()
        ]