"""
AnyCost uploader orchestrator for adhoc CSV uploads.

Coordinates the complete adhoc upload workflow using AnyCost Stream API,
including CSV processing, payload building, and HTTP uploads.
"""

import time
import json
from typing import Dict, List
from pathlib import Path
import polars as pl

from ..models import UploadConfig, WorkflowResult, OrganizationResult
from src.upload.services.data_processing_service import DataProcessingService
from .anycost_client import AnyCostClient, AnyCostAPIError
from .payload_builder import PayloadBuilder
from .anycost_connection_manager import AnyCostConnectionManager
from .connection_resolver import ConnectionResolver


class AnyCostUploader:
    """
    Orchestrate adhoc CSV upload via AnyCost Stream API.

    Single Responsibility: Coordinate complete API upload workflow
    from CSV to CloudZero ingestion.
    """

    def __init__(
        self,
        data_processor: DataProcessingService
    ):
        """
        Initialize with data processing service.

        Args:
            data_processor: Service for CBF mapping and validation
        """
        self.data_processor = data_processor
        self.payload_builder = PayloadBuilder()
        self.connection_resolver = ConnectionResolver()

    def upload_adhoc_csv(
        self,
        config: UploadConfig,
        dry_run: bool = False
    ) -> WorkflowResult:
        """
        Execute complete adhoc upload workflow via API.

        Steps:
        1. Load and process CSV with CBF mapping
        2. Build AnyCost payloads
        3. Ensure connections exist (auto-create if needed)
        4. POST payloads to AnyCost API
        5. Aggregate results

        Args:
            config: Upload configuration
            dry_run: If True, validate only without uploading

        Returns:
            WorkflowResult with upload outcomes
        """
        print(f"🚀 Starting adhoc upload (API method)")
        print(f"📁 CSV file: {config.adhoc.csv_path}")

        # Load and process CSV
        processed_df = self._load_and_process_csv(config)
        if processed_df.is_empty():
            return self._create_failed_result(config, "CSV processing failed")

        print(f"✅ Processed {len(processed_df):,} records with CBF mapping\n")

        if dry_run:
            return self._create_dry_run_result(config, processed_df)

        # Upload to each organization
        org_results = {}
        for org in config.organizations:
            print(f"🏢 {org.name}:")
            org_result = self._upload_to_organization(org, processed_df, config)
            org_results[org.name] = org_result

        return self._create_workflow_result(config, org_results)

    def _load_and_process_csv(self, config: UploadConfig) -> pl.DataFrame:
        """Load CSV(s) and apply CBF mapping. Supports file or folder."""
        csv_path = Path(config.adhoc.csv_path)

        if csv_path.is_dir():
            # Load all CSV files from directory
            csv_files = sorted(csv_path.glob("*.csv"))
            if not csv_files:
                print(f"❌ No CSV files found in directory: {csv_path}")
                return pl.DataFrame()

            print(f"📋 Loading {len(csv_files)} CSV files from: {csv_path}")
            dfs = []
            for csv_file in csv_files:
                print(f"   • {csv_file.name}")
                dfs.append(pl.read_csv(csv_file, infer_schema_length=10000))

            df = pl.concat(dfs)
            print(f"📊 Loaded {len(df):,} total rows from {len(csv_files)} files")
        else:
            # Load single CSV file
            print(f"📋 Loading CSV: {csv_path}")
            df = pl.read_csv(csv_path, infer_schema_length=10000)
            print(f"📊 Loaded {len(df):,} rows, {len(df.columns)} columns")

        print(f"📋 Applying CBF mapping...")

        # Load user mapping if provided
        user_mapping = []
        if config.adhoc.mapping_path:
            mapping_path = Path(config.adhoc.mapping_path)
            user_mapping = self.data_processor.load_user_mapping(mapping_path)

        # Process with CBF mapping
        result = self.data_processor.process_adhoc_dataset(df, user_mapping)

        if not result.is_successful():
            print(f"❌ CBF mapping failed: {result.error_message}")
            return pl.DataFrame()

        return result.processed_data

    def _upload_to_organization(
        self,
        org,
        processed_df: pl.DataFrame,
        config: UploadConfig
    ) -> OrganizationResult:
        """Upload processed data to single organization."""
        try:
            api_key = self._get_org_api_key(org.api_token_env)
            anycost_client, conn_manager = self._create_api_clients(api_key)
            s3_config = self._get_s3_config_from_env()

            total_records, total_batches, providers_uploaded = self._upload_all_providers(
                org, processed_df, anycost_client, conn_manager, s3_config
            )

            return OrganizationResult(
                organization=org.name,
                status="success",
                files_uploaded=total_batches,
                total_rows=total_records,
                providers=providers_uploaded
            )
        except Exception as e:
            print(f"   ❌ Upload failed: {str(e)}")
            return OrganizationResult(
                organization=org.name,
                status="failed",
                error_message=str(e),
                providers=org.providers
            )

    def _create_api_clients(self, api_key: str):
        """Create AnyCost and connection manager clients."""
        anycost_client = AnyCostClient(api_key)
        from .client import CloudZeroAPI
        cz_api = CloudZeroAPI(api_key)
        conn_manager = AnyCostConnectionManager(cz_api)
        return anycost_client, conn_manager

    def _upload_all_providers(
        self, org, processed_df, anycost_client, conn_manager, s3_config
    ):
        """Upload data for all providers in organization."""
        total_records = 0
        total_batches = 0
        providers_uploaded = []

        for provider in org.providers:
            print(f"   📤 {provider}:")
            records, batches = self._upload_single_provider(
                org, provider, processed_df, anycost_client, conn_manager, s3_config
            )
            total_records += records
            total_batches += batches
            providers_uploaded.append(provider)

        return total_records, total_batches, providers_uploaded

    def _upload_single_provider(
        self, org, provider, processed_df, anycost_client, conn_manager, s3_config
    ):
        """
        Upload data for single provider.

        Uses explicit connection IDs for ad-hoc uploads (required).
        """
        month = self._extract_month_from_dataframe(processed_df)

        # Try to load existing payloads from dry-run, otherwise build new ones
        payloads = self._load_or_build_payloads(org.name, provider, processed_df, month)

        # Persist payloads for inspection and debugging
        payloads_dir = self._save_payloads_for_inspection(
            None, org.name, provider, payloads
        )
        print(f"      💾 Payloads saved: {payloads_dir}")

        # Resolve connection IDs - REQUIRED for ad-hoc uploads
        connection_ids = self.connection_resolver.resolve_for_adhoc(org, provider)
        print(f"      ✓ Using {len(connection_ids)} configured connection(s)")

        records_uploaded, batches_sent = self._upload_provider(
            anycost_client, connection_ids, payloads, provider
        )

        print(f"      ✅ Uploaded {records_uploaded:,} records "
              f"({batches_sent} batches, {len(connection_ids)} connections)")

        return records_uploaded, batches_sent

    def _upload_provider(
        self,
        anycost_client: AnyCostClient,
        connection_ids: List[str],
        payloads: List[Dict],
        provider: str
    ) -> tuple[int, int]:
        """
        Upload payloads for single provider across multiple connections.

        Distributes payloads across connections using round-robin pattern,
        with first payload per connection using "replace_drop" operation.
        """
        print(f"      📦 Uploading {len(payloads)} batches...")
        start_time = time.time()

        # Track first payload per connection
        connection_first_payload = {conn_id: True for conn_id in connection_ids}

        # Upload payloads with connection cycling
        for i, payload in enumerate(payloads):
            # Cycle through connections (16 payloads per connection)
            connection_idx = (i // 16) % len(connection_ids)
            connection_id = connection_ids[connection_idx]

            # Set operation based on whether this is first for connection
            if connection_first_payload[connection_id]:
                payload["operation"] = "replace_drop"
                connection_first_payload[connection_id] = False
            else:
                payload["operation"] = "sum"

            try:
                anycost_client.post_billing_drop(connection_id, payload)
                print(f"      📤 Batch {i+1}/{len(payloads)} uploaded "
                      f"(conn {connection_idx+1}/{len(connection_ids)})", end="\r")
            except AnyCostAPIError as e:
                print(f"\n      ❌ Batch {i+1} failed: {str(e)}")
                raise

        elapsed = time.time() - start_time
        records_uploaded = sum(len(p["data"]) for p in payloads)

        print(f"      ✅ {len(payloads)} batches uploaded in {elapsed:.1f}s          ")

        return records_uploaded, len(payloads)

    def _extract_month_from_dataframe(self, df: pl.DataFrame) -> str:
        """Extract month from timestamp column (CBF: time/usage_start or internal: billing_period_start_date)."""
        # Try CBF column name first (for ad-hoc uploads)
        if "time/usage_start" in df.columns:
            first_date = df["time/usage_start"][0]
            # Parse ISO date and extract YYYY-MM
            return str(first_date)[:7]  # "2025-01-15T..." -> "2025-01" or "2025-08-01" -> "2025-08"

        # Fall back to internal column name (for generated datasets)
        if "billing_period_start_date" in df.columns:
            first_date = df["billing_period_start_date"][0]
            # Parse ISO date and extract YYYY-MM
            return str(first_date)[:7]  # "2025-01-15T..." -> "2025-01"

        raise ValueError("DataFrame missing time/usage_start or billing_period_start_date column")

    def _get_org_api_key(self, api_token_env: str) -> str:
        """Get organization API key from ~/.czbg/.env file."""
        from src.services.environment_service import EnvironmentService
        env_service = EnvironmentService()
        return env_service.require(api_token_env)

    def _get_s3_config_from_env(self) -> Dict[str, str]:
        """Get S3 configuration from environment variables."""
        from src.services.environment_service import EnvironmentService
        env_service = EnvironmentService()

        return {
            "bucket": env_service.require("CZBILLGEN_S3_BUCKET"),
            "root_folder": env_service.get("CZBILLGEN_S3_ROOT_FOLDER",
                                          "piedpiper-custom-billing"),
            "role_arn": env_service.require("AWS_IAM_ROLE_ARN"),
            "external_id": env_service.require("AWS_IAM_EXTERNAL_ID")
        }

    def _get_payloads_directory(self, org_name: str, provider: str) -> Path:
        """Get payloads directory path."""
        config_base = Path.cwd()

        # Look for configs directory
        if (config_base / "configs").exists():
            config_base = config_base / "configs"

        payloads_base = config_base / "upload-adhoc" / "orgs" / org_name / "payloads"
        return payloads_base / provider

    def _load_or_build_payloads(
        self,
        org_name: str,
        provider: str,
        processed_df: pl.DataFrame,
        month: str
    ) -> List[Dict]:
        """
        Load existing payloads from dry-run or build new ones.

        Checks for existing payload files in:
        configs/upload-adhoc/orgs/{org_name}/payloads/{provider}/

        If found, loads them. Otherwise, builds new payloads from data.
        """
        payloads_dir = self._get_payloads_directory(org_name, provider)

        # Check for existing payload files
        if payloads_dir.exists():
            payload_files = sorted(payloads_dir.glob("payload_*.json"))
            if payload_files:
                print(f"      📂 Loading existing payloads from dry-run ({len(payload_files)} files)")
                payloads = []
                for payload_file in payload_files:
                    with open(payload_file, "r") as f:
                        payloads.append(json.load(f))
                return payloads

        # No existing payloads, build new ones
        print(f"      🔨 Building new payloads from data")
        return self.payload_builder.build_payloads(processed_df, month)

    def _save_payloads_for_inspection(
        self,
        config: UploadConfig,
        org_name: str,
        provider: str,
        payloads: List[Dict]
    ) -> str:
        """Save payloads to directory for inspection."""
        payloads_dir = self._get_payloads_directory(org_name, provider)
        payloads_dir.mkdir(parents=True, exist_ok=True)

        # Save each payload as numbered JSON file
        for i, payload in enumerate(payloads, start=1):
            payload_file = payloads_dir / f"payload_{i:04d}.json"
            with open(payload_file, "w") as f:
                json.dump(payload, f, indent=2)

        return str(payloads_dir)

    def _create_dry_run_result(
        self,
        config: UploadConfig,
        processed_df: pl.DataFrame
    ) -> WorkflowResult:
        """Create dry-run validation result with payload generation."""
        print("🧪 DRY RUN MODE - Building and saving payloads\n")

        # Build and save payloads for each org/provider
        total_payloads = 0
        total_records = 0

        for org in config.organizations:
            print(f"🏢 {org.name}:")
            for provider in org.providers:
                month = self._extract_month_from_dataframe(processed_df)
                payloads = self.payload_builder.build_payloads(processed_df, month)

                # Save payloads to config directory
                payloads_saved = self._save_payloads_for_inspection(
                    config, org.name, provider, payloads
                )

                provider_records = sum(len(p["data"]) for p in payloads)
                total_payloads += len(payloads)
                total_records += provider_records

                print(f"   📦 {provider}: {len(payloads)} payloads, "
                      f"{provider_records:,} records")
                print(f"      💾 Saved to: {payloads_saved}")

        print(f"\n✅ Dry run complete:")
        print(f"   • Total payloads: {total_payloads}")
        print(f"   • Total records: {total_records:,}")

        return WorkflowResult(
            status="success",
            dataset=config.adhoc.csv_path,
            organizations_processed=len(config.organizations),
            organizations_successful=len(config.organizations),
            results={},
            mode="dry_run"
        )

    def _create_failed_result(
        self,
        config: UploadConfig,
        error_message: str
    ) -> WorkflowResult:
        """Create failed workflow result."""
        return WorkflowResult(
            status="failed",
            dataset=config.adhoc.csv_path,
            organizations_processed=0,
            organizations_successful=0,
            results={},
            error_message=error_message
        )

    def _create_workflow_result(
        self,
        config: UploadConfig,
        org_results: Dict[str, OrganizationResult]
    ) -> WorkflowResult:
        """Create workflow result from organization results."""
        successful = sum(1 for r in org_results.values() if r.is_successful())
        total = len(org_results)

        if successful == total:
            status = "success"
        elif successful > 0:
            status = "partial"
        else:
            status = "failed"

        return WorkflowResult(
            status=status,
            dataset=config.adhoc.csv_path,
            organizations_processed=total,
            organizations_successful=successful,
            results=org_results,
            mode="api"
        )
