app/api/v1/reconcile.py

from __future__ import annotations

import asyncio
import json
import os
import time
import uuid
from typing import Any, Dict, List, Optional

import aiohttp
from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
from pydantic import BaseModel

from app.api.deps import require_auth
from app.config import get_settings
from app.core.logger import get_logger
from app.services.reconciliation_service import (
    SUPPORTED_EXTENSIONS,
    analyze_duplicates,
    analyze_missing_data,
    apply_column_mapping,
    compare_rows,
    compare_schemas,
    download_file_with_retry,
    normalize_dataframe,
    read_to_dataframe,
    reconcile_columns,
    reconcile_date_columns,
    reconcile_numeric_columns,
)

router = APIRouter()
_logger = get_logger(__name__)
_settings = get_settings()
_MAX_UPLOAD_BYTES = _settings.max_upload_bytes
_MAX_PAIRS = 10
semaphore_jobs = asyncio.Semaphore(5)


class ColumnMapping(BaseModel):
    source: str
    destination: str


class ReconciliationPair(BaseModel):
    source: str
    destination: str
    column_mapping: Optional[List[ColumnMapping]] = None


class ReconciliationUrlRequest(BaseModel):
    pairs: List[ReconciliationPair]


class ReconciliationPairResult(BaseModel):
    pair_index: int
    status: str
    error: Optional[str] = None
    schema_report: Optional[Dict[str, Any]] = None
    columns: Optional[List[Dict[str, Any]]] = None
    numeric_columns: Optional[List[Dict[str, Any]]] = None
    date_columns: Optional[List[Dict[str, Any]]] = None
    duplicates: Optional[Dict[str, int]] = None
    missing_values: Optional[List[Dict[str, Any]]] = None


class ReconciliationResponse(BaseModel):
    success: bool
    total_pairs: int
    completed_pairs: int
    failed_pairs: int
    processing_time_ms: float
    results: List[ReconciliationPairResult]


def _extract_extension(filename: str) -> str:
    return os.path.splitext(filename)[1].lstrip(".").lower()


def _failed_result(pair_index: int, error: str) -> Dict[str, Any]:
    return {
        "pair_index": pair_index,
        "status": "FAILED",
        "error": error,
        "schema_report": None,
        "columns": None,
        "numeric_columns": None,
        "date_columns": None,
        "duplicates": None,
        "missing_values": None,
    }


def _validate_extensions(src_ext: str, dst_ext: str, idx: int) -> None:
    if src_ext not in SUPPORTED_EXTENSIONS:
        raise HTTPException(status_code=400, detail={"success": False, "message": f"Pair {idx}: Unsupported source format '{src_ext}'."})
    if dst_ext not in SUPPORTED_EXTENSIONS:
        raise HTTPException(status_code=400, detail={"success": False, "message": f"Pair {idx}: Unsupported destination format '{dst_ext}'."})
    if src_ext != dst_ext:
        raise HTTPException(status_code=400, detail={"success": False, "message": f"Pair {idx}: File type mismatch - source '{src_ext}' != destination '{dst_ext}'."})


def _build_response(results: List[Dict[str, Any]], start_time: float) -> ReconciliationResponse:
    failed = sum(1 for r in results if r["status"] == "FAILED")
    return ReconciliationResponse(
        success=failed == 0,
        total_pairs=len(results),
        completed_pairs=len(results) - failed,
        failed_pairs=failed,
        processing_time_ms=round((time.time() - start_time) * 1000, 2),
        results=results,
    )


def _parse_column_mapping_str(value: Optional[str]) -> Optional[Dict[str, str]]:
    if not value:
        return None
    try:
        parsed = json.loads(value)
        if isinstance(parsed, list):
            return {item["destination"]: item["source"] for item in parsed if "source" in item and "destination" in item}
    except (json.JSONDecodeError, KeyError, TypeError):
        pass
    return None


def _parse_column_mapping_model(mappings: Optional[List[ColumnMapping]]) -> Optional[Dict[str, str]]:
    if not mappings:
        return None
    return {m.destination: m.source for m in mappings}


async def _process_file_pair(
    source_file: bytes,
    dest_file: bytes,
    source_ext: str,
    dest_ext: str,
    column_mapping: Optional[Dict[str, str]] = None,
    pair_index: int = 0,
) -> Dict[str, Any]:
    if source_ext != dest_ext:
        return _failed_result(pair_index, f"File type mismatch. Source: {source_ext}, Destination: {dest_ext}")

    try:
        df_src = read_to_dataframe(source_file, source_ext, f"pair-{pair_index}")
        df_dst = read_to_dataframe(dest_file, dest_ext, f"pair-{pair_index}")

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

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

        if column_mapping:
            df_dst = apply_column_mapping(df_dst, column_mapping)

        schema_report = compare_schemas(df_src, df_dst)
        status = "MATCH" if schema_report["fully_match"] else "SCHEMA_MISMATCH"

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

        _, _, missing_rows, extra_rows = compare_rows(df_src, df_dst, common_cols)
        if (missing_rows > 0 or extra_rows > 0) and status == "MATCH":
            status = "PARTIAL_MATCH"

        col_reports = reconcile_columns(df_src, df_dst, common_cols)
        num_reports = reconcile_numeric_columns(df_src, df_dst, common_cols)
        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)

        return {
            "pair_index": pair_index,
            "status": status,
            "error": None,
            "schema_report": schema_report,
            "columns": col_reports,
            "numeric_columns": num_reports,
            "date_columns": date_reports,
            "duplicates": dup_report,
            "missing_values": missing_report,
        }
    except Exception as e:
        _logger.exception(f"Pair {pair_index} failed during processing.")
        return _failed_result(pair_index, str(e))


@router.post(
    "/reconcile/files",
    response_model=ReconciliationResponse,
    summary="Reconcile uploaded file pairs (up to 10 pairs)",
)
async def reconcile_files(
    source_1: UploadFile = File(None),
    destination_1: UploadFile = File(None),
    column_mapping_1: Optional[str] = Form(None),
    source_2: UploadFile = File(None),
    destination_2: UploadFile = File(None),
    column_mapping_2: Optional[str] = Form(None),
    source_3: UploadFile = File(None),
    destination_3: UploadFile = File(None),
    column_mapping_3: Optional[str] = Form(None),
    source_4: UploadFile = File(None),
    destination_4: UploadFile = File(None),
    column_mapping_4: Optional[str] = Form(None),
    source_5: UploadFile = File(None),
    destination_5: UploadFile = File(None),
    column_mapping_5: Optional[str] = Form(None),
    source_6: UploadFile = File(None),
    destination_6: UploadFile = File(None),
    column_mapping_6: Optional[str] = Form(None),
    source_7: UploadFile = File(None),
    destination_7: UploadFile = File(None),
    column_mapping_7: Optional[str] = Form(None),
    source_8: UploadFile = File(None),
    destination_8: UploadFile = File(None),
    column_mapping_8: Optional[str] = Form(None),
    source_9: UploadFile = File(None),
    destination_9: UploadFile = File(None),
    column_mapping_9: Optional[str] = Form(None),
    source_10: UploadFile = File(None),
    destination_10: UploadFile = File(None),
    column_mapping_10: Optional[str] = Form(None),
    token: str = Depends(require_auth),
) -> ReconciliationResponse:
    start_time = time.time()

    files = [
        (source_1, destination_1, column_mapping_1),
        (source_2, destination_2, column_mapping_2),
        (source_3, destination_3, column_mapping_3),
        (source_4, destination_4, column_mapping_4),
        (source_5, destination_5, column_mapping_5),
        (source_6, destination_6, column_mapping_6),
        (source_7, destination_7, column_mapping_7),
        (source_8, destination_8, column_mapping_8),
        (source_9, destination_9, column_mapping_9),
        (source_10, destination_10, column_mapping_10),
    ]

    pairs: List[tuple] = []
    for idx, (src, dst, cm) in enumerate(files, 1):
        if src is None and dst is None:
            continue
        if src is None or dst is None:
            raise HTTPException(status_code=400, detail={"success": False, "message": f"Pair {idx}: Both source and destination files required."})
        src_ext = _extract_extension(src.filename or "")
        dst_ext = _extract_extension(dst.filename or "")
        _validate_extensions(src_ext, dst_ext, idx)
        pairs.append((src, dst, cm, idx, src_ext, dst_ext))

    if not pairs:
        raise HTTPException(status_code=400, detail={"success": False, "message": "At least one file pair is required."})

    async def process_file_pair(src: UploadFile, dst: UploadFile, cm_str: Optional[str], idx: int, src_ext: str, dst_ext: str) -> Dict[str, Any]:
        async with semaphore_jobs:
            src_data = await src.read()
            dst_data = await dst.read()
            if len(src_data) > _MAX_UPLOAD_BYTES or len(dst_data) > _MAX_UPLOAD_BYTES:
                return _failed_result(idx, "File size exceeds maximum allowed limit")
            mapping = _parse_column_mapping_str(cm_str)
            return await _process_file_pair(src_data, dst_data, src_ext, dst_ext, mapping, idx)

    results = await asyncio.gather(*[process_file_pair(*p) for p in pairs])
    return _build_response(results, start_time)


@router.post(
    "/reconcile/urls",
    response_model=ReconciliationResponse,
    summary="Reconcile file URL pairs (up to 10 pairs)",
)
async def reconcile_urls(
    body: ReconciliationUrlRequest,
    token: str = Depends(require_auth),
) -> ReconciliationResponse:
    start_time = time.time()

    if len(body.pairs) > _MAX_PAIRS:
        raise HTTPException(status_code=400, detail={"success": False, "message": f"Maximum {_MAX_PAIRS} pairs per request."})
    if len(body.pairs) < 1:
        raise HTTPException(status_code=400, detail={"success": False, "message": "At least 1 pair is required."})

    for idx, pair in enumerate(body.pairs, 1):
        _validate_extensions(_extract_extension(pair.source), _extract_extension(pair.destination), idx)

    async with aiohttp.ClientSession() as session:
        async def process_url_pair(pair: ReconciliationPair, idx: int) -> Dict[str, Any]:
            async with semaphore_jobs:
                src_ext = _extract_extension(pair.source)
                dst_ext = _extract_extension(pair.destination)

                try:
                    src_data, dst_data = await asyncio.gather(
                        download_file_with_retry(session, pair.source, str(uuid.uuid4())),
                        download_file_with_retry(session, pair.destination, str(uuid.uuid4())),
                    )
                except Exception as e:
                    return _failed_result(idx, str(e))

                if len(src_data) > _MAX_UPLOAD_BYTES or len(dst_data) > _MAX_UPLOAD_BYTES:
                    return _failed_result(idx, "File size exceeds maximum allowed limit")

                mapping = _parse_column_mapping_model(pair.column_mapping)
                return await _process_file_pair(src_data, dst_data, src_ext, dst_ext, mapping, idx)

        results = await asyncio.gather(*[process_url_pair(pair, idx) for idx, pair in enumerate(body.pairs, 1)])

    return _build_response(results, start_time)