import logging, warnings
from pydantic import BaseModel, Field, model_validator
from typing import List, Literal, Optional, Dict, Any, Union
from pathlib import Path
import polars as pl
from .hf_utils import get_dataset_splits, get_dataset_configs


class ModelConfig(BaseModel):
    repo_id: str
    base_model_dir: Path = Path("models/")


class PeftConfig(BaseModel):
    method: Literal["qlora", "lora", "dora"] = "qlora"
    r: int = 8
    lora_alpha: int = 32
    lora_dropout: float = 0.05


class TabularConfig(BaseModel):
    categorical_columns: List[str]
    numerical_columns: List[str]
    projection_dim: int
    categorical_cardinality: Optional[Dict[str, int]] = None


class VisionConfig(BaseModel):
    image_column: str
    model_name: str = "vit_base_patch16_224"
    projection_dim: int


class AudioConfig(BaseModel):
    audio_column: str
    model_name: str = "openai/whisper-tiny" # Default lightweight model
    projection_dim: int


class DatasetConfig(BaseModel):
    repo_id: Optional[str] = None
    file_path: Optional[Path] = None
    config_name: Optional[str] = None
    subset: Optional[str] = None
    db_uri: Optional[str] = None
    db_query: Optional[str] = None
    split: str = "train"
    instruction_column: str = "instruction"
    input_column: str = "input"
    output_column: str = "output"
    text_columns: List[str] = Field(default_factory=list)
    image_root_path: Optional[Path] = None
    
    @model_validator(mode='after')
    def check_source(self):
        if not self.file_path and not self.repo_id and not self.db_uri:
            raise ValueError("Must provide either 'file_path', 'repo_id', or 'db_uri'.")
        if self.db_uri and not self.db_query:
            raise ValueError("'db_query' must be provided when using 'db_uri'.")
        return self

class DataConfig(BaseModel):
    # Deprecated single-dataset fields, kept for backward compatibility helper
    file_path: Optional[Path] = None
    dataset_repo_id: Optional[str] = None
    dataset_config_name: Optional[str] = None
    subset: Optional[str] = None
    db_uri: Optional[str] = None
    db_query: Optional[str] = None
    split: str = "train"
    
    # New Multi-Dataset support
    datasets: List[DatasetConfig] = Field(default_factory=list)

    # Global config
    instruction_column: str = "instruction"
    input_column: str = "input"
    output_column: str = "output"
    text_columns: List[str] = Field(default_factory=list)
    image_root_path: Optional[Path] = None
    max_seq_length: int = 512
    tabular_config: Optional[TabularConfig] = None
    vision_config: Optional[VisionConfig] = None
    audio_config: Optional[AudioConfig] = None

    @model_validator(mode='after')
    def consolidate_datasets(self):
        # If legacy fields are present and datasets list is empty, convert to a DatasetConfig
        if not self.datasets:
             if self.file_path or self.dataset_repo_id:
                 self.datasets.append(DatasetConfig(
                     repo_id=self.dataset_repo_id,
                     file_path=self.file_path,
                     config_name=self.dataset_config_name,
                     subset=self.subset,
                     db_uri=self.db_uri,
                     db_query=self.db_query,
                     split=self.split,
                     instruction_column=self.instruction_column,
                     input_column=self.input_column,
                     output_column=self.output_column,
                     text_columns=self.text_columns,
                     image_root_path=self.image_root_path
                 ))
             else:
                 pass # Might be valid if we are initializing empty? Or error.
        return self


class EvaluationConfig(BaseModel):
    do_eval: bool = True
    eval_dataset_size: float = 0.1


from .utils import get_auto_device

class TrainerConfig(BaseModel):
    max_epochs: int = 3
    devices: int = 1
    device: str = Field(default_factory=get_auto_device)
    logger: Literal["csv", "tensorboard"] = "tensorboard"
    checkpoint_callback: bool = True
    evaluation: EvaluationConfig = Field(default_factory=EvaluationConfig)
    limit_train_batches: Union[int, float] = 1.0
    limit_val_batches: Union[int, float] = 1.0
    accelerator: str = "auto"


class TrainConfig(BaseModel):
    peft: PeftConfig = Field(default_factory=PeftConfig)
    llm_lr: float = 5e-5
    tower_lr: float = 1e-4
    batch_size: int = 2
    push_to_hub: bool = False
    hub_model_id: Optional[str] = None


class DeploymentConfig(BaseModel):
    port: int = 8000
    use_vllm: bool = True


class PipelineConfig(BaseModel):
    model: ModelConfig
    data: DataConfig
    trainer: TrainerConfig = Field(default_factory=TrainerConfig)
    train: TrainConfig = Field(default_factory=TrainConfig)
    deployment: DeploymentConfig = Field(default_factory=DeploymentConfig)

    @property
    def is_multimodal(self) -> bool:
        return (
            self.data.vision_config is not None 
            or self.data.tabular_config is not None
            or self.data.audio_config is not None
        )

    def get_output_dir(self) -> Path:
        base = self.model.repo_id.replace("/", "_")
        mode = "multimodal" if self.is_multimodal else self.train.peft.method
        return Path("finetuning_runs") / f"{base}_{mode}"

    @staticmethod
    def _analyze_schema(df_sample: pl.DataFrame) -> Dict:
        analysis = {
            "image_column": None,
            "audio_column": None,
            "text_columns": [],
            "numerical_columns": [],
            "categorical_columns": [],
            "output_column": None,
        }
        potential_output_cols = ("output", "target", "label", "completion", "review")
        unstructured_cols = []

        for col, dtype in df_sample.schema.items():
            l_col = col.lower()
            if ("image" in l_col or "path" in l_col) and dtype == pl.Utf8:
                # Naive check, improved below
                analysis["image_column"] = col
            elif ("audio" in l_col or "sound" in l_col) and dtype == pl.Utf8:
                analysis["audio_column"] = col
            elif dtype in [pl.Float32, pl.Float64, pl.Int32, pl.Int64]:
                analysis["numerical_columns"].append(col)
            elif dtype == pl.Utf8:
                unstructured_cols.append(col)

        # Prioritize finding the output column among all unstructured columns
        for col_name in potential_output_cols:
            if col_name in unstructured_cols:
                analysis["output_column"] = col_name
                unstructured_cols.remove(col_name)
                break

        # Classify remaining unstructured columns
        common_text_cols = ("text", "content", "description", "comment", "message", "body")
        for col in unstructured_cols:
            l_col = col.lower()
            avg_str_len = df_sample[col].str.len_chars().mean()
            # If low cardinality and short strings, it's likely categorical
            if any(c in l_col for c in common_text_cols):
                analysis["text_columns"].append(col)
            elif df_sample[col].n_unique() < 50 and avg_str_len < 30:
                analysis["categorical_columns"].append(col)
            else:
                analysis["text_columns"].append(col)

        # Fallback for output column if not explicitly found
        if not analysis["output_column"]:
            if analysis["text_columns"]:
                # Prefer the text column with the longest average string length as output
                longest_text_col = max(
                    analysis["text_columns"],
                    key=lambda c: df_sample[c].str.len_chars().mean(),
                )
                analysis["output_column"] = longest_text_col
                analysis["text_columns"].remove(longest_text_col)
                warnings.warn(
                    f"Could not automatically determine the output column. Using '{analysis['output_column']}' as the target."
                )
            elif analysis["categorical_columns"]:
                # As a last resort, use a categorical column
                analysis["output_column"] = analysis["categorical_columns"].pop()
                warnings.warn(
                    f"Could not find a clear text output column. Using categorical column '{analysis['output_column']}' as the target."
                )

        return analysis

    @staticmethod
    def _suggest_hyperparameters(dataset_size: int) -> Dict:
        if dataset_size < 10000:
            return {"llm_lr": 5e-5, "epochs": 3, "r": 8}
        elif dataset_size < 100000:
            return {"llm_lr": 3e-5, "epochs": 2, "r": 16}
        else:
            return {"llm_lr": 2e-5, "epochs": 1, "r": 32}

    @staticmethod
    def _read_and_sample_dataset(file_path: Union[Path, str]) -> pl.DataFrame:
        file_path = Path(file_path)
        file_suffix = file_path.suffix.lower()
        if file_suffix == ".csv":
            df = pl.read_csv(file_path)
        elif file_suffix in (".json", ".jsonl"):
            df = pl.read_json(file_path)
        else:
            raise ValueError(f"Unsupported file type: {file_suffix}")
        return df, df.head(100)

    def _get_llm_hf_config(model_repo_id: str):
        from transformers import AutoConfig
        return AutoConfig.from_pretrained(model_repo_id, trust_remote_code=True)

    @classmethod
    def from_yaml(cls, path: str) -> "PipelineConfig":
        import yaml
        with open(path, "r") as f:
            config = yaml.safe_load(f)
        return cls(**config)

    @classmethod
    def from_dataset(
        cls,
        model_repo_id: str,
        datasets: Optional[List[Dict[str, Any]]] = None,
        file_path: Optional[Path] = None,
        dataset_repo_id: Optional[str] = None,
        dataset_config_name: Optional[str] = None,
        split: Optional[str] = None,
        **kwargs
    ) -> Union["PipelineConfig", Dict[str, Any]]:
        logging.info(f"🔬 Analyzing dataset(s) to create smart configuration...")

        token = kwargs.get("token")
        
        # Consolidate inputs into a list of "sources"
        sources = []
        if datasets:
            sources.extend(datasets)
        if file_path or dataset_repo_id:
            sources.append({
                "file_path": file_path,
                "repo_id": dataset_repo_id,
                "config_name": dataset_config_name,
                "split": split,
                "subset": kwargs.get("subset")
            })

        if not sources:
             raise ValueError("Must provide at least one dataset (file_path, dataset_repo_id, or datasets list).")

        # We will analyze the FIRST source to determine schema/hyperparams for now,
        # or we could try to error if schemas mismatch.
        
        # Let's analyze the first valid source to get the global schema
        first_source = sources[0]
        
        def get_sample(src):
            fpath = src.get("file_path")
            drepo = src.get("repo_id") or src.get("dataset_repo_id")
            dconf = src.get("config_name")
            dsplit = src.get("split")
            db_uri = src.get("db_uri")
            db_query = src.get("db_query")
            
            if fpath:
                try:
                    df, df_sample = cls._read_and_sample_dataset(Path(fpath))
                    return df_sample, df.height
                except Exception as e:
                     raise ValueError(f"Failed to read local file {fpath}: {e}")
            elif drepo:
                 return cls._sample_hf_dataset(drepo, dconf, dsplit, token)
            elif db_uri and db_query:
                 try:
                     import polars as pl
                     sample_query = db_query
                     if "limit" not in sample_query.lower():
                         sample_query = f"SELECT * FROM ({db_query}) AS subquery LIMIT 100"
                     df_sample = pl.read_database_uri(query=sample_query, uri=db_uri, engine="connectorx")
                     return df_sample, 10000 
                 except Exception as e:
                     raise ValueError(f"Failed to read from database: {e}")
            else:
                 raise ValueError("Source must have file_path, repo_id, or db_uri with db_query")

        df_sample, total_size = get_sample(first_source)
        
        is_alpaca_format = all(c in df_sample.columns for c in ["instruction", "input", "output"])

        if is_alpaca_format:
            schema_analysis = {
                "image_column": None,
                "text_columns": [],
                "numerical_columns": [],
                "categorical_columns": [],
                "output_column": "output",
                "instruction_column": "instruction",
                "input_column": "input",
            }
        else:
            schema_analysis = cls._analyze_schema(df_sample)

        # Create DatasetConfig objects
        dataset_configs = []
        for src in sources:
            d_cfg = DatasetConfig(
                repo_id=src.get("repo_id") or src.get("dataset_repo_id"),
                file_path=src.get("file_path"),
                config_name=src.get("config_name"),
                split=src.get("split") or "train",
                subset=src.get("subset"),
                instruction_column=schema_analysis.get("instruction_column", "instruction"),
                input_column=schema_analysis.get("input_column", "input"),
                output_column=schema_analysis.get("output_column", "output"),
                text_columns=schema_analysis.get("text_columns", []),
                image_root_path=kwargs.get("image_root_path") 
            )
            dataset_configs.append(d_cfg)

        hyperparams = cls._suggest_hyperparameters(total_size)

        llm_hf_config = cls._get_llm_hf_config(model_repo_id)

        data_cfg = DataConfig(
            datasets=dataset_configs,
            # Global settings
            text_columns=schema_analysis["text_columns"],
            output_column=schema_analysis["output_column"],
            image_root_path=kwargs.pop("image_root_path", None),
            **{
                k: v
                for k, v in schema_analysis.items()
                if k in ("instruction_column", "input_column")
            },
        )

        if schema_analysis.get("image_column"):
            data_cfg.vision_config = VisionConfig(
                image_column=schema_analysis["image_column"],
                projection_dim=llm_hf_config.hidden_size,
            )

        if schema_analysis.get("audio_column"):
            data_cfg.audio_config = AudioConfig(
                audio_column=schema_analysis["audio_column"],
                projection_dim=llm_hf_config.hidden_size,
            )
        if (
            schema_analysis["numerical_columns"]
            or schema_analysis["categorical_columns"]
        ):
            data_cfg.tabular_config = TabularConfig(
                numerical_columns=schema_analysis["numerical_columns"],
                categorical_columns=schema_analysis["categorical_columns"],
                projection_dim=llm_hf_config.hidden_size,
            )

        train_cfg = TrainConfig(
            llm_lr=hyperparams["llm_lr"], peft=PeftConfig(r=hyperparams["r"])
        )
        trainer_cfg = TrainerConfig(max_epochs=hyperparams["epochs"])
        
        logging.info(
            f"✅ Analysis complete. Smart Hyperparameters: "
            f"LR={hyperparams['llm_lr']}, Epochs={hyperparams['epochs']}, LoRA Rank={hyperparams['r']}"
        )
        return cls(
            model=ModelConfig(repo_id=model_repo_id),
            data=data_cfg,
            train=train_cfg,
            trainer=trainer_cfg,
            **kwargs,
        )

    @staticmethod
    def _sample_hf_dataset(repo_id, config_name, split, token):
         from datasets import load_dataset
         try:
             ds = load_dataset(repo_id, name=config_name, split=split or "train", streaming=True, token=token)
             sample_data = []
             for i, row in enumerate(ds):
                 if i >= 100: break
                 sample_data.append(row)
             return pl.from_dicts(sample_data), 10000 
         except Exception as e:
             raise ValueError(f"Failed to stream HF dataset {repo_id}: {e}")
