"""
Model architecture construction, pipeline assembly, and training orchestration.
Integrates text vectorization, character n-grams, and cyber threat indicators.
"""

import time
import joblib
from typing import Dict, Any, Tuple
from sklearn.pipeline import Pipeline, FeatureUnion
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import StandardScaler
from sklearn.dummy import DummyClassifier
from sklearn.linear_model import LogisticRegression
from xgboost import XGBClassifier

from src.feature_engineering import ThreatFeatureExtractor, EmailTextCombiner


def build_pipeline(model_type: str = "logistic_regression", params: Dict[str, Any] = None) -> Pipeline:
    """
    Constructs a complete end-to-end scikit-learn pipeline combining:
    1. Word-level TF-IDF (1, 2 n-grams)
    2. Character-level TF-IDF (3, 5 n-grams for obfuscation detection)
    3. Structural & behavioral threat features (urgency, credential, url counts)
    4. Selected classification algorithm
    """
    if params is None:
        params = {}

    word_max_features = params.get("word_max_features", 4000)
    char_max_features = params.get("char_max_features", 2000)

    # Word-level TF-IDF pipeline
    word_tfidf_pipeline = Pipeline([
        ("combiner", EmailTextCombiner(subject_weight=2)),
        ("word_tfidf", TfidfVectorizer(
            ngram_range=(1, 2),
            max_features=word_max_features,
            sublinear_tf=True,
            strip_accents="unicode"
        ))
    ])

    # Character-level TF-IDF pipeline for lexical obfuscation/typosquatting
    char_tfidf_pipeline = Pipeline([
        ("combiner", EmailTextCombiner(subject_weight=1)),
        ("char_tfidf", TfidfVectorizer(
            analyzer="char",
            ngram_range=(3, 5),
            max_features=char_max_features,
            sublinear_tf=True
        ))
    ])

    # Cyber domain behavioral feature extraction
    threat_features_pipeline = Pipeline([
        ("extractor", ThreatFeatureExtractor()),
        ("scaler", StandardScaler())
    ])

    # Unified feature representation
    feature_union = FeatureUnion([
        ("word_features", word_tfidf_pipeline),
        ("char_features", char_tfidf_pipeline),
        ("threat_indicators", threat_features_pipeline)
    ])

    # Instantiate chosen classifier
    if model_type == "dummy":
        clf = DummyClassifier(
            strategy=params.get("strategy", "stratified"),
            random_state=params.get("random_state", 42)
        )
    elif model_type == "logistic_regression":
        clf = LogisticRegression(
            C=params.get("C", 1.0),
            max_iter=params.get("max_iter", 1000),
            class_weight=params.get("class_weight", "balanced"),
            solver="lbfgs",
            random_state=params.get("random_state", 42)
        )
    elif model_type == "logistic_regression_security_tuned":
        # Tuned with higher penalty weighting on Suspicious (Class 2) for zero-trust defense
        custom_weights = params.get("class_weight", {0: 1.0, 1: 1.0, 2: 1.4})
        clf = LogisticRegression(
            C=params.get("C", 2.5),
            max_iter=params.get("max_iter", 1500),
            class_weight=custom_weights,
            solver="lbfgs",
            random_state=params.get("random_state", 42)
        )
    elif model_type == "xgboost":
        clf = XGBClassifier(
            n_estimators=params.get("n_estimators", 100),
            max_depth=params.get("max_depth", 5),
            learning_rate=params.get("learning_rate", 0.1),
            subsample=params.get("subsample", 0.85),
            colsample_bytree=params.get("colsample_bytree", 0.85),
            eval_metric="mlogloss",
            random_state=params.get("random_state", 42),
            n_jobs=4
        )
    else:
        raise ValueError(f"Unsupported model_type: '{model_type}'")

    pipeline = Pipeline([
        ("features", feature_union),
        ("classifier", clf)
    ])

    return pipeline


def train_model(
    model_type: str,
    train_df,
    params: Dict[str, Any] = None
) -> Tuple[Pipeline, float]:
    """
    Fits the complete pipeline and measures training execution latency.
    """
    pipeline = build_pipeline(model_type, params)
    
    start_time = time.time()
    pipeline.fit(train_df, train_df["label_id"].values)
    duration = time.time() - start_time

    return pipeline, duration


def save_pipeline(pipeline: Pipeline, output_path: str) -> str:
    """
    Serializes complete fitted pipeline to file using joblib.
    """
    import os
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    joblib.dump(pipeline, output_path)
    return output_path


def load_pipeline(pipeline_path: str) -> Pipeline:
    """
    Loads saved pipeline artifact.
    """
    return joblib.load(pipeline_path)