import json
from collections.abc import Generator
from typing import TYPE_CHECKING
import os
import gradio as gr
from datetime import datetime
from typing import Optional, Dict, List, Any
import uuid
import asyncio

from gradio.components import Component
from browser_use.browser.browser import Browser
from browser_use.browser.context import BrowserContext
from browser_use.agent.service import Agent
from src.browser.custom_browser import CustomBrowser
from src.browser.custom_context import CustomBrowserContext
from src.controller.custom_controller import CustomController
from src.agent.deep_research.deep_research_agent import DeepResearchAgent


class WebuiManager:
    def __init__(self, settings_save_dir: str = "./tmp/webui_settings"):
        self.id_to_component: dict[str, Component] = {}
        self.component_to_id: dict[Component, str] = {}

        self.settings_save_dir = settings_save_dir
        os.makedirs(self.settings_save_dir, exist_ok=True)

        # Add type annotation for vayner_pdf_report
        self.vayner_pdf_report: Optional[str] = None

    def init_browser_use_agent(self) -> None:
        """
        init browser use agent
        """
        self.bu_agent: Optional[Agent] = None
        self.bu_browser: Optional[CustomBrowser] = None
        self.bu_browser_context: Optional[CustomBrowserContext] = None
        self.bu_controller: Optional[CustomController] = None
        self.bu_chat_history: List[Dict[str, Optional[str]]] = []
        self.bu_response_event: Optional[asyncio.Event] = None
        self.bu_user_help_response: Optional[str] = None
        self.bu_current_task: Optional[asyncio.Task] = None
        self.bu_agent_task_id: Optional[str] = None
        self.bu_task_metrics: Optional[Dict[str, Any]] = None
        
    def init_vayner_client_research(self) -> None:
        """
        Initialize Vayner Client Research components and state
        """
        if not hasattr(self, "vayner_chat_history"):
            self.vayner_chat_history = []
        
        if not hasattr(self, "vayner_pdf_report"):
            self.vayner_pdf_report = None
            
        if not hasattr(self, "vayner_controller"):
            self.vayner_controller = None
            
        if not hasattr(self, "vayner_browser"):
            self.vayner_browser = None
            
        if not hasattr(self, "vayner_browser_context"):
            self.vayner_browser_context = None
            
        if not hasattr(self, "vayner_agent"):
            self.vayner_agent = None
            
        if not hasattr(self, "vayner_current_task"):
            self.vayner_current_task = None
            
        # Initialize data collections for PDF report
        if not hasattr(self, "vayner_screenshots"):
            self.vayner_screenshots = []
            
        if not hasattr(self, "vayner_business_info"):
            self.vayner_business_info = []
            
        if not hasattr(self, "vayner_keyword_data"):
            self.vayner_keyword_data = []
            
        if not hasattr(self, "vayner_ranking_data"):
            self.vayner_ranking_data = []
            
        if not hasattr(self, "vayner_current_business"):
            self.vayner_current_business = "Unknown Business"
            
        # New: keyword table rows for third page
        if not hasattr(self, "vayner_keyword_table_rows"):
            self.vayner_keyword_table_rows = []
            
        # Queue for updates during task execution
        if not hasattr(self, "update_queue"):
            self.update_queue = []

    def init_deep_research_agent(self) -> None:
        """
        init deep research agent
        """
        self.dr_agent: Optional[DeepResearchAgent] = None
        self.dr_current_task = None
        self.dr_agent_task_id: Optional[str] = None
        self.dr_save_dir: Optional[str] = None

    def add_components(self, tab_name: str, components_dict: dict[str, "Component"]) -> None:
        """
        Add tab components
        """
        for comp_name, component in components_dict.items():
            comp_id = f"{tab_name}.{comp_name}"
            self.id_to_component[comp_id] = component
            self.component_to_id[component] = comp_id

    def get_components(self) -> list["Component"]:
        """
        Get all components
        """
        return list(self.id_to_component.values())

    def get_component_by_id(self, comp_id: str) -> Optional["Component"]:
        """
        Get component by id. Returns None if not found.
        """
        return self.id_to_component.get(comp_id, None)

    def get_id_by_component(self, comp: "Component") -> str:
        """
        Get id by component. Raises KeyError if not found.
        """
        return self.component_to_id[comp]

    def save_config(self, components: Dict["Component", str]) -> str:
        """
        Save config
        """
        cur_settings = {}
        for comp in components:
            if not isinstance(comp, gr.Button) and not isinstance(comp, gr.File) and str(
                    getattr(comp, "interactive", True)).lower() != "false":
                comp_id = self.get_id_by_component(comp)
                cur_settings[comp_id] = components[comp]

        config_name = datetime.now().strftime("%Y%m%d-%H%M%S")
        with open(os.path.join(self.settings_save_dir, f"{config_name}.json"), "w") as fw:
            json.dump(cur_settings, fw, indent=4)

        return os.path.join(self.settings_save_dir, f"{config_name}.json")

    def load_config(self, config_path: str):
        """
        Load config
        """
        with open(config_path, "r") as fr:
            ui_settings = json.load(fr)

        update_components = {}
        for comp_id, comp_val in ui_settings.items():
            if comp_id in self.id_to_component:
                comp = self.id_to_component[comp_id]
                update_components[comp] = gr.update(value=comp_val)

        config_status = self.id_to_component["load_save_config.config_status"]
        update_components.update(
            {
                config_status: gr.update(value=f"Successfully loaded config: {config_path}")
            }
        )
        yield update_components