"""CSVAnswers Studio full-stack A2A agent.

Deterministic, no-LLM CSV analysis with PlatformUserAuth, typed uploads,
bounded browser base64 bridge, managed Postgres persistence, and receipt rows.
"""
from __future__ import annotations

import base64
import csv
import hashlib
import io
import json
import os
import re
import uuid
from datetime import datetime, timezone
from decimal import Decimal, InvalidOperation
from typing import Annotated, Any

from pydantic import BaseModel, Field

import a2a_pack as a2a
from a2a_pack import (
    A2AAgent,
    AgentDatabase,
    AgentDatabaseEnv,
    AgentDatabaseMigrations,
    AgentPlatformResources,
    FileUpload,
    PlatformUserAuth,
    Pricing,
    Resources,
    RunContext,
    State,
    UploadedFile,
)

MAX_CSV_BYTES = 64 * 1024
MAX_ROWS = 500
MAX_COLUMNS = 50
MAX_CELL_CHARS = 2_000
CSV_MEDIA_TYPES = {"text/csv", "application/csv", "text/plain", "application/vnd.ms-excel"}
SAFE_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.:-]{0,127}$")

# Local-only fallback lets sandbox/unit tests prove behavior without receiving
# managed DATABASE_URL. Hosted deployments use managed Postgres exclusively.
_LOCAL_ANALYSES: dict[tuple[str, str], dict[str, Any]] = {}
_LOCAL_RECEIPTS: list[dict[str, Any]] = []


class CsvAnswersStudioV1Config(BaseModel):
    pass


class BrowserCsvUpload(BaseModel):
    filename: str = Field(..., min_length=1, max_length=180)
    media_type: str = Field(..., min_length=1, max_length=120)
    data_base64: str = Field(..., min_length=1, max_length=((MAX_CSV_BYTES * 4) // 3) + 16)


class CsvAnalysisResult(BaseModel):
    ok: bool
    analysis_id: str
    answer: str | None = None
    row_count: int = 0
    column_count: int = 0
    question: str | None = None
    source_rows: list[dict[str, Any]] = Field(default_factory=list)
    chart_data: list[dict[str, Any]] = Field(default_factory=list)
    numeric_column: str | None = None
    row_label_column: str | None = None
    saved: bool = False
    receipt_id: str | None = None
    code: str | None = None
    message: str | None = None
    warnings: list[str] = Field(default_factory=list)


class CsvAnalysisListResult(BaseModel):
    ok: bool
    analyses: list[dict[str, Any]] = Field(default_factory=list)
    code: str | None = None
    message: str | None = None


class CsvAnswersStudioV1(A2AAgent[CsvAnswersStudioV1Config, PlatformUserAuth]):
    name = "csv-answers-studio-v1"
    description = (
        "CSVAnswers: upload or paste a bounded CSV, ask which row has the "
        "highest numeric value, persist the grounded answer in user-scoped "
        "managed Postgres, and reopen it from a polished one-page app."
    )
    version = "0.1.0"

    config_model = CsvAnswersStudioV1Config
    auth_model = PlatformUserAuth

    state = State.DURABLE
    resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=120)
    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes="Deterministic CSV parsing and managed Postgres persistence. No LLM credential required.",
    )
    tools_used = ("python-csv", "managed-postgres")
    platform_resources = AgentPlatformResources(
        databases=(
            AgentDatabase(
                name="csv-answers-studio-v1-data",
                scope="user",
                access_mode="read_write",
                env=AgentDatabaseEnv(url="DATABASE_URL"),
                migrations=AgentDatabaseMigrations(path="db/migrations"),
            ),
        )
    )

    @a2a.tool(
        description="Analyze bounded pasted CSV text, answer the highest numeric row question, save the result, and persist a receipt.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="cheap",
    )
    async def analyze_csv(
        self,
        ctx: RunContext[PlatformUserAuth],
        analysis_id: str,
        csv_text: str,
        question: str,
    ) -> CsvAnalysisResult:
        tenant = _tenant_key(ctx)
        result = _analyze_csv_text(analysis_id=analysis_id, csv_text=csv_text, question=question)
        if not result.ok:
            return result
        saved, receipt_id = _persist_analysis_and_receipt(
            tenant_key=tenant,
            analysis_id=analysis_id,
            question=question,
            csv_text=csv_text,
            result=result.model_dump(mode="json"),
            source="pasted_csv",
        )
        result.saved = saved
        result.receipt_id = receipt_id
        await ctx.emit_progress(f"saved analysis {analysis_id} for {tenant}")
        return result

    @a2a.tool(
        description="Analyze a bounded browser CSV upload encoded as base64 JSON, then save the result and receipt.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="cheap",
    )
    async def analyze_csv_upload_base64(
        self,
        ctx: RunContext[PlatformUserAuth],
        analysis_id: str,
        upload: BrowserCsvUpload,
        question: str,
    ) -> CsvAnalysisResult:
        decoded = _decode_browser_upload(upload)
        if isinstance(decoded, CsvAnalysisResult):
            return decoded.model_copy(update={"analysis_id": analysis_id})
        return await self.analyze_csv(ctx, analysis_id=analysis_id, csv_text=decoded, question=question)

    @a2a.tool(
        description="Analyze a typed external FileUpload CSV, then save the result and receipt.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="cheap",
    )
    async def analyze_csv_file(
        self,
        ctx: RunContext[PlatformUserAuth],
        analysis_id: str,
        file: Annotated[
            UploadedFile,
            FileUpload(
                accept=("text/csv", "text/plain", "application/vnd.ms-excel"),
                max_bytes=MAX_CSV_BYTES,
                description="Bounded CSV file for CSVAnswers analysis.",
            ),
        ],
        question: str,
    ) -> CsvAnalysisResult:
        if file.size_bytes and file.size_bytes > MAX_CSV_BYTES:
            return _error(analysis_id, "csv_too_large", f"CSV must be at most {MAX_CSV_BYTES} bytes.")
        if file.media_type not in CSV_MEDIA_TYPES:
            return _error(analysis_id, "unsupported_media_type", "Upload must be a CSV or plain text file.")
        try:
            raw = await _read_uploaded_file(ctx, file)
        except Exception:
            return _error(analysis_id, "upload_unreadable", "The uploaded CSV could not be read from the workspace.")
        try:
            csv_text = raw.decode("utf-8-sig")
        except UnicodeDecodeError:
            return _error(analysis_id, "invalid_encoding", "CSV must be valid UTF-8 text.")
        return await self.analyze_csv(ctx, analysis_id=analysis_id, csv_text=csv_text, question=question)

    @a2a.tool(
        description="Reopen a saved user-scoped CSV analysis by id.",
        timeout_seconds=15,
        idempotent=True,
        cost_class="cheap",
    )
    async def get_analysis(
        self,
        ctx: RunContext[PlatformUserAuth],
        analysis_id: str,
    ) -> CsvAnalysisResult:
        tenant = _tenant_key(ctx)
        record = _load_analysis(tenant, analysis_id)
        if record is None:
            return _error(analysis_id, "not_found", "No saved analysis with that id exists for this signed-in user.")
        return CsvAnalysisResult.model_validate(record)

    @a2a.tool(
        description="List recent saved CSV analyses for the signed-in user.",
        timeout_seconds=15,
        idempotent=True,
        cost_class="cheap",
    )
    async def list_analyses(
        self,
        ctx: RunContext[PlatformUserAuth],
        limit: int = 10,
    ) -> CsvAnalysisListResult:
        tenant = _tenant_key(ctx)
        safe_limit = min(max(int(limit), 1), 50)
        return CsvAnalysisListResult(ok=True, analyses=_list_analyses(tenant, safe_limit))


def _tenant_key(ctx: RunContext[PlatformUserAuth]) -> str:
    stable_id = ctx.auth.user_id if ctx.auth.user_id is not None else ctx.auth.sub
    if not stable_id:
        raise PermissionError("stable platform identity required")
    return f"user:{stable_id}"


def _validate_analysis_id(analysis_id: str) -> str | None:
    cleaned = (analysis_id or "").strip()
    if not SAFE_ID_RE.fullmatch(cleaned):
        return "analysis_id must be 1-128 characters and contain only letters, numbers, dots, underscores, colons, or hyphens."
    return None


def _error(analysis_id: str, code: str, message: str) -> CsvAnalysisResult:
    return CsvAnalysisResult(ok=False, analysis_id=analysis_id, code=code, message=message, answer=None)


def _decode_browser_upload(upload: BrowserCsvUpload) -> str | CsvAnalysisResult:
    filename = upload.filename.replace("\\", "/").split("/")[-1]
    if not filename or filename in {".", ".."}:
        return _error("", "invalid_filename", "Upload filename is invalid.")
    if upload.media_type not in CSV_MEDIA_TYPES:
        return _error("", "unsupported_media_type", "Upload must be a CSV or plain text file.")
    try:
        raw = base64.b64decode(upload.data_base64, validate=True)
    except Exception:
        return _error("", "invalid_base64", "Upload data_base64 must be valid base64.")
    if len(raw) > MAX_CSV_BYTES:
        return _error("", "csv_too_large", f"CSV must be at most {MAX_CSV_BYTES} bytes.")
    try:
        return raw.decode("utf-8-sig")
    except UnicodeDecodeError:
        return _error("", "invalid_encoding", "CSV must be valid UTF-8 text.")


async def _read_uploaded_file(ctx: RunContext[PlatformUserAuth], file: UploadedFile) -> bytes:
    workspace = ctx.workspace
    reader = getattr(workspace, "read_bytes", None)
    if reader is not None:
        value = reader(file.path)
        if hasattr(value, "__await__"):
            return await value
        return value
    view = await workspace.open_view(
        purpose="Read uploaded CSV",
        hints=(file.path,),
        max_files=1,
        reason="Read the uploaded CSV staged for this invocation.",
    )
    return await view.read(file.path)


def _analyze_csv_text(*, analysis_id: str, csv_text: str, question: str) -> CsvAnalysisResult:
    id_error = _validate_analysis_id(analysis_id)
    if id_error:
        return _error(analysis_id, "invalid_analysis_id", id_error)
    if not question or not question.strip():
        return _error(analysis_id, "missing_question", "Question is required.")
    try:
        encoded = csv_text.encode("utf-8")
    except UnicodeEncodeError:
        return _error(analysis_id, "invalid_encoding", "CSV must be valid UTF-8 text.")
    if not encoded:
        return _error(analysis_id, "empty_csv", "CSV text is required.")
    if len(encoded) > MAX_CSV_BYTES:
        return _error(analysis_id, "csv_too_large", f"CSV must be at most {MAX_CSV_BYTES} bytes.")

    try:
        rows = _parse_csv_strict(csv_text)
    except csv.Error:
        return _error(analysis_id, "malformed_csv", "CSV is malformed: unterminated quoted field or invalid CSV structure.")
    if not rows:
        return _error(analysis_id, "empty_csv", "CSV must contain a header row and at least one data row.")
    header = rows[0]
    data_rows = rows[1:]
    if len(header) < 2 or not any(cell.strip() for cell in header):
        return _error(analysis_id, "missing_header", "CSV must include a header row with at least two columns.")
    if len(header) > MAX_COLUMNS:
        return _error(analysis_id, "too_many_columns", f"CSV may contain at most {MAX_COLUMNS} columns.")
    if not data_rows:
        return _error(analysis_id, "no_data_rows", "CSV must include at least one data row.")
    if len(data_rows) > MAX_ROWS:
        return _error(analysis_id, "too_many_rows", f"CSV may contain at most {MAX_ROWS} data rows.")
    if len(set(header)) != len(header):
        return _error(analysis_id, "duplicate_headers", "CSV headers must be unique.")
    for row in data_rows:
        if len(row) != len(header):
            return _error(analysis_id, "ragged_rows", "Every CSV row must have the same number of columns as the header.")
        if any(len(cell) > MAX_CELL_CHARS for cell in row):
            return _error(analysis_id, "cell_too_large", f"Each CSV cell must be at most {MAX_CELL_CHARS} characters.")

    records = [dict(zip(header, row, strict=True)) for row in data_rows]
    numeric_columns = _numeric_columns(records, header)
    if not numeric_columns:
        return _error(analysis_id, "no_numeric_values", "CSV must contain at least one numeric column.")
    numeric_column = _choose_numeric_column(question, numeric_columns)
    row_label_column = _choose_label_column(header, numeric_columns)

    best_record: dict[str, str] | None = None
    best_value: Decimal | None = None
    for record in records:
        value = _parse_number(record.get(numeric_column, ""))
        if value is None:
            continue
        if best_value is None or value > best_value:
            best_value = value
            best_record = record
    if best_record is None or best_value is None:
        return _error(analysis_id, "no_numeric_values", f"Column {numeric_column!r} does not contain numeric values.")

    label = best_record.get(row_label_column) or f"row {records.index(best_record) + 1}"
    value_text = _format_decimal(best_value)
    answer = f"{label} has the highest {numeric_column} at {value_text}."
    chart_data = []
    for idx, record in enumerate(records):
        value = _parse_number(record.get(numeric_column, ""))
        if value is not None:
            chart_data.append(
                {
                    "label": record.get(row_label_column) or f"row {idx + 1}",
                    "value": float(value),
                    "source_row_number": idx + 1,
                }
            )
    source_rows = [
        {
            "row_number": records.index(best_record) + 1,
            "values": best_record,
            "matched_column": numeric_column,
            "matched_value": float(best_value),
        }
    ]
    return CsvAnalysisResult(
        ok=True,
        analysis_id=analysis_id,
        answer=answer,
        row_count=len(records),
        column_count=len(header),
        question=question,
        source_rows=source_rows,
        chart_data=chart_data,
        numeric_column=numeric_column,
        row_label_column=row_label_column,
        saved=False,
    )


def _parse_csv_strict(csv_text: str) -> list[list[str]]:
    stream = io.StringIO(csv_text, newline="")
    sample = csv_text[:2048]
    try:
        dialect = csv.Sniffer().sniff(sample) if sample.strip() else csv.excel
    except csv.Error:
        dialect = csv.excel
    reader = csv.reader(stream, dialect=dialect, strict=True)
    rows = [[cell.strip() for cell in row] for row in reader]
    if csv_text.count('"') % 2:
        raise csv.Error("unterminated quoted field")
    return [row for row in rows if any(cell != "" for cell in row)]


def _parse_number(value: Any) -> Decimal | None:
    text = str(value).strip().replace(",", "")
    if not text:
        return None
    if text.endswith("%"):
        text = text[:-1]
    if text.startswith("$"):
        text = text[1:]
    try:
        number = Decimal(text)
    except (InvalidOperation, ValueError):
        return None
    return number if number.is_finite() else None


def _numeric_columns(records: list[dict[str, str]], header: list[str]) -> list[str]:
    out: list[str] = []
    for column in header:
        if any(_parse_number(record.get(column, "")) is not None for record in records):
            out.append(column)
    return out


def _choose_numeric_column(question: str, numeric_columns: list[str]) -> str:
    q = question.lower()
    for column in numeric_columns:
        normalized = column.lower().replace("_", " ").replace("-", " ")
        if normalized in q or column.lower() in q:
            return column
    return numeric_columns[0]


def _choose_label_column(header: list[str], numeric_columns: list[str]) -> str:
    for column in header:
        if column not in numeric_columns:
            return column
    return header[0]


def _format_decimal(value: Decimal) -> str:
    if value == value.to_integral_value():
        return str(value.quantize(Decimal(1)))
    return format(value.normalize(), "f")


def _connect():
    url = os.environ.get("DATABASE_URL")
    if not url:
        return None
    import psycopg

    return psycopg.connect(
        url,
        options="-c statement_timeout=10000 -c lock_timeout=5000 -c idle_in_transaction_session_timeout=30000",
    )


def _persist_analysis_and_receipt(
    *,
    tenant_key: str,
    analysis_id: str,
    question: str,
    csv_text: str,
    result: dict[str, Any],
    source: str,
) -> tuple[bool, str]:
    receipt_id = _receipt_id(tenant_key, analysis_id, csv_text, question)
    payload = dict(result)
    payload.update({"saved": True, "receipt_id": receipt_id})
    receipt = {
        "receipt_id": receipt_id,
        "analysis_id": analysis_id,
        "tenant_key_hash": hashlib.sha256(tenant_key.encode()).hexdigest(),
        "input_sha256": hashlib.sha256(csv_text.encode("utf-8")).hexdigest(),
        "question_sha256": hashlib.sha256(question.encode("utf-8")).hexdigest(),
        "source": source,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "ok": True,
    }
    conn = _connect()
    if conn is None:
        _LOCAL_ANALYSES[(tenant_key, analysis_id)] = payload
        _LOCAL_RECEIPTS.append(receipt)
        return True, receipt_id
    with conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO csv_analyses (tenant_key, analysis_id, question, answer, row_count, column_count, payload, updated_at)
                VALUES (%s, %s, %s, %s, %s, %s, %s::jsonb, NOW())
                ON CONFLICT (tenant_key, analysis_id) DO UPDATE SET
                    question = EXCLUDED.question,
                    answer = EXCLUDED.answer,
                    row_count = EXCLUDED.row_count,
                    column_count = EXCLUDED.column_count,
                    payload = EXCLUDED.payload,
                    updated_at = NOW()
                """,
                (
                    tenant_key,
                    analysis_id,
                    question,
                    payload.get("answer"),
                    int(payload.get("row_count") or 0),
                    int(payload.get("column_count") or 0),
                    json.dumps(payload),
                ),
            )
            cur.execute(
                """
                INSERT INTO execution_receipts (tenant_key, receipt_id, analysis_id, skill_name, input_hash, payload)
                VALUES (%s, %s, %s, %s, %s, %s::jsonb)
                ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET payload = EXCLUDED.payload
                """,
                (
                    tenant_key,
                    receipt_id,
                    analysis_id,
                    "analyze_csv",
                    receipt["input_sha256"],
                    json.dumps(receipt),
                ),
            )
    return True, receipt_id


def _load_analysis(tenant_key: str, analysis_id: str) -> dict[str, Any] | None:
    conn = _connect()
    if conn is None:
        return _LOCAL_ANALYSES.get((tenant_key, analysis_id))
    with conn:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT payload FROM csv_analyses WHERE tenant_key = %s AND analysis_id = %s",
                (tenant_key, analysis_id),
            )
            row = cur.fetchone()
    if not row:
        return None
    return row[0]


def _list_analyses(tenant_key: str, limit: int) -> list[dict[str, Any]]:
    conn = _connect()
    if conn is None:
        rows = [payload for (tenant, _), payload in _LOCAL_ANALYSES.items() if tenant == tenant_key]
        return [
            {
                "analysis_id": item["analysis_id"],
                "answer": item.get("answer"),
                "row_count": item.get("row_count", 0),
                "updated_at": item.get("updated_at"),
            }
            for item in rows[:limit]
        ]
    with conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT analysis_id, answer, row_count, updated_at
                FROM csv_analyses
                WHERE tenant_key = %s
                ORDER BY updated_at DESC
                LIMIT %s
                """,
                (tenant_key, limit),
            )
            rows = cur.fetchall()
    return [
        {
            "analysis_id": row[0],
            "answer": row[1],
            "row_count": row[2],
            "updated_at": row[3].isoformat() if hasattr(row[3], "isoformat") else str(row[3]),
        }
        for row in rows
    ]


def _receipt_id(tenant_key: str, analysis_id: str, csv_text: str, question: str) -> str:
    digest = hashlib.sha256(f"{tenant_key}\0{analysis_id}\0{question}\0{csv_text}".encode("utf-8")).hexdigest()
    return f"rcpt_{digest[:32]}"


def _reset_local_state_for_tests() -> None:
    _LOCAL_ANALYSES.clear()
    _LOCAL_RECEIPTS.clear()