"""Template serialization schema for NIKAME."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel, Field, field_validator

CURRENT_SCHEMA_VERSION = 1

class WizardSession(BaseModel):
    """Pydantic model representing a serialized wizard state/template."""
    schema_version: int = CURRENT_SCHEMA_VERSION
    
    # 1. Project Basics
    name: str = "my-app"
    target: Literal["local", "kubernetes", "aws", "gcp"] = "local"
    profile: Literal["local", "staging", "production"] = "local"
    scale: str = "small"
    access_pattern: str = "balanced"
    
    # 2. Infrastructure
    databases: list[str] = Field(default_factory=list)
    cache: str = "dragonfly"
    messaging: str = "none"
    gateway: str = "traefik"
    observability: str = "none"
    ci_cd: list[str] = Field(default_factory=list)
    
    # 2.1 Contextual Infra Follow-ups
    pg_replicas: int = 1
    msg_mps: str = "under 1k"
    msg_partitions: int = 3
    auth_pattern: str = "JWT stateless"
    storage_buckets: list[str] = Field(default_factory=list)
    alert_channels: list[str] = Field(default_factory=list)
    cloud_provider: str = "bare-metal"
    
    # 2.2 Credentials
    pg_user: str = "postgres"
    pg_password: str = ""
    pg_db: str = "app"
    redis_password: str = ""
    minio_user: str = "minioadmin"
    minio_password: str = ""
    keycloak_admin: str = "admin"
    keycloak_password: str = ""
    grafana_password: str = ""
    
    # 3. Features & Components
    api_framework: str = "fastapi"
    features: list[str] = Field(default_factory=list)
    generate_guide: bool = True
    
    # 3.1 Advanced Feature Follow-ups
    tenancy_model: str = "separate schemas"
    max_rps: int = 100
    
    # 4. MLOps (Optional)
    enable_mlops: bool = False
    ml_serving: list[str] = Field(default_factory=list)
    ml_tracking: list[str] = Field(default_factory=list)
    ml_orchestration: list[str] = Field(default_factory=list)
    ml_monitoring: list[str] = Field(default_factory=list)
    ml_vector_dbs: list[str] = Field(default_factory=list)
    ml_agents: list[str] = Field(default_factory=list)
    ml_caching: list[str] = Field(default_factory=list)
    
    # 4.1 MLOps Follow-ups
    ml_training_enabled: bool = False
    ml_orchestrator: str = "prefect"
    ml_embedding_dim: int = 1536
    
    # 5. Data Models
    models: dict[str, Any] = Field(default_factory=dict)
    
    # Tracking dirty steps for cascading invalidation
    completed_steps: list[int] = Field(default_factory=list)
    auto_added_modules: list[str] = Field(default_factory=list)

    def to_config_dict(self) -> dict[str, Any]:
        """Convert the flat wizard session into a structured NikameConfig dict."""
        from nikame.config.schema import (
            CICDConfig,
            DragonflyConfig,
            KeycloakConfig,
            PostgresConfig,
            RedisConfig,
        )

        config: dict[str, Any] = {
            "name": self.name,
            "description": f"Generated from template for {self.name}",
            "project": {"scale": self.scale, "access_pattern": self.access_pattern},
            "environment": {"target": self.target, "profile": self.profile},
            "models": self.models,
            "features": self.features,
            "generate_guide": self.generate_guide,
        }

        # Databases
        db_config: dict[str, Any] = {}
        if "postgres" in self.databases:
            db_config["postgres"] = PostgresConfig(
                user=self.pg_user,
                password=self.pg_password,
                database=self.pg_db,
                replicas=self.pg_replicas,
            ).model_dump()
        if "redis" in self.databases:
            db_config["redis"] = RedisConfig(password=self.redis_password).model_dump()
        if db_config:
            config["databases"] = db_config

        # Cache
        if self.cache == "dragonfly":
            config["cache"] = {
                "provider": "dragonfly",
                "dragonfly": DragonflyConfig().model_dump(),
            }
        elif self.cache == "redis":
            config["cache"] = {
                "provider": "redis",
            }

        # Storage
        if "storage" in self.features:
            config["storage"] = {
                "provider": "minio",
                "root_user": self.minio_user,
                "root_password": self.minio_password,
                "buckets": self.storage_buckets,
            }

        # Auth
        if "auth" in self.features:
            config["auth"] = {
                "provider": "keycloak" if "keycloak" in self.auth_pattern.lower() else "postgres",
                "keycloak": KeycloakConfig(
                    admin_user=self.keycloak_admin, admin_password=self.keycloak_password
                ).model_dump(),
            }

        # CI/CD
        config["ci_cd"] = CICDConfig(
            github_actions="github" in "".join(self.ci_cd).lower(),
            argocd="argo" in "".join(self.ci_cd).lower(),
        ).model_dump()

        # Gateway
        if self.gateway != "none":
            config["gateway"] = {"provider": self.gateway}

        # Observability
        if self.observability != "none":
            config["observability"] = {"stack": self.observability}

        # MLOps
        if self.enable_mlops:
            config["mlops"] = {
                "models": [],
                "serving": self.ml_serving,
                "vector_dbs": self.ml_vector_dbs,
                "monitoring": self.ml_monitoring,
                "tracking": self.ml_tracking,
                "orchestration": self.ml_orchestration,
            }

        return config

    @classmethod
    def from_domain_model(cls, models: dict[str, Any], project_name: str = "app") -> WizardSession:
        """Generate a starting WizardSession heuristically from a domain model."""
        session = cls(name=project_name, models=models)
        
        # Heuristics based on field types
        for model_name, model_attrs in models.items():
            # Handle both dicts and DataModelConfig objects
            fields = model_attrs.get("fields", {}) if isinstance(model_attrs, dict) else getattr(model_attrs, "fields", {})
            for field_name, field_data in fields.items():
                field_type = field_data.get("type", "") if isinstance(field_data, dict) else getattr(field_data, "type", "")
                if field_type == "Vector":
                    session.enable_mlops = True
                    if "qdrant" not in session.ml_vector_dbs:
                        session.ml_vector_dbs.append("qdrant")
                    if "ollama" not in session.ml_serving:
                        session.ml_serving.append("ollama")
        
        # Provide base defaults
        if not session.databases:
            session.databases.append("postgres")
        if not session.features:
            session.features.append("auth")
            
        return session

def migrate_session(raw_data: dict, from_version: int) -> dict:
    """Migrate an older session dictionary to the current schema version."""
    # Example migration logic for future use
    # if from_version < 2:
    #     raw_data["new_field"] = "default_value"
    
    # Pydantic will ignore extra fields and supply defaults for missing fields automatically 
    # if we just pass the dict directly to the constructor since default= is specified.
    # We update the version to avoid continuous migration loops.
    raw_data["schema_version"] = CURRENT_SCHEMA_VERSION
    return raw_data

def load_template(path: Path) -> WizardSession | None:
    """Load a template defensively, skipping invalid ones."""
    if not path.exists():
        return None
        
    try:
        with open(path) as f:
            data = json.load(f)
            
        version = data.get("schema_version", 1)
        if version < CURRENT_SCHEMA_VERSION:
            data = migrate_session(data, version)
            
        return WizardSession(**data)
    except Exception as e:
        # We don't crash, we just return nothing
        return None
