import os
from datetime import datetime
from enum import Enum
from typing import Optional

from mojentic.llm.gateways import OllamaGateway, OpenAIGateway
from pydantic import BaseModel, Field


class ModelGateway(str, Enum):
    OLLAMA = "ollama"
    OPENAI = "openai"


def get_available_models(gateway: ModelGateway = ModelGateway.OLLAMA) -> list[str]:
    if gateway == ModelGateway.OLLAMA:
        g = OllamaGateway()
    elif gateway == ModelGateway.OPENAI:
        openai_key = os.environ.get("OPENAI_API_KEY")
        if not openai_key:
            print("Error: OPENAI_API_KEY environment variable is not set.")
            return []
        g = OpenAIGateway(openai_key)
    return g.get_available_models()


def select_model(gateway: ModelGateway = ModelGateway.OLLAMA, is_visual: bool = False) -> str:
    model_type = "visual analysis" if is_visual else "chat"
    models = get_available_models(gateway)
    if not models:
        if gateway == ModelGateway.OLLAMA:
            return input(
                f"No models found in Ollama. Please enter {model_type} model name manually: ")
        else:
            return input(f"No models available. Please enter {model_type} model name manually: ")

    print(f"\nAvailable {gateway.value} models for {model_type}:")
    for idx, model in enumerate(models, 1):
        print(f"{idx}. {model}")

    while True:
        try:
            choice = int(input(f"\nSelect a {model_type} model (enter number): "))
            if 1 <= choice <= len(models):
                return models[choice - 1]
        except ValueError:
            pass
        print("Invalid selection. Please try again.")


def get_config_path(vault_path: str) -> str:
    """Get the path to the config file in the vault directory."""
    return os.path.join(vault_path, ".zk_chat")


class Config(BaseModel):
    vault: str
    model: str  # Chat model
    visual_model: str | None = None  # Visual analysis model
    gateway: ModelGateway = ModelGateway.OLLAMA
    chunk_size: int = 500
    chunk_overlap: int = 100
    last_indexed: datetime | None = None  # Deprecated, kept for backward compatibility
    gateway_last_indexed: dict[str, datetime] = Field(default_factory=dict)

    def get_last_indexed(self, gateway: ModelGateway | None = None) -> datetime | None:
        """
        Get the last indexed time for the specified gateway or the current gateway if not specified.
        Falls back to the deprecated last_indexed field if the gateway-specific value is not found.
        """
        gateway_value = gateway.value if gateway else self.gateway.value
        if gateway_value in self.gateway_last_indexed:
            return self.gateway_last_indexed[gateway_value]
        return self.last_indexed

    def set_last_indexed(self, timestamp: datetime, gateway: ModelGateway | None = None):
        """
        Set the last indexed time for the specified gateway or the current gateway if not specified.
        """
        gateway_value = gateway.value if gateway else self.gateway.value
        self.gateway_last_indexed[gateway_value] = timestamp

    @classmethod
    def load(cls, vault_path: str) -> Optional['Config']:
        config_path = get_config_path(vault_path)
        if os.path.exists(config_path):
            with open(config_path) as f:
                return cls.model_validate_json(f.read())
        else:
            return None

    @classmethod
    def load_or_initialize(cls, vault_path: str, gateway: ModelGateway = ModelGateway.OLLAMA,
                           model: str = None, visual_model: str = None) -> 'Config':
        config = cls.load(vault_path)
        if config:
            # If config exists, just return it - visual_model can be None
            return config

        # Initialize new config
        if model is None:
            print("Please select a model for chat:")
            model = select_model(gateway)

        if visual_model is None:
            print("Would you like to select a model for visual analysis? (y/n): ")
            choice = input().strip().lower()
            if choice == 'y':
                print("Please select a model for visual analysis:")
                visual_model = select_model(gateway, is_visual=True)
            else:
                visual_model = None
                print("Visual analysis will be disabled.")

        config = cls(vault=vault_path, model=model, visual_model=visual_model, gateway=gateway)

        config.save()
        return config

    def save(self) -> None:
        config_path = get_config_path(self.vault)
        with open(config_path, 'w') as f:
            f.write(self.model_dump_json(indent=2))

    def update_model(self, model_name: str = None, gateway: ModelGateway = None,
                     is_visual: bool = False) -> None:
        """
        Update the model in config. If model_name is None, interactive selection will be used.

        Args:
            model_name: Name of the model to use
            gateway: Gateway to use
            is_visual: If True, update the visual model instead of the chat model
        """
        # Update gateway if specified
        if gateway is not None:
            self.gateway = gateway

        if model_name:
            available_models = get_available_models(self.gateway)
            if model_name in available_models:
                if is_visual:
                    self.visual_model = model_name
                else:
                    self.model = model_name
            else:
                print(f"Model '{model_name}' not found in available models.")
                selected_model = select_model(self.gateway)
                if is_visual:
                    self.visual_model = selected_model
                else:
                    self.model = selected_model
        else:
            selected_model = select_model(self.gateway)
            if is_visual:
                self.visual_model = selected_model
            else:
                self.model = selected_model

        model_type = "Visual model" if is_visual else "Chat model"
        model_name = self.visual_model if is_visual else self.model
        print(f"{model_type} selected: {model_name} (using {self.gateway.value} gateway)")
        self.save()
