import glob
import json
import re
from abc import ABC, abstractmethod
from pathlib import Path

import pandas as pd
import requests
from urllib3.exceptions import ProtocolError

from .config import settings
from .helpers import (
    parse_json_to_dfs,
    parse_trial_data,
    unnest_trial_level_data,
    verify_dataframe_parsing,
)
from .log import log_error, log_info
from .utils import get_filename_timestamp


# === Importer Classes ===
class BaseImporter(ABC):
    def __init__(self, importer_name: str):
        self.importer_name = importer_name
        self.source_name = importer_name

    @abstractmethod
    def load(self, source_path: str):
        pass

    def _process(self, df: pd.DataFrame, activity_name_col: str):
        grouped = parse_json_to_dfs(df, activity_name_col)
        validation, activity_names = verify_dataframe_parsing(
            df, grouped, activity_name_col
        )
        return df, grouped, validation, activity_names, get_filename_timestamp()


class MetricWireImporter(BaseImporter):
    def __init__(self):
        super().__init__("metricwire")

    def load(self, filepath: str = "metricwire/data/unzipped/*/*/*.json"):
        json_files = glob.glob(filepath)
        log_info("Found MetricWire JSON files", {"count": len(json_files)})

        data = [json.load(open(fp)) for fp in json_files]
        flattened = []

        for record_list in data:
            for record in record_list:
                # Extract top-level fields (excluding 'data')
                top_level = {k: v for k, v in record.items() if k != "data"}

                # Process each trial in the data array
                for entry in record.get("data", []):
                    trial = {**top_level, **entry}

                    # Handle new MetricWire format: extract activityId from MW_META_DATA if needed
                    if "activityId" not in trial and "MW_META_DATA" in entry:
                        mw_meta = entry.get("MW_META_DATA", {})
                        if "activityId" in mw_meta:
                            trial["activityId"] = mw_meta["activityId"]

                    flattened.append(trial)

        return self._process(pd.DataFrame(flattened), activity_name_col="activityName")


class MongoDBImporter(BaseImporter):
    def __init__(self):
        super().__init__("mongodb")

    def load(self, source_path: str):
        df_flat = pd.read_json(source_path)
        return self._process(df_flat, activity_name_col="activity_name")

    def _process(self, df: pd.DataFrame, activity_name_col: str):
        grouped = parse_json_to_dfs(df, activity_name_col)
        # for each grouped dataset, apply unnest
        for name, group in grouped.items():
            grouped[name] = unnest_trial_level_data(group, drop_duplicates=True)
        validation, activity_names = verify_dataframe_parsing(
            df, grouped, activity_name_col
        )
        return df, grouped, validation, activity_names, get_filename_timestamp()


class M2C2KitAPIImporter(BaseImporter):
    def __init__(self):
        super().__init__("m2c2kit-api")

    def load(
        self,
        source_path: str,
        study_id: str,
        username: str,
        password: str,
        payload: dict = None,
        pipeline_name: str = "all_uploaded_data",
    ):
        """
        Load data from M2C2 API with authentication and streaming NDJSON response.

        Args:
            source_path: API base URL (e.g., "https://api.m2c2kit.com")
            study_id: Study identifier
            username: API username
            password: API password
            payload: Optional filters (e.g., {"skip": 1, "limit": 2, "activity_ids": [...], "participant_ids": [...]})
            pipeline_name: Pipeline endpoint name. Options:
                - "all_uploaded_data" (default): Raw data without aggregation
                - "assessment_data": Grouped sessions (may have streaming issues)
                - "participant_summaries": Participant summaries

        Returns:
            Tuple of (df, grouped, validation, activity_names, timestamp)

        Note:
            The "assessment_data" pipeline may experience streaming timeouts on some servers.
            Use "all_uploaded_data" for most reliable results.
        """
        # Step 1: Authenticate and get access token
        log_info(
            "Authenticating with M2C2 API", {"username": username, "study_id": study_id}
        )
        token = self._get_access_token(source_path, username, password)

        # Step 2: Submit pipeline job and stream NDJSON response
        log_info(
            "Submitting pipeline job", {"study_id": study_id, "pipeline": pipeline_name}
        )
        df_flat = self._stream_pipeline_data(
            source_path, study_id, pipeline_name, token, payload or {}
        )

        # Step 3: Process the data
        log_info("Processing streamed data", {"rows": len(df_flat)})
        return self._process(df_flat, activity_name_col="activity_name")

    def _get_access_token(self, base_url: str, username: str, password: str) -> str:
        """
        Authenticate with the API and retrieve an access token.

        Args:
            base_url: API base URL
            username: API username
            password: API password

        Returns:
            Access token string
        """
        auth_url = f"{base_url}/auth/token"

        try:
            response = requests.post(
                auth_url,
                data={"username": username, "password": password},
                headers={"Content-Type": "application/x-www-form-urlencoded"},
            )
            response.raise_for_status()

            token_data = response.json()
            access_token = token_data.get("access_token")

            if not access_token:
                raise ValueError("No access_token in authentication response")

            log_info("Authentication successful", {"username": username})
            return access_token

        except requests.RequestException as e:
            log_error("Authentication failed", {"error": str(e), "username": username})
            raise ValueError(f"Failed to authenticate: {e}")

    def _stream_pipeline_data(
        self,
        base_url: str,
        study_id: str,
        pipeline_name: str,
        access_token: str,
        payload: dict,
    ) -> pd.DataFrame:
        """
        Submit pipeline job and parse streaming NDJSON response.

        Args:
            base_url: API base URL
            study_id: Study identifier
            pipeline_name: Pipeline endpoint name
            access_token: Bearer token for authentication
            payload: Request filters/parameters

        Returns:
            DataFrame with parsed NDJSON data
        """
        pipeline_url = f"{base_url}/studies/study/{study_id}/pipelines/{pipeline_name}/submit?stream_response=true"

        headers = {
            "Authorization": f"Bearer {access_token}",
            "Content-Type": "application/json",
        }

        session = requests.Session()
        records = []

        try:
            with session.post(
                pipeline_url,
                headers=headers,
                json=payload,
                stream=True,
                timeout=300,  # 5 minute timeout
            ) as response:
                response.raise_for_status()

                line_count = 0

                try:
                    # Parse NDJSON
                    for line_bytes in response.iter_lines(
                        chunk_size=1024, decode_unicode=False
                    ):
                        if line_bytes:  # Skip empty lines
                            line_count += 1
                            # Manually decoding bytes to strings
                            line = line_bytes.decode("utf-8", errors="ignore").strip()

                            try:
                                record = json.loads(line)
                                records.append(record)
                            except json.JSONDecodeError as e:
                                log_error(
                                    "Failed to parse JSON line",
                                    {"error": str(e), "line_preview": line[:200]},
                                )
                                continue

                except (requests.exceptions.ChunkedEncodingError, ProtocolError) as e:
                    # If we got some records before the error, use them
                    if records:
                        log_error(
                            "Streaming error but partial data received",
                            {"error": str(e), "records_count": len(records)},
                        )
                    else:
                        # No data received
                        log_error(
                            "API streaming failed with no data",
                            {
                                "error": str(e),
                                "study_id": study_id,
                                "pipeline": pipeline_name,
                                "payload": payload,
                            },
                        )
                        raise ValueError(
                            f"M2C2 API streaming error for study '{study_id}'. "
                            f"Stream closed without data."
                            f"Error: {e}"
                        )

            if not records:
                log_error(
                    "No data returned from API",
                    {"study_id": study_id, "pipeline": pipeline_name},
                )
                raise ValueError("No data returned from API pipeline")

            log_info("Successfully streamed data", {"records": len(records)})
            return pd.DataFrame(records)

        except requests.RequestException as e:
            # Only raise if we don't have any records
            if not records:
                log_error(
                    "Pipeline request failed", {"error": str(e), "study_id": study_id}
                )
                raise ValueError(f"Failed to retrieve pipeline data: {e}")
            else:
                log_error(
                    "Request error but partial data received",
                    {"error": str(e), "records_count": len(records)},
                )
                return pd.DataFrame(records)
        finally:
            session.close()

    def _process(self, df: pd.DataFrame, activity_name_col: str):
        """
        Process the flat dataframe into grouped datasets by activity.

        Args:
            df: Flat dataframe with all records
            activity_name_col: Column name containing activity identifiers

        Returns:
            Tuple of (df, grouped, validation, activity_names, timestamp)
        """
        grouped = parse_json_to_dfs(df, activity_name_col)

        # For each grouped dataset, apply unnest to flatten trial-level data
        for name, group in grouped.items():
            grouped[name] = unnest_trial_level_data(group, drop_duplicates=True)

        validation, activity_names = verify_dataframe_parsing(
            df, grouped, activity_name_col
        )

        return df, grouped, validation, activity_names, get_filename_timestamp()


class UASImporter(BaseImporter):
    def __init__(self):
        super().__init__("uas")

    def load(self, url: str):
        try:
            response = requests.get(url)
            response.raise_for_status()
        except requests.RequestException as e:
            log_error("Failed to fetch UAS data", {"error": str(e)})
            return pd.DataFrame(), {}, False, []

        lines = [
            line.strip().rstrip(",")
            for line in response.text.splitlines()
            if line.strip()
        ]
        parsed = [json.loads(line) for line in lines if line]
        df = pd.DataFrame(parsed)

        if "data" not in df.columns:
            raise ValueError("Expected 'data' column not found in UAS export.")

        expanded = df["data"].apply(pd.Series)
        full_df = pd.concat([df.drop(columns=["data"]), expanded], axis=1)
        return self._process(full_df, activity_name_col="taskname")


class QualtricsImporter(BaseImporter):
    def __init__(self):
        super().__init__("qualtrics")

    def load(self, source_path: str):
        df = pd.read_csv(source_path, header=1, skiprows=[2])
        df = df.rename(columns={"Response ID": "ResponseId"}).reset_index()
        return self._process(df, activity_name_col="activity_name")

    def _process(self, df: pd.DataFrame, activity_name_col: str):
        # === Step 2: Identify M2C2 Trial Data Columns ===
        pattern = re.compile(r"(M2C2_ASSESSMENT_\d+)_TRIAL_DATA_(\d+)")
        # settings.QUALTRICS_TRIAL_DATA_REGEX
        trial_columns = [col for col in df.columns if pattern.match(col)]

        # Group columns by assessment
        assessment_map = {}
        for col in trial_columns:
            match = pattern.match(col)
            if match:
                assessment_name, _ = match.groups()
                assessment_map.setdefault(assessment_name, []).append(col)

        # === Step 4: Long-Format Parsing Across All Assessments ===
        long_format_data = []

        for assessment_name, columns in assessment_map.items():
            melted = df.melt(
                id_vars=[
                    "index",
                    "ResponseId",
                    "M2C2_ASSESSMENT_ORDER",
                    "M2C2_AUTO_ADVANCE",
                    "M2C2_LANGUAGE",
                ],
                value_vars=columns,
                var_name="trial_label",
                value_name="trial_data_json",
            )

            melted["assessment"] = assessment_name
            melted["trial_number"] = melted["trial_label"].str.extract(
                r"TRIAL_DATA_(\d+)"
            )
            melted["trial_number_int"] = pd.to_numeric(
                melted["trial_number"], errors="coerce"
            ).astype("Int64")

            melted = melted.dropna(subset=["trial_data_json"])

            parsed = melted["trial_data_json"].apply(parse_trial_data)
            expanded = pd.json_normalize(parsed)

            combined = pd.concat(
                [
                    melted.drop(columns=["trial_data_json"]).reset_index(drop=True),
                    expanded.reset_index(drop=True),
                ],
                axis=1,
            )

            long_format_data.append(combined)

        # === Step 5: Final Combined Long Format Dataset ===
        final_df = pd.concat(long_format_data, ignore_index=True)
        final_df = final_df.drop("index", axis=1)

        grouped_dataframes = parse_json_to_dfs(
            final_df, activity_name_col=activity_name_col
        )
        return (
            df,
            grouped_dataframes,
            True,
            assessment_map.keys(),
            get_filename_timestamp(),
        )


class MultiCSVImporter(BaseImporter):
    def __init__(self):
        super().__init__("multicsv")

    def load(self, source_map: dict, activity_name_col: str = "activity_name"):
        """
        Load and stack multiple CSV files tagged with activity names.

        Args:
            source_map (dict): Mapping from activity name to file path.
            key_name (str): Column to tag source of each row (e.g., "activity_name").

        Returns:
            pd.DataFrame: Combined dataframe with all sources stacked.
        """
        dfs = []
        for activity, path in source_map.items():
            df = pd.read_csv(Path(path))
            df[activity_name_col] = activity
            dfs.append(df)

        combined_df = pd.concat(dfs, ignore_index=True)
        return self._process(combined_df, activity_name_col=activity_name_col)


class DataImporter:
    SOURCES = {
        "metricwire": MetricWireImporter,
        "mongodb": MongoDBImporter,
        "m2c2kit-api": M2C2KitAPIImporter,
        "uas": UASImporter,
        "multicsv": MultiCSVImporter,
        "qualtrics": QualtricsImporter,
    }

    @staticmethod
    def load_from(source_name: str, source_path: str):
        name = source_name.lower()
        if name not in DataImporter.SOURCES:
            log_error("Unsupported source", {"source_name": source_name})
            raise ValueError(f"Unsupported source: '{source_name}'")

        log_info("Initializing data load", {"source": source_name, "path": source_path})
        return DataImporter.SOURCES[name]().load(source_path)
