IPSExtractor/services/cleaning/api/main.py

"""API REST del microservicio de limpieza."""
from __future__ import annotations

from contextlib import asynccontextmanager
import json
from typing import Any

import pandas as pd
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from pydantic import BaseModel, Field
from sqlalchemy import delete, select

from domain.csv_cleaning import (
  CLEANING_RULES,
  analyze_dataframe,
  apply_rules,
  load_csv_bytes,
  sample_dataframe_rows,
)
from infrastructure.database import (
  CleanedRow,
  CleaningJob,
  ExtractionBatch,
  ExtractionRow,
  Job,
  Task,
  init_db,
  session_scope,
  utcnow,
)

_SOURCE_BATCH_KEY = "_cleaning_source_batch_id"
_SOURCE_FILE_KEY = "_cleaning_source_file"


@asynccontextmanager
async def lifespan(app: FastAPI):
  init_db()
  yield


app = FastAPI(
  title="IPSExtractor Cleaning Service",
  version="1.0.0",
  lifespan=lifespan,
)


class CleaningRequest(BaseModel):
  extraction_batch_id: str | None = None
  extraction_batch_ids: list[str] = Field(default_factory=list)
  job_id: str | None = None
  rules_per_column: dict[str, list[str]] = Field(default_factory=dict)
  cop_mod_fecha_column: dict[str, str] = Field(default_factory=dict)
  cop_mod_bruto_column: dict[str, str] = Field(default_factory=dict)
  cop_mod_neto_column: dict[str, str] = Field(default_factory=dict)
  cop_mod_config: dict[str, Any] | None = None


class AnalyzeBatchesRequest(BaseModel):
  extraction_batch_ids: list[str] = Field(default_factory=list)
  extraction_batch_id: str | None = None


class PreviewRequest(BaseModel):
  rows: list[dict[str, Any]] = Field(default_factory=list)
  extraction_batch_id: str | None = None
  extraction_batch_ids: list[str] = Field(default_factory=list)
  rules_per_column: dict[str, list[str]] = Field(default_factory=dict)
  cop_mod_fecha_column: dict[str, str] = Field(default_factory=dict)
  cop_mod_bruto_column: dict[str, str] = Field(default_factory=dict)
  cop_mod_neto_column: dict[str, str] = Field(default_factory=dict)
  cop_mod_config: dict[str, Any] | None = None
  sample_size: int = 8


def dataframe_to_records(df: pd.DataFrame) -> list[dict[str, Any]]:
  return json.loads(df.to_json(orient="records", date_format="iso"))


def serialize_job(job: CleaningJob) -> dict[str, Any]:
  # started/finished derivados: created_at = inicio; updated_at = fin si ya terminó
  finished = (
    job.updated_at.isoformat()
    if job.status in {"completed", "failed"}
    else None
  )
  return {
    "id": job.id,
    "extraction_batch_id": job.extraction_batch_id,
    "job_id": job.job_id,
    "source": job.source,
    "status": job.status,
    "request": job.request,
    "log": job.log,
    "error": job.error,
    "created_at": job.created_at.isoformat(),
    "updated_at": job.updated_at.isoformat(),
    "started_at": job.created_at.isoformat(),
    "finished_at": finished,
  }


def _validate_job_id(session: Any, job_id: str | None) -> None:
  if not job_id:
    return
  if session.get(Job, job_id) is None:
    raise HTTPException(status_code=404, detail="Trabajo (job) no encontrado")


def _normalize_batch_ids(
  extraction_batch_id: str | None,
  extraction_batch_ids: list[str] | None,
) -> list[str]:
  ids = [str(x).strip() for x in (extraction_batch_ids or []) if str(x).strip()]
  if not ids and extraction_batch_id and extraction_batch_id.strip():
    ids = [extraction_batch_id.strip()]
  # dedupe preservando orden
  seen: set[str] = set()
  out: list[str] = []
  for bid in ids:
    if bid not in seen:
      seen.add(bid)
      out.append(bid)
  return out


def _create_cleaning_task(
  session: Any,
  *,
  job_id: str | None,
  cleaning_job_id: str,
  source: str,
  status: str,
  request: dict[str, Any],
  error: str | None = None,
  title: str | None = None,
  started_at: Any | None = None,
) -> str | None:
  """Crea Task(type=cleaning) vinculada al job Desktop. Devuelve task.id o None."""
  if not job_id:
    return None
  now = utcnow()
  task = Task(
    job_id=job_id,
    type="cleaning",
    status=status,
    title=title or f"Limpieza ({source})",
    request=request,
    result={"cleaning_job_id": cleaning_job_id},
    error=error,
    started_at=started_at or now,
    finished_at=now,
    updated_at=now,
  )
  session.add(task)
  session.flush()
  return task.id


def _persist_cleaned_rows(
  session: Any,
  *,
  job_id: str,
  extraction_batch_id: str | None,
  output_rows: list[dict[str, Any]],
) -> None:
  session.execute(delete(CleanedRow).where(CleanedRow.cleaning_job_id == job_id))
  for idx, row in enumerate(output_rows, start=1):
    row_batch = row.get(_SOURCE_BATCH_KEY)
    batch_for_row = str(row_batch).strip() if row_batch else extraction_batch_id
    cleaned = {k: v for k, v in row.items() if k not in {_SOURCE_BATCH_KEY, _SOURCE_FILE_KEY}}
    session.add(
      CleanedRow(
        cleaning_job_id=job_id,
        extraction_batch_id=batch_for_row or None,
        row_index=idx,
        row_json=cleaned,
      )
    )


def _run_rules_on_dataframe(
  df: pd.DataFrame,
  rules_per_column: dict[str, list[str]],
  cop_mod_fecha_column: dict[str, str],
  cop_mod_config: dict[str, Any] | None,
  cop_mod_bruto_column: dict[str, str] | None = None,
  cop_mod_neto_column: dict[str, str] | None = None,
) -> tuple[list[dict[str, Any]], list[str]]:
  # Conservar trazabilidad fuera de las reglas
  meta_cols = [c for c in (_SOURCE_BATCH_KEY, _SOURCE_FILE_KEY) if c in df.columns]
  meta = df[meta_cols].copy() if meta_cols else None
  work = df.drop(columns=meta_cols, errors="ignore")
  result_df, log = apply_rules(
    work,
    rules_per_column,
    cop_mod_fecha_column=cop_mod_fecha_column,
    cop_mod_config=cop_mod_config,
    cop_mod_bruto_column=cop_mod_bruto_column,
    cop_mod_neto_column=cop_mod_neto_column,
  )
  if meta is not None and not meta.empty:
    result_df = pd.concat([result_df.reset_index(drop=True), meta.reset_index(drop=True)], axis=1)
  return dataframe_to_records(result_df), log


def _parse_json_form(raw: str | None, default: Any) -> Any:
  if not raw or not raw.strip():
    return default
  try:
    return json.loads(raw)
  except json.JSONDecodeError as exc:
    raise HTTPException(status_code=400, detail=f"JSON inválido: {exc}") from exc


def _analysis_payload(df: pd.DataFrame, **extra: Any) -> dict[str, Any]:
  work = df.drop(columns=[_SOURCE_BATCH_KEY, _SOURCE_FILE_KEY], errors="ignore")
  return {
    **extra,
    "row_count": len(work),
    "columns": [summary.__dict__ for summary in analyze_dataframe(work)],
    "sample_rows": sample_dataframe_rows(work, 8),
  }


def _cell_str(value: Any) -> str:
  if value is None:
    return ""
  text = str(value)
  if text in ("None", "nan", "NaT"):
    return ""
  if len(text) > 80:
    return text[:77] + "..."
  return text


def _preview_from_dataframe(
  df: pd.DataFrame,
  rules_per_column: dict[str, list[str]],
  cop_mod_fecha_column: dict[str, str],
  cop_mod_config: dict[str, Any] | None,
  cop_mod_bruto_column: dict[str, str] | None = None,
  cop_mod_neto_column: dict[str, str] | None = None,
) -> dict[str, Any]:
  work = df.drop(columns=[_SOURCE_BATCH_KEY, _SOURCE_FILE_KEY], errors="ignore")
  rows_before = dataframe_to_records(work)
  result_df, _log = apply_rules(
    work,
    rules_per_column,
    cop_mod_fecha_column=cop_mod_fecha_column,
    cop_mod_config=cop_mod_config,
    cop_mod_bruto_column=cop_mod_bruto_column,
    cop_mod_neto_column=cop_mod_neto_column,
  )
  rows_after = dataframe_to_records(result_df)

  n = min(len(rows_before), len(rows_after))
  all_cols = list(dict.fromkeys([*work.columns.astype(str), *result_df.columns.astype(str)]))
  columns_preview: list[dict[str, Any]] = []
  for col in all_cols:
    before_vals = [_cell_str(rows_before[i].get(col)) for i in range(n)]
    after_vals = [_cell_str(rows_after[i].get(col)) for i in range(n)]
    columns_preview.append({
      "name": col,
      "before": before_vals,
      "after": after_vals,
      "changed": before_vals != after_vals,
    })

  return {
    "rows_before": rows_before,
    "rows_after": rows_after,
    "columns": columns_preview,
  }


def _load_batch_records(session: Any, batch_id: str) -> list[dict[str, Any]]:
  batch = session.get(ExtractionBatch, batch_id)
  if batch is None:
    raise HTTPException(status_code=404, detail=f"Lote de extracción no encontrado: {batch_id}")
  rows = session.execute(
    select(ExtractionRow.row_json)
    .where(ExtractionRow.batch_id == batch_id)
    .order_by(ExtractionRow.row_index, ExtractionRow.id)
  ).all()
  records = [row[0] for row in rows]
  if not records:
    raise HTTPException(status_code=400, detail=f"El lote {batch_id} no tiene filas extraídas")
  return records


def _concat_extraction_batches(session: Any, batch_ids: list[str]) -> pd.DataFrame:
  frames: list[pd.DataFrame] = []
  for bid in batch_ids:
    records = _load_batch_records(session, bid)
    part = pd.DataFrame(records)
    part[_SOURCE_BATCH_KEY] = bid
    frames.append(part)
  return pd.concat(frames, ignore_index=True, sort=False)


async def _load_csv_uploads(uploads: list[UploadFile]) -> tuple[pd.DataFrame, list[str], list[str]]:
  if not uploads:
    raise HTTPException(status_code=400, detail="Se requiere al menos un archivo CSV")
  frames: list[pd.DataFrame] = []
  filenames: list[str] = []
  delimiters: list[str] = []
  for upload in uploads:
    filename = upload.filename or "upload.csv"
    content = await upload.read()
    if not content:
      raise HTTPException(status_code=400, detail=f"El archivo {filename} está vacío")
    try:
      df, delimiter = load_csv_bytes(content, filename)
    except Exception as exc:  # noqa: BLE001
      raise HTTPException(
        status_code=400,
        detail=f"No se pudo leer {filename}: {exc}",
      ) from exc
    if df.empty:
      raise HTTPException(status_code=400, detail=f"El CSV {filename} no tiene filas")
    part = df.copy()
    part[_SOURCE_FILE_KEY] = filename
    frames.append(part)
    filenames.append(filename)
    delimiters.append(delimiter)
  return pd.concat(frames, ignore_index=True, sort=False), filenames, delimiters


def _cleaning_title_extraction(n: int) -> str:
  if n <= 1:
    return "Limpieza desde extracción"
  return f"Limpieza ({n} extracciones)"


def _cleaning_title_csv(filenames: list[str]) -> str:
  n = len(filenames)
  if n <= 1:
    return f"Limpieza CSV — {filenames[0] if filenames else 'upload.csv'}"
  return f"Limpieza ({n} CSV)"


@app.get("/health")
def health() -> dict[str, str]:
  return {"status": "ok"}


@app.get("/rules")
def list_rules() -> list[dict[str, Any]]:
  return [
    {
      "id": rule["id"],
      "name": rule["name"],
      "description": rule.get("description", ""),
    }
    for rule in CLEANING_RULES
  ]


@app.get("/batches/{batch_id}/analysis")
def analyze_batch(batch_id: str) -> dict[str, Any]:
  with session_scope() as session:
    records = _load_batch_records(session, batch_id)
  df = pd.DataFrame(records)
  return _analysis_payload(df, batch_id=batch_id, extraction_batch_ids=[batch_id])


@app.post("/analyze-batches")
def analyze_batches(request: AnalyzeBatchesRequest) -> dict[str, Any]:
  batch_ids = _normalize_batch_ids(request.extraction_batch_id, request.extraction_batch_ids)
  if not batch_ids:
    raise HTTPException(status_code=400, detail="extraction_batch_ids es requerido (mínimo 1)")
  with session_scope() as session:
    df = _concat_extraction_batches(session, batch_ids)
  return _analysis_payload(
    df,
    batch_id=batch_ids[0],
    extraction_batch_ids=batch_ids,
  )


@app.post("/analyze-csv")
async def analyze_csv(
  file: UploadFile | None = File(None),
  files: list[UploadFile] | None = File(None),
) -> dict[str, Any]:
  uploads: list[UploadFile] = []
  if files:
    uploads.extend(files)
  if file is not None:
    uploads.append(file)
  df, filenames, delimiters = await _load_csv_uploads(uploads)
  return _analysis_payload(
    df,
    filename=filenames[0] if filenames else None,
    filenames=filenames,
    delimiter=delimiters[0] if delimiters else None,
    delimiters=delimiters,
  )


@app.post("/preview")
def preview_cleaning(request: PreviewRequest) -> dict[str, Any]:
  """Aplica reglas sobre una muestra pequeña y devuelve Antes/Después."""
  records: list[dict[str, Any]] = list(request.rows or [])
  batch_ids = _normalize_batch_ids(request.extraction_batch_id, request.extraction_batch_ids)

  if not records and batch_ids:
    with session_scope() as session:
      df_full = _concat_extraction_batches(session, batch_ids)
    if df_full.empty:
      raise HTTPException(status_code=400, detail="Los lotes no tienen filas")
    records = sample_dataframe_rows(df_full, max(1, min(request.sample_size, 20)))

  if not records:
    raise HTTPException(
      status_code=400,
      detail="Se requieren rows o extraction_batch_ids con filas",
    )

  df = pd.DataFrame(records)
  try:
    return _preview_from_dataframe(
      df,
      request.rules_per_column,
      request.cop_mod_fecha_column,
      request.cop_mod_config,
      request.cop_mod_bruto_column,
      request.cop_mod_neto_column,
    )
  except Exception as exc:  # noqa: BLE001
    raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.post("/jobs")
def run_cleaning(request: CleaningRequest) -> dict[str, Any]:
  batch_ids = _normalize_batch_ids(request.extraction_batch_id, request.extraction_batch_ids)
  if not batch_ids:
    raise HTTPException(
      status_code=400,
      detail="extraction_batch_ids es requerido (use /jobs/from-csv para CSV)",
    )

  with session_scope() as session:
    _validate_job_id(session, request.job_id)
    df = _concat_extraction_batches(session, batch_ids)

    request_dump = request.model_dump()
    request_dump["extraction_batch_ids"] = batch_ids
    request_dump["extraction_batch_id"] = batch_ids[0]

    job = CleaningJob(
      extraction_batch_id=batch_ids[0],
      job_id=request.job_id,
      source="extraction",
      status="running",
      request=request_dump,
      updated_at=utcnow(),
    )
    session.add(job)
    session.flush()
    job_id = job.id
    primary_batch_id = batch_ids[0]
    desktop_job_id = request.job_id
    title = _cleaning_title_extraction(len(batch_ids))

  try:
    output_rows, log = _run_rules_on_dataframe(
      df,
      request.rules_per_column,
      request.cop_mod_fecha_column,
      request.cop_mod_config,
      request.cop_mod_bruto_column,
      request.cop_mod_neto_column,
    )

    with session_scope() as session:
      _persist_cleaned_rows(
        session,
        job_id=job_id,
        extraction_batch_id=primary_batch_id,
        output_rows=output_rows,
      )
      job = session.get(CleaningJob, job_id)
      if job is None:
        raise RuntimeError("No se encontro el job de limpieza recien creado.")
      job.status = "completed"
      job.log = log
      job.updated_at = utcnow()
      _create_cleaning_task(
        session,
        job_id=desktop_job_id,
        cleaning_job_id=job_id,
        source="extraction",
        status="completed",
        request=request_dump,
        title=title,
        started_at=job.created_at,
      )
      session.flush()
      return serialize_job(job)
  except HTTPException:
    raise
  except Exception as exc:
    with session_scope() as session:
      job = session.get(CleaningJob, job_id)
      started = job.created_at if job is not None else None
      if job is not None:
        job.status = "failed"
        job.error = str(exc)
        job.updated_at = utcnow()
      _create_cleaning_task(
        session,
        job_id=desktop_job_id,
        cleaning_job_id=job_id,
        source="extraction",
        status="failed",
        request=request_dump,
        error=str(exc),
        title=title,
        started_at=started,
      )
    raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.post("/jobs/from-csv")
async def run_cleaning_from_csv(
  file: UploadFile | None = File(None),
  files: list[UploadFile] | None = File(None),
  job_id: str | None = Form(None),
  rules_per_column: str | None = Form(None),
  cop_mod_fecha_column: str | None = Form(None),
  cop_mod_bruto_column: str | None = Form(None),
  cop_mod_neto_column: str | None = Form(None),
  cop_mod_config: str | None = Form(None),
) -> dict[str, Any]:
  uploads: list[UploadFile] = []
  if files:
    uploads.extend(files)
  if file is not None:
    uploads.append(file)

  df, filenames, delimiters = await _load_csv_uploads(uploads)

  rules = _parse_json_form(rules_per_column, {})
  fecha_map = _parse_json_form(cop_mod_fecha_column, {})
  bruto_map = _parse_json_form(cop_mod_bruto_column, {})
  neto_map = _parse_json_form(cop_mod_neto_column, {})
  cop_cfg = _parse_json_form(cop_mod_config, None)
  if not isinstance(rules, dict):
    raise HTTPException(status_code=400, detail="rules_per_column debe ser un objeto JSON")
  if not isinstance(fecha_map, dict):
    raise HTTPException(status_code=400, detail="cop_mod_fecha_column debe ser un objeto JSON")
  if not isinstance(bruto_map, dict):
    raise HTTPException(status_code=400, detail="cop_mod_bruto_column debe ser un objeto JSON")
  if not isinstance(neto_map, dict):
    raise HTTPException(status_code=400, detail="cop_mod_neto_column debe ser un objeto JSON")

  request_dump: dict[str, Any] = {
    "source": "csv",
    "filename": filenames[0] if filenames else None,
    "filenames": filenames,
    "delimiter": delimiters[0] if delimiters else None,
    "delimiters": delimiters,
    "job_id": job_id,
    "rules_per_column": rules,
    "cop_mod_fecha_column": fecha_map,
    "cop_mod_bruto_column": bruto_map,
    "cop_mod_neto_column": neto_map,
    "cop_mod_config": cop_cfg,
  }
  title = _cleaning_title_csv(filenames)

  with session_scope() as session:
    _validate_job_id(session, job_id)
    job = CleaningJob(
      extraction_batch_id=None,
      job_id=job_id,
      source="csv",
      status="running",
      request=request_dump,
      updated_at=utcnow(),
    )
    session.add(job)
    session.flush()
    cleaning_job_id = job.id
    desktop_job_id = job_id

  try:
    output_rows, log = _run_rules_on_dataframe(
      df, rules, fecha_map, cop_cfg, bruto_map, neto_map
    )

    with session_scope() as session:
      _persist_cleaned_rows(
        session,
        job_id=cleaning_job_id,
        extraction_batch_id=None,
        output_rows=output_rows,
      )
      job = session.get(CleaningJob, cleaning_job_id)
      if job is None:
        raise RuntimeError("No se encontro el job de limpieza recien creado.")
      job.status = "completed"
      job.log = log
      job.updated_at = utcnow()
      _create_cleaning_task(
        session,
        job_id=desktop_job_id,
        cleaning_job_id=cleaning_job_id,
        source="csv",
        status="completed",
        request=request_dump,
        title=title,
        started_at=job.created_at,
      )
      session.flush()
      return serialize_job(job)
  except HTTPException:
    raise
  except Exception as exc:
    with session_scope() as session:
      job = session.get(CleaningJob, cleaning_job_id)
      started = job.created_at if job is not None else None
      if job is not None:
        job.status = "failed"
        job.error = str(exc)
        job.updated_at = utcnow()
      _create_cleaning_task(
        session,
        job_id=desktop_job_id,
        cleaning_job_id=cleaning_job_id,
        source="csv",
        status="failed",
        request=request_dump,
        error=str(exc),
        title=title,
        started_at=started,
      )
    raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.get("/jobs")
def list_jobs(job_id: str | None = None) -> list[dict[str, Any]]:
  with session_scope() as session:
    stmt = select(CleaningJob).order_by(CleaningJob.created_at.desc())
    if job_id:
      stmt = stmt.where(CleaningJob.job_id == job_id)
    jobs = session.scalars(stmt).all()
    return [serialize_job(job) for job in jobs]


@app.get("/jobs/{job_id}")
def get_job(job_id: str) -> dict[str, Any]:
  if job_id in {"from-csv", "analyze-batches"}:
    raise HTTPException(status_code=405, detail=f"Use POST /jobs/{job_id}")
  with session_scope() as session:
    job = session.get(CleaningJob, job_id)
    if job is None:
      raise HTTPException(status_code=404, detail="Job de limpieza no encontrado")
    return serialize_job(job)


@app.get("/jobs/{job_id}/rows")
def get_cleaned_rows(job_id: str, limit: int = 500, offset: int = 0) -> dict[str, Any]:
  with session_scope() as session:
    job = session.get(CleaningJob, job_id)
    if job is None:
      raise HTTPException(status_code=404, detail="Job de limpieza no encontrado")
    rows = session.scalars(
      select(CleanedRow)
      .where(CleanedRow.cleaning_job_id == job_id)
      .order_by(CleanedRow.row_index, CleanedRow.id)
      .offset(offset)
      .limit(limit)
    ).all()
    return {
      "job_id": job_id,
      "rows": [row.row_json for row in rows],
      "limit": limit,
      "offset": offset,
    }