app/services/reconciliation_service.py

from __future__ import annotations

import asyncio
import io
import os
import re
import time
import unicodedata
from collections import Counter
from datetime import datetime
from typing import Any, Dict, List, Optional, Set, Tuple

import aiohttp
import pandas as pd

from app.config import get_settings
from app.core.logger import get_logger

_logger = get_logger(__name__)
_settings = get_settings()

DOWNLOAD_TIMEOUT = 60
DOWNLOAD_MAX_RETRIES = 3
DOWNLOAD_BACKOFF_FACTOR = 0.5
SUPPORTED_EXTENSIONS: Set[str] = {"csv", "xlsx", "xls", "tsv", "parquet"}
_MAX_FILE_SIZE = _settings.max_upload_bytes

DEFAULT_NUMERIC_KEYWORDS: Set[str] = {
    "amount", "total", "debit", "credit", "tax",
    "net amount", "gross amount", "balance", "quantity", "price", "rate", "value"
}


class ReconciliationError(Exception): pass
class DownloadError(ReconciliationError): pass
class FileTypeMismatchError(ReconciliationError): pass
class SchemaMismatchError(ReconciliationError): pass
class EmptyDatasetError(ReconciliationError): pass
class UnsupportedFormatError(ReconciliationError): pass


def extract_extension(url: str) -> str:
    path = url.split("?")[0]
    return os.path.splitext(path)[1].lstrip(".").lower()


def _normalize_str(x: str) -> str:
    if not pd.notna(x) or x == 'nan':
        return x
    return re.sub(r'\s+', ' ', unicodedata.normalize('NFKC', x))


def normalize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    for col in df.columns:
        if pd.api.types.is_string_dtype(df[col]):
            s = df[col].astype(str).str.strip()
            df[col] = s.apply(_normalize_str).replace({'nan': pd.NA, 'None': pd.NA, '': pd.NA, 'null': pd.NA})
        elif pd.api.types.is_numeric_dtype(df[col]):
            df[col] = df[col].replace({pd.NA: None})
    return df


def apply_column_mapping(df: pd.DataFrame, mapping: Dict[str, str]) -> pd.DataFrame:
    valid_mapping = {old: new for old, new in mapping.items() if old in df.columns}
    if valid_mapping:
        df = df.rename(columns=valid_mapping)
    return df


async def download_file_with_retry(session: aiohttp.ClientSession, url: str, correlation_id: str) -> bytes:
    for attempt in range(DOWNLOAD_MAX_RETRIES):
        try:
            async with session.get(url, timeout=aiohttp.ClientTimeout(total=DOWNLOAD_TIMEOUT)) as response:
                response.raise_for_status()
                return await response.read()
        except aiohttp.ClientError as e:
            wait_time = DOWNLOAD_BACKOFF_FACTOR * (2 ** attempt)
            _logger.warning(
                f"Download attempt {attempt+1} failed for {url}. Retrying in {wait_time}s. Error: {e}",
                extra={"correlation_id": correlation_id}
            )
            await asyncio.sleep(wait_time)
    raise DownloadError(f"Failed to download {url} after {DOWNLOAD_MAX_RETRIES} retries.")


def read_to_dataframe(data: bytes, ext: str, correlation_id: str) -> pd.DataFrame:
    try:
        if ext == "csv":
            return pd.read_csv(io.BytesIO(data))
        elif ext == "tsv":
            return pd.read_csv(io.BytesIO(data), sep="\t")
        elif ext == "xlsx":
            return pd.read_excel(io.BytesIO(data), engine="openpyxl")
        elif ext == "xls":
            return pd.read_excel(io.BytesIO(data), engine="xlrd")
        elif ext == "parquet":
            return pd.read_parquet(io.BytesIO(data))
        else:
            raise UnsupportedFormatError(f"File format '{ext}' is not supported.")
    except Exception as e:
        _logger.error(f"Failed to parse {ext} file: {str(e)}", extra={"correlation_id": correlation_id})
        raise ReconciliationError(f"Corrupted or unparseable file: {str(e)}")


def compare_schemas(df_src: pd.DataFrame, df_dst: pd.DataFrame) -> Dict[str, Any]:
    src_cols = set(df_src.columns)
    dst_cols = set(df_dst.columns)
    missing_in_src = list(dst_cols - src_cols)
    missing_in_dst = list(src_cols - dst_cols)
    type_mismatches = []
    common_cols = src_cols & dst_cols
    for col in common_cols:
        if df_src[col].dtype != df_dst[col].dtype:
            type_mismatches.append({
                "column": col,
                "source_type": str(df_src[col].dtype),
                "destination_type": str(df_dst[col].dtype)
            })
    fully_match = not missing_in_src and not missing_in_dst and not type_mismatches
    return {
        "fully_match": fully_match,
        "source_columns": len(df_src.columns),
        "destination_columns": len(df_dst.columns),
        "missing_in_source": missing_in_src,
        "missing_in_destination": missing_in_dst,
        "type_mismatches": type_mismatches
    }


def compare_rows(df_src: pd.DataFrame, df_dst: pd.DataFrame, common_cols: List[str]) -> Tuple[int, int, int, int]:
    src_subset = df_src[common_cols].astype(str).agg('|'.join, axis=1)
    dst_subset = df_dst[common_cols].astype(str).agg('|'.join, axis=1)
    src_counts = Counter(src_subset)
    dst_counts = Counter(dst_subset)
    identical = sum(min(count, dst_counts.get(row, 0)) for row, count in src_counts.items())
    missing = len(df_src) - identical
    extra = len(df_dst) - identical
    return identical, 0, missing, extra


def reconcile_columns(df_src: pd.DataFrame, df_dst: pd.DataFrame, common_cols: List[str]) -> List[Dict[str, Any]]:
    results = []
    for col in common_cols:
        src_nn = int(df_src[col].notna().sum())
        dst_nn = int(df_dst[col].notna().sum())
        status = "fully_match" if src_nn == dst_nn else "partial_match"
        results.append({
            "column_name": col,
            "status": status,
            "source_non_null": src_nn,
            "destination_non_null": dst_nn,
            "matching_values": 0,
            "mismatching_values": 0,
            "difference_count": abs(src_nn - dst_nn)
        })
    return results


def reconcile_numeric_columns(df_src: pd.DataFrame, df_dst: pd.DataFrame, common_cols: List[str], numeric_keywords: Optional[Set[str]] = None) -> List[Dict[str, Any]]:
    if numeric_keywords is None:
        numeric_keywords = DEFAULT_NUMERIC_KEYWORDS
    results = []
    for col in common_cols:
        is_numeric_name = any(kw in col.lower() for kw in numeric_keywords)
        is_numeric_dtype = pd.api.types.is_numeric_dtype(df_src[col]) and pd.api.types.is_numeric_dtype(df_dst[col])
        if is_numeric_name and is_numeric_dtype:
            src_sum = df_src[col].sum()
            dst_sum = df_dst[col].sum()
            diff = src_sum - dst_sum
            abs_diff = abs(diff)
            pct_diff = (abs_diff / src_sum * 100) if src_sum != 0 else 0.0
            src_nn = int(df_src[col].notna().sum())
            dst_nn = int(df_dst[col].notna().sum())
            status = "fully_match" if abs_diff < 1e-9 else "partial_match"
            results.append({
                "column_name": col,
                "status": status,
                "source_non_null": src_nn,
                "destination_non_null": dst_nn,
                "matching_values": src_nn,
                "mismatching_values": abs(src_nn - dst_nn),
                "difference_count": abs(src_nn - dst_nn),
                "source_total": round(float(src_sum), 2),
                "destination_total": round(float(dst_sum), 2),
                "total_difference": round(float(diff), 2),
                "absolute_difference": round(float(abs_diff), 2),
                "percentage_difference": round(float(pct_diff), 5),
                "source_min": round(float(df_src[col].min()), 2),
                "source_max": round(float(df_src[col].max()), 2),
                "source_avg": round(float(df_src[col].mean()), 2),
                "dest_min": round(float(df_dst[col].min()), 2),
                "dest_max": round(float(df_dst[col].max()), 2),
                "dest_avg": round(float(df_dst[col].mean()), 2)
            })
    return results


def reconcile_date_columns(df_src: pd.DataFrame, df_dst: pd.DataFrame, common_cols: List[str]) -> List[Dict[str, Any]]:
    results = []
    for col in common_cols:
        is_date = pd.api.types.is_datetime64_any_dtype(df_src[col]) and pd.api.types.is_datetime64_any_dtype(df_dst[col])
        if is_date:
            src_nn = int(df_src[col].notna().sum())
            dst_nn = int(df_dst[col].notna().sum())
            results.append({
                "column_name": col,
                "status": "fully_match" if df_src[col].equals(df_dst[col]) else "partial_match",
                "source_non_null": src_nn,
                "destination_non_null": dst_nn,
                "matching_values": min(src_nn, dst_nn),
                "mismatching_values": abs(src_nn - dst_nn),
                "difference_count": abs(src_nn - dst_nn),
                "source_min_date": str(df_src[col].min()),
                "source_max_date": str(df_src[col].max()),
                "destination_min_date": str(df_dst[col].min()),
                "destination_max_date": str(df_dst[col].max()),
                "missing_dates_source": int(df_src[col].isna().sum()),
                "missing_dates_destination": int(df_dst[col].isna().sum())
            })
    return results


def analyze_duplicates(df_src: pd.DataFrame, df_dst: pd.DataFrame) -> Dict[str, int]:
    src_dups = int(df_src.duplicated().sum())
    dst_dups = int(df_dst.duplicated().sum())
    return {
        "duplicate_rows_source": src_dups,
        "duplicate_rows_destination": dst_dups,
        "duplicate_difference": abs(src_dups - dst_dups)
    }


def analyze_missing_data(df_src: pd.DataFrame, df_dst: pd.DataFrame, common_cols: List[str]) -> List[Dict[str, Any]]:
    results = []
    for col in common_cols:
        src_nulls = int(df_src[col].isna().sum())
        dst_nulls = int(df_dst[col].isna().sum())
        results.append({
            "column_name": col,
            "source_nulls": src_nulls,
            "destination_nulls": dst_nulls,
            "difference": abs(src_nulls - dst_nulls)
        })
    return results


def _build_failed_pair_result(job_id: str, src_url: str, dst_url: str, started_at: str, status: str, errors: List[str], schema: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    return {
        "job_id": job_id,
        "source_file": src_url,
        "destination_file": dst_url,
        "started_at": started_at,
        "completed_at": datetime.utcnow().isoformat(),
        "processing_time_ms": 0,
        "status": status,
        "summary": {"overall_status": status, "overall_match_percentage": 0.0},
        "schema": schema or {"fully_match": False, "source_columns": 0, "destination_columns": 0},
        "columns": [], "numeric_columns": [], "date_columns": [],
        "duplicates": {}, "missing_values": [],
        "errors": errors,
        "warnings": []
    }


async def process_pair(
    src_url: str,
    dst_url: str,
    src_data: bytes,
    dst_data: bytes,
    ext: str,
    job_id: str,
    column_mapping: Optional[Dict[str, str]] = None,
    numeric_keywords: Optional[Set[str]] = None,
) -> Dict[str, Any]:
    start_time = time.time()
    started_at = datetime.utcnow().isoformat()
    errors = []
    warnings = []
    status = "MATCH"

    try:
        df_src = read_to_dataframe(src_data, ext, job_id)
        df_dst = read_to_dataframe(dst_data, ext, job_id)

        if df_src.empty or df_dst.empty:
            raise EmptyDatasetError("One or both datasets are empty.")

        df_src = normalize_dataframe(df_src.copy())
        df_dst = normalize_dataframe(df_dst.copy())
        df_src.infer_objects()
        df_dst.infer_objects()

        if column_mapping:
            df_dst = apply_column_mapping(df_dst, column_mapping)

        schema_report = compare_schemas(df_src, df_dst)
        if not schema_report["fully_match"]:
            status = "SCHEMA_MISMATCH"
            warnings.append("Schema mismatch detected.")

        common_cols = df_src.columns.intersection(df_dst.columns).tolist()

        identical_rows, _, missing_rows, extra_rows = compare_rows(df_src, df_dst, common_cols)

        if missing_rows > 0 or extra_rows > 0:
            status = "PARTIAL_MATCH" if status == "MATCH" else status

        col_reports = reconcile_columns(df_src, df_dst, common_cols)
        num_reports = reconcile_numeric_columns(df_src, df_dst, common_cols, numeric_keywords)
        date_reports = reconcile_date_columns(df_src, df_dst, common_cols)

        dup_report = analyze_duplicates(df_src, df_dst)
        missing_report = analyze_missing_data(df_src, df_dst, common_cols)

        full_cols = sum(1 for c in col_reports if c["status"] == "fully_match")
        part_cols = sum(1 for c in col_reports if c["status"] == "partial_match")

        if part_cols > 0 and status == "MATCH":
            status = "PARTIAL_MATCH"

        total_max_rows = max(len(df_src), len(df_dst))
        match_pct = (identical_rows / total_max_rows * 100) if total_max_rows > 0 else 100.0

        summary = {
            "total_columns": len(common_cols),
            "fully_matched_columns": full_cols,
            "partially_matched_columns": part_cols,
            "unmatched_columns": len(col_reports) - full_cols - part_cols,
            "total_rows_source": len(df_src),
            "total_rows_destination": len(df_dst),
            "matching_rows": identical_rows,
            "partial_rows": 0,
            "unmatched_rows": missing_rows + extra_rows,
            "duplicate_rows": dup_report["duplicate_difference"],
            "overall_match_percentage": round(match_pct, 2),
            "overall_status": status
        }
    except Exception as e:
        _logger.exception(f"Job {job_id} failed during processing.", extra={"correlation_id": job_id})
        status = "FAILED"
        errors.append(str(e))
        summary = {"overall_status": "FAILED", "overall_match_percentage": 0.0}
        schema_report, col_reports, num_reports, date_reports, dup_report, missing_report = {}, [], [], [], {}, []

    return {
        "job_id": job_id,
        "source_file": src_url,
        "destination_file": dst_url,
        "started_at": started_at,
        "completed_at": datetime.utcnow().isoformat(),
        "processing_time_ms": round((time.time() - start_time) * 1000, 2),
        "status": status,
        "summary": summary,
        "schema": schema_report,
        "columns": col_reports,
        "numeric_columns": num_reports,
        "date_columns": date_reports,
        "duplicates": dup_report,
        "missing_values": missing_report,
        "errors": errors,
        "warnings": warnings
    }


async def reconcile_pair(
    pair: Dict[str, Any],
    job_id: str,
    numeric_keywords: Optional[Set[str]] = None,
) -> Dict[str, Any]:
    src_url = pair["source"]
    dst_url = pair["destination"]
    started_at = datetime.utcnow().isoformat()

    src_ext = extract_extension(src_url)
    dst_ext = extract_extension(dst_url)

    if src_ext not in SUPPORTED_EXTENSIONS or dst_ext not in SUPPORTED_EXTENSIONS:
        return _build_failed_pair_result(
            job_id, src_url, dst_url, started_at, "FAILED",
            [f"Unsupported format. Source: {src_ext}, Dest: {dst_ext}"]
        )

    if src_ext != dst_ext:
        return _build_failed_pair_result(
            job_id, src_url, dst_url, started_at, "FILE_TYPE_MISMATCH",
            [f"File type mismatch. Source: {src_ext}, Dest: {dst_ext}"]
        )

    column_mapping: Optional[Dict[str, str]] = pair.get("column_mapping")

    async with aiohttp.ClientSession() as session:
        try:
            src_data, dst_data = await asyncio.gather(
                download_file_with_retry(session, src_url, job_id),
                download_file_with_retry(session, dst_url, job_id)
            )
        except DownloadError as e:
            _logger.error(f"Download failed for job {job_id}: {str(e)}", extra={"correlation_id": job_id})
            return _build_failed_pair_result(
                job_id, src_url, dst_url, started_at, "FAILED", [str(e)]
            )

    if len(src_data) > _MAX_FILE_SIZE or len(dst_data) > _MAX_FILE_SIZE:
        return _build_failed_pair_result(
            job_id, src_url, dst_url, started_at, "FAILED",
            ["File size exceeds maximum allowed limit"]
        )

    return await process_pair(src_url, dst_url, src_data, dst_data, src_ext, job_id, column_mapping, numeric_keywords)