app/services/csv_analysis_service.py

from __future__ import annotations

import asyncio
import json
import logging
import shutil
import sys
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

from app.services.dataset_metadata_service import extract_metadata
from app.utils.http_utils import download_url
from app.utils.subprocess_utils import run_subprocess

logger = logging.getLogger(__name__)

_PYTHON = getattr(sys, "executable", None) or "python3"
_MAX_OUTPUT_BYTES = 65536
_MAX_CONCURRENT = 8
_DOWNLOAD_TIMEOUT = 120

_semaphore = asyncio.Semaphore(_MAX_CONCURRENT)


class CSVAnalysisError(Exception):
    pass


async def _resolve_source(source: Union[str, bytes]) -> Tuple[bytes, Optional[str]]:
    if isinstance(source, str) and source.lower().startswith(("http://", "https://")):
        return await download_url(source, timeout_seconds=_DOWNLOAD_TIMEOUT)
    elif isinstance(source, bytes):
        return source, None
    else:
        raise TypeError(f"Unsupported source type: {type(source)}")


def _run_subprocess(cmd: List[str], timeout: float, max_output: int) -> Dict[str, Any]:
    return run_subprocess(cmd, timeout, max_output)


_CHAT_SCRIPT = """\
import json, sys, io, base64, traceback
import pandas as pd, numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns

df = pd.read_csv(r"{csv_path}")

with open(r"{blocks_path}", "r") as _f:
    _data = json.load(_f)

_results = {{"analyze": [], "visualization": []}}
_ns = {{"df": df, "pd": pd, "np": np, "plt": plt, "sns": sns}}
_EXCLUDED = {{"pd", "np", "plt", "sns", "df", "json", "sys", "io", "base64", "traceback", "matplotlib", "seaborn"}}

for _b in _data.get("analyze", []):
    _code = (_b.get("python_code") or "").strip()
    if not _code:
        _results["analyze"].append({{"success": True, "output": "", "error": None}})
        continue
    _old = sys.stdout
    sys.stdout = io.StringIO()
    _pre = {{k for k in _ns if not k.startswith("_")}} - _EXCLUDED
    try:
        exec(_code, _ns)
        _output = sys.stdout.getvalue()
        if not _output:
            _post = {{k for k in _ns if not k.startswith("_")}} - _EXCLUDED
            _new = sorted(_post - _pre)
            if "final_result" in _ns:
                _output += f"final_result = {{repr(_ns['final_result'])}}\\n"
            for _k in _new:
                if _k == "final_result":
                    continue
                _v = _ns[_k]
                try:
                    _s = repr(_v)
                except Exception:
                    _s = "<unrepresentable>"
                _output += f"{{_k}} = {{_s}}\\n"
        _results["analyze"].append({{"success": True, "output": _output, "error": None}})
    except Exception:
        _results["analyze"].append({{"success": False, "output": sys.stdout.getvalue(), "error": traceback.format_exc()}})
    finally:
        sys.stdout = _old

for _b in _data.get("visualization", []):
    _code = (_b.get("python_code") or "").strip()
    if not _code:
        _results["visualization"].append({{"success": True, "image_base64": "", "error": None}})
        continue
    _full = _code + (
        "\\nfrom io import BytesIO\\nimport base64\\n"
        "_buf = BytesIO()\\nplt.savefig(_buf, format='png', bbox_inches='tight', dpi=150)\\n"
        "_buf.seek(0)\\nprint(base64.b64encode(_buf.read()).decode(), end='')\\n"
        "plt.close('all')\\n"
    )
    _old = sys.stdout
    sys.stdout = io.StringIO()
    try:
        exec(_full, _ns)
        _results["visualization"].append({{"success": True, "image_base64": sys.stdout.getvalue().strip(), "error": None}})
    except Exception:
        _results["visualization"].append({{"success": False, "image_base64": None, "error": traceback.format_exc()}})
    finally:
        sys.stdout = _old
    plt.close("all")

print(json.dumps(_results))
"""


async def execute_csv_chat_blocks(
    source: Union[str, bytes],
    analyze_blocks: List[Dict[str, Any]],
    viz_blocks: List[Dict[str, Any]],
    timeout: int = 60,
) -> Dict[str, Any]:
    data, _ = await _resolve_source(source)
    if not data:
        return {"success": False, "results": None, "error": "No data provided"}

    async with _semaphore:
        run_dir = None
        start = time.monotonic()
        try:
            run_dir = Path(tempfile.mkdtemp())
            csv_path = run_dir / "data.csv"
            csv_path.write_bytes(data)

            blocks_path = run_dir / "blocks.json"
            blocks_path.write_text(
                json.dumps({"analyze": analyze_blocks, "visualization": viz_blocks}),
                encoding="utf-8",
            )

            script = _CHAT_SCRIPT.format(
                csv_path=csv_path.as_posix(),
                blocks_path=blocks_path.as_posix(),
            )

            script_path = run_dir / "chat_exec.py"
            script_path.write_text(script, encoding="utf-8")

            cmd = [_PYTHON, str(script_path)]
            result = await asyncio.to_thread(_run_subprocess, cmd, timeout, _MAX_OUTPUT_BYTES)
            elapsed_ms = (time.monotonic() - start) * 1000

            if result["exit_code"] != 0:
                return {
                    "success": False,
                    "results": None,
                    "error": result["stderr"] or "Subprocess failed",
                    "execution_time_ms": round(elapsed_ms, 2),
                }

            parsed = json.loads(result["stdout"])
            return {
                "success": True,
                "results": parsed,
                "error": None,
                "execution_time_ms": round(elapsed_ms, 2),
            }
        except json.JSONDecodeError as exc:
            elapsed_ms = (time.monotonic() - start) * 1000
            return {"success": False, "results": None, "error": f"Failed to parse output: {exc}", "execution_time_ms": round(elapsed_ms, 2)}
        except FileNotFoundError:
            elapsed_ms = (time.monotonic() - start) * 1000
            return {"success": False, "results": None, "error": f"Python runtime ({_PYTHON}) not found", "execution_time_ms": round(elapsed_ms, 2)}
        except Exception as exc:
            elapsed_ms = (time.monotonic() - start) * 1000
            logger.exception("CSV chat execution error")
            return {"success": False, "results": None, "error": f"Execution error: {exc}", "execution_time_ms": round(elapsed_ms, 2)}
        finally:
            if run_dir and run_dir.exists():
                shutil.rmtree(run_dir, ignore_errors=True)


async def get_dataset_info(source: Union[str, bytes, Any]) -> Dict[str, Any]:
    return await extract_metadata(source)