from typing import Dict, Any, Optional, List, Literal
from pydantic import BaseModel, model_validator
from enum import Enum


class HuggingFaceConfig(BaseModel):
    repo_id: str


class StorageConfig(BaseModel):
    storage_ref: str
    storage_path: str


class FoundationModel(BaseModel):
    name: str
    source: Literal["huggingface", "storage"]
    huggingface: Optional[HuggingFaceConfig] = None
    storage: Optional[StorageConfig] = None


class BaseModelConfig(BaseModel):
    foundation_model: FoundationModel
    lora_adapter: Optional[StorageConfig] = None
    quantization: Optional[str] = None


class TrainingTechnique(str, Enum):
    SFT = "SFT"
    DPO = "DPO"
    GRPO = "GRPO"
    PPO = "PPO"


class LoRAConfig(BaseModel):
    lora_alpha: Optional[int] = 64
    lora_dropout: Optional[float] = 0.05
    r: Optional[float] = 16
    bias: Optional[str] = "none"
    task_type: Optional[str] = "CAUSAL_LM"
    target_modules: Optional[List[str]] = [
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj"
    ]


class TrainingOutput(BaseModel):
    storage_ref: str
    base_path: str


class TrainConfig(BaseModel):
    training_technique: TrainingTechnique
    task_type: str
    training_method: str = "LoRA"
    terminate_on_failure: Optional[bool] = True
    grpo_configs: Optional[Dict[str, Any]] = None
    hyper_parameters: Optional[Dict[str, Any]] = None
    lora_config: Optional[LoRAConfig] = None
    training_output: Optional[TrainingOutput] = None