import argparse
import json
import logging
import re
import sys
from datetime import datetime
from getpass import getpass
from pathlib import Path
from typing import Dict, Any

import pandas as pd

from amplify_client import AmplifyClient
from model_field_parser import ModelFieldParser

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

CONFIG_DIR = Path.home() / ".amplify-migrator"
CONFIG_FILE = CONFIG_DIR / "config.json"


class ExcelToAmplifyMigrator:
    def __init__(self, excel_file_path: str):
        self._current_sheet = None
        self.failed_records_by_sheet = {}
        self.model_field_parser = ModelFieldParser()
        self.excel_file_path = excel_file_path
        self.amplify_client = None

    def init_client(
        self,
        api_endpoint: str,
        region: str,
        user_pool_id: str,
        is_aws_admin: bool = False,
        client_id: str = None,
        username: str = None,
        aws_profile: str = None,
    ):

        self.amplify_client = AmplifyClient(
            api_endpoint=api_endpoint,
            user_pool_id=user_pool_id,
            region=region,
            client_id=client_id,
        )

        try:
            self.amplify_client.init_cognito_client(
                is_aws_admin=is_aws_admin, username=username, aws_profile=aws_profile
            )

        except RuntimeError or Exception:
            sys.exit(1)

    def authenticate(self, username: str, password: str) -> bool:
        return self.amplify_client.authenticate(username, password)

    def run(self):
        all_sheets = self.read_excel()

        total_success = 0

        for sheet_name, df in all_sheets.items():
            logger.info(f"Processing {sheet_name} sheet with {len(df)} rows")
            total_success += self.process_sheet(df, sheet_name)

        self._display_summary(len(all_sheets), total_success)

    def _display_summary(self, sheets_processed: int, total_success: int) -> None:
        total_failed = sum(len(failures) for failures in self.failed_records_by_sheet.values())

        print("\n" + "=" * 60)
        print("MIGRATION SUMMARY")
        print("=" * 60)
        print(f"📊 Sheets processed: {sheets_processed}")
        print(f"✅ Total successful: {total_success}")
        print(f"❌ Total failed: {total_failed}")
        print(
            f"📈 Success rate: {(total_success / (total_success + total_failed) * 100):.1f}%"
            if (total_success + total_failed) > 0
            else "📈 Success rate: N/A"
        )

        if total_failed > 0:
            print("\n" + "=" * 60)
            print("FAILED RECORDS DETAILS")
            print("=" * 60)

            for sheet_name, failed_records in self.failed_records_by_sheet.items():
                if not failed_records:
                    continue

                print(f"\n📄 {sheet_name}:")
                print("-" * 60)
                for record in failed_records:
                    primary_field_value = record.get("primary_field_value", "Unknown")
                    error = record.get("error", "Unknown error")
                    print(f"  • Record: {primary_field_value}")
                    print(f"    Error: {error}")

            print("\n" + "=" * 60)

            export_confirm = input("\nExport failed records to Excel? (yes/no): ")
            if export_confirm.lower() == "yes":
                failed_records_file = self._write_failed_records_to_excel()
                if failed_records_file:
                    print(f"📁 Failed records exported to: {failed_records_file}")
                    print("=" * 60)

                    update_config = input("\nUpdate config to use this failed records file for next run? (yes/no): ")
                    if update_config.lower() == "yes":
                        self._update_excel_path_in_config(failed_records_file)
                        print(f"✅ Config updated! Next 'migrate' will use: {Path(failed_records_file).name}")
                        print("=" * 60)
            else:
                print("Failed records export skipped.")
                print("=" * 60)
        else:
            print("\n✨ No failed records!")

        print("=" * 60)

    @staticmethod
    def _update_excel_path_in_config(new_excel_path: str) -> None:
        config = load_cached_config()
        config["excel_path"] = new_excel_path
        save_config(config)

    def _write_failed_records_to_excel(self) -> str | None:
        if not self.failed_records_by_sheet or all(
            len(failures) == 0 for failures in self.failed_records_by_sheet.values()
        ):
            return None

        input_path = Path(self.excel_file_path)

        base_name = input_path.stem
        if "_failed_records_" in base_name:
            base_name = base_name.split("_failed_records_")[0]

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_filename = f"{base_name}_failed_records_{timestamp}.xlsx"
        output_path = input_path.parent / output_filename

        logger.info(f"Writing failed records to {output_path}")

        with pd.ExcelWriter(output_path, engine="openpyxl") as writer:
            for sheet_name, failed_records in self.failed_records_by_sheet.items():
                if not failed_records:
                    continue

                rows_data = []
                for record in failed_records:
                    row_data = record.get("original_row", {}).copy()
                    row_data["ERROR"] = record["error"]
                    rows_data.append(row_data)

                df = pd.DataFrame(rows_data)

                df.to_excel(writer, sheet_name=sheet_name, index=False)

        logger.info(f"Successfully wrote failed records to {output_path}")
        return str(output_path)

    def read_excel(self) -> Dict[str, Any]:
        logger.info(f"Reading Excel file: {self.excel_file_path}")
        all_sheets = pd.read_excel(self.excel_file_path, sheet_name=None)
        logger.info(f"Loaded {len(all_sheets)} sheets from Excel")
        return all_sheets

    def process_sheet(self, df: pd.DataFrame, sheet_name: str) -> int:
        self._current_sheet = sheet_name
        self.failed_records_by_sheet[sheet_name] = []

        parsed_model_structure = self.get_parsed_model_structure(sheet_name)

        records, row_dict_by_primary = self.transform_rows_to_records(df, parsed_model_structure, sheet_name)

        confirm = input(f"\nUpload {len(records)} records of {sheet_name} to Amplify? (yes/no): ")
        if confirm.lower() != "yes":
            logger.info(f"Upload cancelled for {sheet_name} sheet")
            return 0

        success_count, upload_error_count, failed_uploads = self.amplify_client.upload(
            records, sheet_name, parsed_model_structure
        )

        for failed_upload in failed_uploads:
            primary_value = str(failed_upload["primary_field_value"])
            original_row = row_dict_by_primary.get(primary_value, {})

            self._record_failure(
                primary_field=failed_upload["primary_field"],
                primary_field_value=failed_upload["primary_field_value"],
                error=failed_upload["error"],
                original_row=original_row,
            )

        print(f"=== Upload of Excel sheet: {sheet_name} Complete ===")
        print(f"✅ Success: {success_count}")
        print(
            f"❌ Failed: {len(self.failed_records_by_sheet[sheet_name])} "
            f"(Parsing: {len(self.failed_records_by_sheet[sheet_name])}, Upload: {upload_error_count})"
        )
        print(f"📊 Total: {len(df)}")

        return success_count

    def transform_rows_to_records(
        self,
        df: pd.DataFrame,
        parsed_model_structure: Dict[str, Any],
        sheet_name: str,
    ) -> tuple[list[Any], Dict[str, Dict]]:
        records = []
        row_dict_by_primary = {}  # Maps primary_field_value to original row_dict
        row_count = 0
        df.columns = [self.to_camel_case(c) for c in df.columns]
        primary_field, _, _ = self.amplify_client.get_primary_field_name(sheet_name, parsed_model_structure)

        fk_lookup_cache = {}
        if self.amplify_client:
            logger.info("🚀 Pre-fetching foreign key lookups...")
            fk_lookup_cache = self.amplify_client.build_foreign_key_lookups(df, parsed_model_structure)

        for row_tuple in df.itertuples(index=False, name="Row"):
            row_count += 1
            row_dict = {col: getattr(row_tuple, col) for col in df.columns}
            primary_field_value = row_dict.get(primary_field, f"Row {row_count}")

            row_dict_by_primary[str(primary_field_value)] = row_dict.copy()

            try:
                record = self.transform_row_to_record(row_dict, parsed_model_structure, fk_lookup_cache)
                if record:
                    records.append(record)
            except Exception as e:
                error_msg = str(e)
                logger.error(f"Error transforming row {row_count} ({primary_field}={primary_field_value}): {error_msg}")
                self._record_failure(
                    primary_field=primary_field,
                    primary_field_value=primary_field_value,
                    error=f"Parsing error: {error_msg}",
                    original_row=row_dict,
                )

        logger.info(f"Prepared {len(records)} records for upload")

        return records, row_dict_by_primary

    def get_parsed_model_structure(self, sheet_name: str) -> Dict[str, Any]:
        model_structure = self.amplify_client.get_model_structure(sheet_name)
        return self.model_field_parser.parse_model_structure(model_structure)

    def _record_failure(
        self,
        primary_field: str,
        primary_field_value: str,
        error: str,
        original_row: Dict = None,
    ) -> None:
        if self._current_sheet not in self.failed_records_by_sheet:
            self.failed_records_by_sheet[self._current_sheet] = []

        failure_record = {
            "primary_field": primary_field,
            "primary_field_value": primary_field_value,
            "error": error,
        }

        if original_row is not None:
            failure_record["original_row"] = original_row

        self.failed_records_by_sheet[self._current_sheet].append(failure_record)

    def transform_row_to_record(
        self, row_dict: Dict, parsed_model_structure: Dict[str, Any], fk_lookup_cache: Dict[str, Dict[str, str]]
    ) -> dict[Any, Any] | None:
        """Transform a DataFrame row to Amplify model format"""

        model_record = {}

        for field in parsed_model_structure["fields"]:
            input = self.parse_input(row_dict, field, parsed_model_structure, fk_lookup_cache)
            if input:
                model_record[field["name"]] = input

        return model_record

    def parse_input(
        self,
        row_dict: Dict,
        field: Dict[str, Any],
        parsed_model_structure: Dict[str, Any],
        fk_lookup_cache: Dict[str, Dict[str, str]],
    ) -> Any | None:
        field_name = field["name"][:-2] if field["is_id"] else field["name"]

        if field_name not in row_dict or pd.isna(row_dict[field_name]):
            if field["is_required"]:
                raise ValueError(f"Required field '{field_name}' is missing")
            return None

        value = self.model_field_parser.clean_input(row_dict[field_name])

        if field["is_id"]:
            if "related_model" in field:
                related_model = field["related_model"]
            else:
                related_model = (temp := field["name"][:-2])[0].upper() + temp[1:]

            if related_model in fk_lookup_cache:
                lookup_dict = fk_lookup_cache[related_model]["lookup"]
                record_id = lookup_dict.get(str(value))

                if record_id:
                    return record_id
                elif field["is_required"]:
                    raise ValueError(f"{related_model}: {value} does not exist")
                return None
            else:
                logger.warning(f"No pre-fetched data for {related_model}, falling back to API call")
                record = self.amplify_client.get_record(
                    related_model, parsed_model_structure=parsed_model_structure, value=value
                )
                if record and record.get("id"):
                    return record["id"]
                elif field["is_required"]:
                    raise ValueError(f"{related_model}: {value} does not exist")
                return None
        elif field["is_list"] and field["is_scalar"]:
            return self.model_field_parser.parse_scalar_array(field, field_name, row_dict[field_name])
        else:
            return self.model_field_parser.parse_field_input(field, field_name, value)

    def _parse_custom_type_array(self, row: pd.Series, field: Dict[str, Any]) -> Any:
        field_name = field["name"]

        if field_name in row.index and pd.notna(row[field_name]):
            value = row[field_name]
            if isinstance(value, str) and value.strip().startswith(("[", "{")):
                try:
                    return json.loads(value)
                except json.JSONDecodeError:
                    logger.warning(f"Failed to parse JSON for '{field_name}', trying column-based parsing")

        custom_type_name = field["type"]
        parsed_custom_type = self.get_parsed_model_structure(custom_type_name)
        custom_type_fields = parsed_custom_type["fields"]

        return self.model_field_parser.build_custom_type_from_columns(row, custom_type_fields, custom_type_name)

    @staticmethod
    def to_camel_case(s: str) -> str:
        # Handle PascalCase
        s_with_spaces = re.sub(r"(?<!^)(?=[A-Z])", " ", s)

        parts = re.split(r"[\s_\-]+", s_with_spaces.strip())
        return parts[0].lower() + "".join(word.capitalize() for word in parts[1:])


def get_config_value(prompt: str, default: str = "", secret: bool = False) -> str:
    if default:
        prompt = f"{prompt} [{default}]: "
    else:
        prompt = f"{prompt}: "

    if secret:
        value = getpass(prompt)
    else:
        value = input(prompt)

    return value.strip() if value.strip() else default


def save_config(config: Dict[str, str]) -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)

    cache_config = {k: v for k, v in config.items() if k not in ["password", "ADMIN_PASSWORD"]}

    with open(CONFIG_FILE, "w") as f:
        json.dump(cache_config, f, indent=2)

    logger.info(f"✅ Configuration saved to {CONFIG_FILE}")


def load_cached_config() -> Dict[str, str]:
    if not CONFIG_FILE.exists():
        return {}

    try:
        with open(CONFIG_FILE, "r") as f:
            return json.load(f)
    except Exception as e:
        logger.warning(f"Failed to load cached config: {e}")
        return {}


def get_cached_or_prompt(key: str, prompt: str, cached_config: Dict, default: str = "", secret: bool = False) -> str:
    if key in cached_config:
        return cached_config[key]

    return get_config_value(prompt, default, secret)


def cmd_show(args=None):
    print(
        """
    ╔════════════════════════════════════════════════════╗
    ║        Amplify Migrator - Current Configuration    ║
    ╚════════════════════════════════════════════════════╝
    """
    )

    cached_config = load_cached_config()

    if not cached_config:
        print("\n❌ No configuration found!")
        print("💡 Run 'amplify-migrator config' first to set up your configuration.")
        return

    print("\n📋 Cached Configuration:")
    print("-" * 54)
    print(f"Excel file path:      {cached_config.get('excel_path', 'N/A')}")
    print(f"API endpoint:         {cached_config.get('api_endpoint', 'N/A')}")
    print(f"AWS Region:           {cached_config.get('region', 'N/A')}")
    print(f"User Pool ID:         {cached_config.get('user_pool_id', 'N/A')}")
    print(f"Client ID:            {cached_config.get('client_id', 'N/A')}")
    print(f"Admin Username:       {cached_config.get('username', 'N/A')}")
    print("-" * 54)
    print(f"\n📍 Config location: {CONFIG_FILE}")
    print(f"💡 Run 'amplify-migrator config' to update configuration.")


def cmd_config(args=None):
    print(
        """
    ╔════════════════════════════════════════════════════╗
    ║        Amplify Migrator - Configuration Setup      ║
    ╚════════════════════════════════════════════════════╝
    """
    )

    cached_config = load_cached_config()

    config = {
        "excel_path": get_config_value("Excel file path", cached_config.get("excel_path")),
        "api_endpoint": get_config_value("AWS Amplify API endpoint", cached_config.get("api_endpoint")),
        "region": get_config_value("AWS Region", cached_config.get("region")),
        "user_pool_id": get_config_value("Cognito User Pool ID", cached_config.get("user_pool_id")),
        "client_id": get_config_value("Cognito Client ID", cached_config.get("client_id")),
        "username": get_config_value("Admin Username", cached_config.get("username")),
    }

    save_config(config)
    print("\n✅ Configuration saved successfully!")
    print(f"💡 You can now run 'amplify-migrator migrate' to start the migration.")


def cmd_migrate(args=None):
    print(
        """
    ╔════════════════════════════════════════════════════╗
    ║             Migrator Tool for Amplify              ║
    ╠════════════════════════════════════════════════════╣
    ║   This tool requires admin privileges to execute   ║
    ╚════════════════════════════════════════════════════╝
    """
    )

    cached_config = load_cached_config()

    if not cached_config:
        print("\n❌ No configuration found!")
        print("💡 Run 'amplify-migrator config' first to set up your configuration.")
        sys.exit(1)

    excel_path = get_cached_or_prompt("excel_path", "Excel file path", cached_config, "data.xlsx")
    api_endpoint = get_cached_or_prompt("api_endpoint", "AWS Amplify API endpoint", cached_config)
    region = get_cached_or_prompt("region", "AWS Region", cached_config, "us-east-1")
    user_pool_id = get_cached_or_prompt("user_pool_id", "Cognito User Pool ID", cached_config)
    client_id = get_cached_or_prompt("client_id", "Cognito Client ID", cached_config)
    username = get_cached_or_prompt("username", "Admin Username", cached_config)

    print("\n🔐 Authentication:")
    print("-" * 54)
    password = get_config_value("Admin Password", secret=True)

    migrator = ExcelToAmplifyMigrator(excel_path)
    migrator.init_client(api_endpoint, region, user_pool_id, client_id=client_id, username=username)
    if not migrator.authenticate(username, password):
        return

    migrator.run()


def main():
    parser = argparse.ArgumentParser(
        description="Amplify Excel Migrator - Migrate Excel data to AWS Amplify GraphQL API",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    config_parser = subparsers.add_parser("config", help="Configure the migration tool")
    config_parser.set_defaults(func=cmd_config)

    show_parser = subparsers.add_parser("show", help="Show current configuration")
    show_parser.set_defaults(func=cmd_show)

    migrate_parser = subparsers.add_parser("migrate", help="Run the migration")
    migrate_parser.set_defaults(func=cmd_migrate)

    args = parser.parse_args()

    if args.command is None:
        parser.print_help()
        sys.exit(1)

    args.func(args)


if __name__ == "__main__":
    # For IDE debugging: set the command you want to test
    # Uncomment and modify one of these lines:

    # sys.argv = ["migrator.py", "config"]  # Test config command
    # sys.argv = ['migrator.py', 'show']    # Test show command
    sys.argv = ["migrator.py", "migrate"]  # Test migrate command

    main()
