import os
import sys
import json
import asyncio
import functools
import nest_asyncio
from typing import Optional, Dict, Any, List

import click

from AgentCrew.setup import ApplicationSetup, PROVIDER_LIST
from AgentCrew.modules.config import ConfigManagement
from AgentCrew.modules.llm.service_manager import ServiceManager
from AgentCrew.modules.agents.local_agent import LocalAgent

nest_asyncio.apply()


def common_options(func):
    @click.option(
        "--provider",
        type=click.Choice(PROVIDER_LIST),
        default=None,
        help="LLM provider to use (claude, groq, openai, google, github_copilot, or deepinfra)",
    )
    @click.option(
        "--agent-config",
        default=None,
        help="Path/URL to the agent configuration file.",
    )
    @click.option(
        "--mcp-config", default=None, help="Path to the mcp servers configuration file."
    )
    @click.option(
        "--memory-llm",
        type=click.Choice(
            ["claude", "groq", "openai", "google", "deepinfra", "github_copilot"]
        ),
        default=None,
        help="LLM Model use for analyzing and processing memory",
    )
    @click.option(
        "--memory-path", default=None, help="Path to the memory database location"
    )
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)

    return wrapper


class AgentCrewApplication:
    def __init__(self):
        self.config_manager = ConfigManagement()
        self.setup = ApplicationSetup(self.config_manager)
        self.setup.load_api_keys_from_config()

    @property
    def services(self) -> Optional[Dict[str, Any]]:
        return self.setup.services

    @property
    def agent_manager(self):
        return self.setup.agent_manager

    def run_console(
        self,
        provider: Optional[str] = None,
        agent_config: Optional[str] = None,
        mcp_config: Optional[str] = None,
        memory_llm: Optional[str] = None,
        with_voice: bool = False,
    ) -> None:
        from AgentCrew.modules.console import ConsoleUI
        from AgentCrew.modules.chat import MessageHandler
        from AgentCrew.modules.mcpclient import MCPSessionManager

        try:
            if provider is None:
                provider = self.setup.detect_provider()
                if provider is None:
                    raise ValueError(
                        "No LLM API key found. Please set either ANTHROPIC_API_KEY, GEMINI_API_KEY, OPENAI_API_KEY, GROQ_API_KEY, or DEEPINFRA_API_KEY"
                    )

            services = self.setup.setup_services(provider, memory_llm)

            if mcp_config:
                os.environ["MCP_CONFIG_PATH"] = mcp_config

            self.setup.setup_agents(services, agent_config)
            self.setup.restore_last_agent()

            message_handler = MessageHandler(
                services["memory"], services["context_persistent"], with_voice
            )
            global_config = self.config_manager.read_global_config_data()

            ui = ConsoleUI(
                message_handler,
                global_config.get("global_settings", {}).get("swap_enter", False),
            )
            ui.start()
        except Exception as e:
            import traceback

            print(traceback.format_exc())
            click.echo(f"❌ Error: {str(e)}", err=True)
        finally:
            MCPSessionManager.get_instance().cleanup()

    def run_gui(
        self,
        provider: Optional[str] = None,
        agent_config: Optional[str] = None,
        mcp_config: Optional[str] = None,
        memory_llm: Optional[str] = None,
        with_voice: bool = False,
    ) -> None:
        from PySide6.QtCore import QCoreApplication
        from PySide6.QtCore import Qt
        from PySide6.QtWidgets import QApplication
        from AgentCrew.modules.gui import ChatWindow
        from AgentCrew.modules.chat import MessageHandler
        from AgentCrew.modules.mcpclient import MCPSessionManager

        try:
            if provider is None:
                provider = self.setup.detect_provider()
                if provider is None:
                    from AgentCrew.modules.gui.widgets.config_window import ConfigWindow

                    app = QApplication(sys.argv)
                    config_window = ConfigWindow()
                    config_window.tab_widget.setCurrentIndex(3)
                    config_window.show()
                    sys.exit(app.exec())

            services = self.setup.setup_services(provider, memory_llm)

            if mcp_config:
                os.environ["MCP_CONFIG_PATH"] = mcp_config

            self.setup.setup_agents(services, agent_config)
            self.setup.restore_last_agent()

            message_handler = MessageHandler(
                services["memory"], services["context_persistent"], with_voice
            )

            QCoreApplication.setAttribute(Qt.ApplicationAttribute.AA_UseOpenGLES)
            app = QApplication(sys.argv)
            chat_window = ChatWindow(message_handler)
            chat_window.show()
            sys.exit(app.exec())
        except Exception as e:
            import traceback

            print(traceback.format_exc())
            click.echo(f"❌ Error: {str(e)}", err=True)
        finally:
            MCPSessionManager.get_instance().cleanup()

    def run_server(
        self,
        host: str = "0.0.0.0",
        port: int = 41241,
        base_url: Optional[str] = None,
        provider: Optional[str] = None,
        model_id: Optional[str] = None,
        agent_config: Optional[str] = None,
        api_key: Optional[str] = None,
        mcp_config: Optional[str] = None,
        memory_llm: Optional[str] = None,
        store_type: str = "memory",
        store_options: Optional[dict] = None,
    ) -> None:
        from AgentCrew.modules.a2a.server import A2AServer
        from AgentCrew.modules.mcpclient import MCPSessionManager

        try:
            if not base_url:
                base_url = f"http://{host}:{port}"

            if provider is None:
                provider = self.setup.detect_provider()
                if provider is None:
                    raise ValueError(
                        "No LLM API key found. Please set either ANTHROPIC_API_KEY, GEMINI_API_KEY, OPENAI_API_KEY, GROQ_API_KEY, or DEEPINFRA_API_KEY"
                    )

            services = self.setup.setup_services(
                provider, memory_llm, need_memory=False
            )

            if mcp_config:
                os.environ["MCP_CONFIG_PATH"] = mcp_config

            os.environ["AGENTCREW_DISABLE_GUI"] = "true"

            self.setup.setup_agents(services, agent_config, provider, model_id)

            if self.agent_manager is None:
                raise ValueError("Agent manager is not initialized")

            self.agent_manager.enforce_transfer = False

            server = A2AServer(
                agent_manager=self.agent_manager,
                host=host,
                port=port,
                base_url=base_url,
                api_key=api_key,
                store_type=store_type,
                store_options=store_options,
            )

            click.echo(f"Starting A2A server on {host}:{port}")
            click.echo(
                f"Available agents: {', '.join(self.agent_manager.agents.keys())}"
            )
            server.start()
        except Exception as e:
            import traceback

            print(traceback.format_exc())
            click.echo(f"❌ Error: {str(e)}", err=True)
        finally:
            MCPSessionManager.get_instance().cleanup()

    def _parse_output_schema(self, schema_input: str) -> tuple[str, dict]:
        try:
            from AgentCrew.modules.prompts.constants import SCHEMA_ENFORCEMENT_PROMPT

            if os.path.exists(schema_input):
                with open(schema_input, "r", encoding="utf-8") as f:
                    schema_dict = json.load(f)
            else:
                schema_dict = json.loads(schema_input)

            schema_json = json.dumps(schema_dict, indent=2)

            enforcement_prompt = SCHEMA_ENFORCEMENT_PROMPT.replace(
                "{schema_json}", schema_json
            )
            return enforcement_prompt, schema_dict

        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON schema: {e}")
        except Exception as e:
            raise ValueError(f"Failed to load output schema: {e}")

    def _clean_json_response(self, response: str) -> str:
        import re

        cleaned = response.strip()
        pattern = r"^```(?:json)?\s*\n?(.*?)\n?```$"
        match = re.match(pattern, cleaned, re.DOTALL)
        if match:
            cleaned = match.group(1).strip()
        return cleaned

    def _validate_response_against_schema(
        self, response: str, schema_dict: Dict[str, Any]
    ) -> tuple[bool, Optional[str]]:
        from jsonschema import validate, ValidationError

        try:
            cleaned_response = self._clean_json_response(response)
            response_json = json.loads(cleaned_response)
        except json.JSONDecodeError as e:
            return (
                False,
                f"Response is not valid JSON: {e}\n\nResponse received:\n{response[:500]}",
            )

        try:
            validate(instance=response_json, schema=schema_dict)
            return True, None
        except ValidationError as e:
            error_details = (
                "JSON Schema Validation Error:\n"
                f"  - Path: {' -> '.join(str(p) for p in e.path) if e.path else 'root'}\n"
                f"  - Error: {e.message}\n"
                f"  - Failed value: {json.dumps(e.instance, indent=2)}\n"
            )
            if e.schema_path:
                error_details += (
                    f"  - Schema path: {' -> '.join(str(p) for p in e.schema_path)}\n"
                )
            return False, error_details

    def run_job(
        self,
        agent: str,
        task: str,
        files: Optional[List[str]] = None,
        provider: Optional[str] = None,
        model_id: Optional[str] = None,
        agent_config: Optional[str] = None,
        mcp_config: Optional[str] = None,
        memory_llm: Optional[str] = None,
        output_schema: Optional[str] = None,
    ) -> str:
        from AgentCrew.modules.chat import MessageHandler
        from AgentCrew.modules.mcpclient import MCPSessionManager
        from AgentCrew.modules.llm.model_registry import ModelRegistry

        try:
            if provider is None:
                provider = self.setup.detect_provider()
                if provider is None:
                    raise ValueError(
                        "No LLM API key found. Please set either ANTHROPIC_API_KEY, GEMINI_API_KEY, OPENAI_API_KEY, GROQ_API_KEY, or DEEPINFRA_API_KEY"
                    )

            services = self.setup.setup_services(
                provider, memory_llm, need_memory=False
            )

            if mcp_config:
                os.environ["MCP_CONFIG_PATH"] = mcp_config

            os.environ["AGENTCREW_DISABLE_GUI"] = "true"

            self.setup.setup_agents(services, agent_config)

            llm_manager = ServiceManager.get_instance()

            llm_service = llm_manager.get_service(provider)
            if model_id:
                llm_service.model = model_id

            if self.agent_manager is None:
                raise ValueError("Agent manager is not initialized")

            self.agent_manager.update_llm_service(llm_service)

            for local_agent in self.agent_manager.agents:
                if isinstance(local_agent, LocalAgent):
                    local_agent.is_remoting_mode = True

            self.agent_manager.enforce_transfer = False
            self.agent_manager.one_turn_process = True

            current_agent = self.agent_manager.get_local_agent(agent)

            if current_agent:
                schema_dict = None
                if output_schema:
                    schema_prompt, schema_dict = self._parse_output_schema(
                        output_schema
                    )
                    if "structured_output" in ModelRegistry.get_model_capabilities(
                        f"{provider}/{model_id}"
                    ):
                        current_agent.llm.structured_output = schema_dict
                    else:
                        current_agent.set_custom_system_prompt(schema_prompt)

                self.agent_manager.select_agent(current_agent.name)

                message_handler = MessageHandler(
                    services["memory"], services["context_persistent"]
                )
                message_handler.is_non_interactive = True
                message_handler.agent = current_agent

                if files:
                    for file_path in files:
                        asyncio.run(
                            message_handler.process_user_input(f"/file {file_path}")
                        )

                max_attempts = 4
                attempt = 0
                response = None

                asyncio.run(message_handler.process_user_input(task))

                while attempt < max_attempts:
                    attempt += 1
                    response, _, _ = asyncio.run(
                        message_handler.get_assistant_response()
                    )
                    if not output_schema or not schema_dict:
                        break

                    if response is None:
                        asyncio.run(
                            message_handler.process_user_input(
                                "No response was generated. Please try again."
                            )
                        )
                        continue

                    success, retry_message = self._validate_response_against_schema(
                        response, schema_dict
                    )
                    if success:
                        break
                    else:
                        if retry_message:
                            asyncio.run(
                                message_handler.process_user_input(retry_message)
                            )

                MCPSessionManager.get_instance().cleanup()
                return response.strip() if response else ""
            else:
                raise ValueError(f"Agent '{agent}' not found")

        except Exception:
            import traceback

            print(traceback.format_exc())
            raise

    def login(self) -> bool:
        return self.setup.login()
