import logging
import sys
from logging.handlers import RotatingFileHandler

import dspy
from dspy.utils.callback import BaseCallback

from config_loader import load_config
from experts.retrieval_agent import DnDRAG

CFG = load_config()
RETRIEVE_MODEL = CFG["models"]["retrieval"]
API_BASE = CFG["api"]["base_url"]
API_VERSION = CFG["api"]["api_version"]


class CallbackHandler(BaseCallback):
    """Custom callback class for logging agent interactions."""

    def __init__(self, logger):
        """Initialize the callback with a logger instance."""
        super().__init__()
        self.logger = logger

    def on_module_end(self, call_id, outputs, exception):
        """Handle module end events for logging."""
        step = "Reasoning" if self._is_reasoning_output(outputs) else "Acting"
        self.logger.debug(f"== {step} Step ===")
        for k, v in outputs.items():
            self.logger.debug(f"  {k}: {v}")

    def on_lm_start(self, call_id, instance, inputs):
        """Handle language model start events for logging."""
        self.logger.debug(f"LM is called with inputs: {inputs}")

    def on_tool_start(self, call_id, instance, inputs):
        """Handle tool start events for logging."""
        self.logger.debug(f"Tool {instance} called with inputs: {inputs}")

    def on_tool_end(self, call_id, outputs, exception):
        """Handle tool end events for logging."""
        self.logger.debug(f"Tool finished with outputs: {outputs}")

    def on_lm_end(self, call_id, outputs, exception):
        """Handle language model end events for logging."""
        self.logger.debug(f"LM is finished with outputs: {outputs}")

    def _is_reasoning_output(self, outputs):
        return any(k.startswith("Thought") for k in outputs)


def setup_logging():
    """Set up logging configuration for Merlin."""
    # Create a custom logger
    logger = logging.getLogger(__name__)

    # Set the minimum level for the logger
    logger.setLevel(logging.DEBUG)

    # Create a console handler
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    # Create a file handler with rotation every 5MB
    file_handler = RotatingFileHandler("data/dmv.log", maxBytes=5 * 1024 * 1024, backupCount=3)
    file_handler.setLevel(logging.DEBUG)

    # Create a formatter
    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")

    # Set the formatter for the handler
    console_handler.setFormatter(formatter)
    file_handler.setFormatter(formatter)

    # Add the handler to the logger
    logger.addHandler(console_handler)
    logger.addHandler(file_handler)

    return logger


def main():
    logger = setup_logging()
    logger.debug("main application started")

    # Add verbose callback
    dspy.configure(verbose_errors=True)
    dspy.configure(callbacks=[CallbackHandler(logger)])
    # 1. Setup the LLM
    print("🚀 Initializing Qwen-8B via LM Studio...")
    lm = dspy.LM(RETRIEVE_MODEL, api_base=API_BASE + API_VERSION)
    dspy.configure(lm=lm)

    # 2. Load the RAG System (only happens once!)
    print("📚 Loading campaign notes...")
    try:
        rag_system = DnDRAG()
        print("✅ Ready! Ask me anything about the campaign. (Type 'exit' or 'q' to quit)")
    except Exception as e:
        print(f"❌ Failed to initialize: {e}")
        return

    # 3. Interactive Loop
    while True:
        try:
            print("\n" + "─" * 30)
            query = input("📝 Query: ").strip()

            # Exit conditions
            if query.lower() in ["exit", "quit", "q", "bye"]:
                print("Farewell, traveler. Good luck on your quest!")
                break

            if not query:
                continue

            print("🔍 Searching and thinking...")
            response = rag_system(question=query)

            # Print Response
            print("\n📜 AI RESPONSE:")
            print(response.answer)

        except KeyboardInterrupt:
            print("\n\nRude?!.... Exiting...")
            sys.exit(0)
        except Exception as e:
            print(f"\n⚠️ An error occurred: {e}")


if __name__ == "__main__":
    main()