"""
DocuLens - Supabase persistence layer.

Handles document storage, pipeline run tracking, and stats aggregation.
Falls back gracefully when Supabase is not configured (persist_results=False).
"""

import os
import logging
from datetime import datetime, timedelta, timezone
from typing import Optional
from uuid import uuid4

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Supabase client singleton
# ---------------------------------------------------------------------------

_client = None

SUPABASE_URL = os.environ.get("SUPABASE_URL", "")
SUPABASE_KEY = os.environ.get("SUPABASE_KEY", "")


def get_client():
    """Lazy-init Supabase client. Returns None if not configured."""
    global _client
    if _client is not None:
        return _client
    if not SUPABASE_URL or not SUPABASE_KEY:
        logger.info("Supabase not configured — persistence disabled")
        return None
    try:
        from supabase import create_client
        _client = create_client(SUPABASE_URL, SUPABASE_KEY)
        logger.info("Supabase client initialized")
        return _client
    except Exception as e:
        logger.warning(f"Failed to init Supabase: {e}")
        return None


def is_enabled() -> bool:
    """Check if persistence is available."""
    return get_client() is not None


# ---------------------------------------------------------------------------
# Documents
# ---------------------------------------------------------------------------

def save_document(
    filename: str,
    file_type: str,
    file_size: int,
    page_count: int = 1,
    org_id: Optional[str] = None,
    user_id: Optional[str] = None,
) -> Optional[str]:
    """Save document metadata. Returns document_id or None."""
    client = get_client()
    if not client:
        return None

    doc_id = str(uuid4())
    try:
        client.table("documents").insert({
            "id": doc_id,
            "org_id": org_id,
            "user_id": user_id,
            "filename": filename,
            "file_type": file_type,
            "file_size": file_size,
            "page_count": page_count,
        }).execute()
        return doc_id
    except Exception as e:
        logger.warning(f"Failed to save document: {e}")
        return None


# ---------------------------------------------------------------------------
# Pipeline Runs
# ---------------------------------------------------------------------------

def create_run(
    document_id: Optional[str],
    use_case: str = "invoice",
    org_id: Optional[str] = None,
    user_id: Optional[str] = None,
) -> Optional[str]:
    """Create a pipeline run. Returns run_id or None."""
    client = get_client()
    if not client:
        return None

    run_id = str(uuid4())
    try:
        client.table("pipeline_runs").insert({
            "id": run_id,
            "org_id": org_id,
            "user_id": user_id,
            "document_id": document_id,
            "use_case": use_case,
            "status": "running",
            "started_at": datetime.now(timezone.utc).isoformat(),
        }).execute()
        return run_id
    except Exception as e:
        logger.warning(f"Failed to create run: {e}")
        return None


def complete_run(
    run_id: str,
    status: str,
    overall_result: str,
    extraction_data: dict,
    processing_time_ms: int,
    error_message: Optional[str] = None,
) -> bool:
    """Mark a run as completed/failed with its results."""
    client = get_client()
    if not client or not run_id:
        return False

    try:
        client.table("pipeline_runs").update({
            "status": status,
            "overall_result": overall_result,
            "extraction_data": extraction_data,
            "processing_time_ms": processing_time_ms,
            "error_message": error_message,
            "completed_at": datetime.now(timezone.utc).isoformat(),
        }).eq("id", run_id).execute()
        return True
    except Exception as e:
        logger.warning(f"Failed to complete run: {e}")
        return False


def save_stage_result(
    run_id: str,
    stage_type: str,
    stage_name: str,
    status: str,
    sort_order: int,
    output: Optional[dict] = None,
    error_message: Optional[str] = None,
    duration_ms: Optional[int] = None,
) -> bool:
    """Save a single stage result within a run."""
    client = get_client()
    if not client or not run_id:
        return False

    try:
        client.table("stage_results").insert({
            "id": str(uuid4()),
            "run_id": run_id,
            "stage_type": stage_type,
            "stage_name": stage_name,
            "status": status,
            "sort_order": sort_order,
            "output": output,
            "error_message": error_message,
            "duration_ms": duration_ms,
            "started_at": datetime.now(timezone.utc).isoformat(),
            "completed_at": datetime.now(timezone.utc).isoformat(),
        }).execute()
        return True
    except Exception as e:
        logger.warning(f"Failed to save stage result: {e}")
        return False


# ---------------------------------------------------------------------------
# Queries — Runs listing + stats
# ---------------------------------------------------------------------------

def list_runs(
    org_id: Optional[str] = None,
    user_id: Optional[str] = None,
    status: Optional[str] = None,
    result: Optional[str] = None,
    use_case: Optional[str] = None,
    search: Optional[str] = None,
    page: int = 1,
    per_page: int = 25,
) -> dict:
    """List pipeline runs with filtering and pagination."""
    client = get_client()
    if not client:
        return {"runs": [], "pagination": {"page": 1, "per_page": per_page, "total": 0, "pages": 0}}

    try:
        # Select only the columns needed for the list view — deliberately
        # exclude extraction_data which contains large base64 page images.
        list_columns = (
            "id, document_id, use_case, status, overall_result, "
            "processing_time_ms, started_at, completed_at, error_message"
        )

        # Use !inner join when searching by filename so PostgREST can
        # filter on the related table; regular join otherwise.
        if search:
            doc_join = "documents!inner(filename, file_type, file_size)"
        else:
            doc_join = "documents(filename, file_type, file_size)"

        query = client.table("pipeline_runs").select(
            f"{list_columns}, {doc_join}",
            count="exact",
        )

        if user_id:
            query = query.eq("user_id", user_id)
        if org_id:
            query = query.eq("org_id", org_id)
        if status:
            query = query.eq("status", status)
        if result:
            query = query.eq("overall_result", result)
        if use_case:
            query = query.eq("use_case", use_case)
        if search:
            query = query.ilike("documents.filename", f"%{search}%")

        # Pagination
        offset = (page - 1) * per_page
        query = query.order("started_at", desc=True).range(offset, offset + per_page - 1)

        response = query.execute()
        total = response.count or 0

        # Batch-fetch stage_results for all runs in one query instead of N+1
        run_ids = [row["id"] for row in response.data]
        stages_by_run: dict[str, list] = {rid: [] for rid in run_ids}
        if run_ids:
            stages_resp = client.table("stage_results").select(
                "run_id, stage_type, stage_name, status, duration_ms, sort_order"
            ).in_("run_id", run_ids).order("sort_order").execute()
            for s in stages_resp.data:
                stages_by_run.setdefault(s["run_id"], []).append(s)

        runs = []
        for row in response.data:
            doc = row.get("documents") or {}
            runs.append({
                "id": row["id"],
                "document": {
                    "id": row.get("document_id"),
                    "filename": doc.get("filename", "unknown"),
                    "file_type": doc.get("file_type", ""),
                },
                "use_case": row.get("use_case", "invoice"),
                "status": row.get("status", "pending"),
                "overall_result": row.get("overall_result"),
                "stages": [
                    {
                        "type": s["stage_type"],
                        "name": s["stage_name"],
                        "status": s["status"],
                        "duration_ms": s.get("duration_ms"),
                    }
                    for s in stages_by_run.get(row["id"], [])
                ],
                "processing_time_ms": row.get("processing_time_ms"),
                "started_at": row.get("started_at"),
                "completed_at": row.get("completed_at"),
                "error_message": row.get("error_message"),
            })

        return {
            "runs": runs,
            "pagination": {
                "page": page,
                "per_page": per_page,
                "total": total,
                "pages": max(1, (total + per_page - 1) // per_page),
            },
        }
    except Exception as e:
        logger.warning(f"Failed to list runs: {e}")
        return {"runs": [], "pagination": {"page": 1, "per_page": per_page, "total": 0, "pages": 0}}


def get_run(run_id: str) -> Optional[dict]:
    """Get a single run with full details."""
    client = get_client()
    if not client:
        return None

    try:
        response = client.table("pipeline_runs").select(
            "*, documents(filename, file_type, file_size)"
        ).eq("id", run_id).single().execute()

        row = response.data
        if not row:
            return None

        stages_resp = client.table("stage_results").select("*").eq(
            "run_id", run_id
        ).order("sort_order").execute()

        doc = row.get("documents") or {}
        return {
            "id": row["id"],
            "document": {
                "id": row.get("document_id"),
                "filename": doc.get("filename", "unknown"),
                "file_type": doc.get("file_type", ""),
            },
            "use_case": row.get("use_case", "invoice"),
            "status": row.get("status"),
            "overall_result": row.get("overall_result"),
            "extraction_data": row.get("extraction_data"),
            "stages": [
                {
                    "type": s["stage_type"],
                    "name": s["stage_name"],
                    "status": s["status"],
                    "output": s.get("output"),
                    "error_message": s.get("error_message"),
                    "duration_ms": s.get("duration_ms"),
                }
                for s in stages_resp.data
            ],
            "processing_time_ms": row.get("processing_time_ms"),
            "started_at": row.get("started_at"),
            "completed_at": row.get("completed_at"),
        }
    except Exception as e:
        logger.warning(f"Failed to get run: {e}")
        return None


def get_stats(
    org_id: Optional[str] = None,
    user_id: Optional[str] = None,
    period_days: int = 30,
) -> dict:
    """Get aggregate stats for dashboard."""
    client = get_client()
    if not client:
        return _empty_stats(period_days)

    try:
        since = (datetime.now(timezone.utc) - timedelta(days=period_days)).isoformat()

        query = client.table("pipeline_runs").select("*").gte("started_at", since)
        if user_id:
            query = query.eq("user_id", user_id)
        if org_id:
            query = query.eq("org_id", org_id)

        response = query.execute()
        runs = response.data or []

        # Aggregate totals
        total = len(runs)
        passed = sum(1 for r in runs if r.get("overall_result") == "passed")
        failed = sum(1 for r in runs if r.get("overall_result") == "failed")
        needs_review = sum(1 for r in runs if r.get("overall_result") == "needs_review")
        times = [r.get("processing_time_ms", 0) for r in runs if r.get("processing_time_ms")]
        avg_ms = int(sum(times) / len(times)) if times else 0

        # Group by day
        by_day: dict[str, dict] = {}
        for r in runs:
            day = r.get("started_at", "")[:10]
            if day not in by_day:
                by_day[day] = {"date": day, "passed": 0, "failed": 0, "needs_review": 0}
            result = r.get("overall_result", "")
            if result in by_day[day]:
                by_day[day][result] += 1

        # Group by use case
        by_uc: dict[str, int] = {}
        for r in runs:
            uc = r.get("use_case", "invoice")
            by_uc[uc] = by_uc.get(uc, 0) + 1

        end_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
        start_date = (datetime.now(timezone.utc) - timedelta(days=period_days)).strftime("%Y-%m-%d")

        return {
            "period": {"start": start_date, "end": end_date},
            "totals": {
                "runs": total,
                "passed": passed,
                "failed": failed,
                "needs_review": needs_review,
                "avg_processing_ms": avg_ms,
            },
            "by_day": sorted(by_day.values(), key=lambda d: d["date"]),
            "by_use_case": [
                {"slug": slug, "name": slug.replace("-", " ").title(), "count": count}
                for slug, count in sorted(by_uc.items(), key=lambda x: -x[1])
            ],
        }
    except Exception as e:
        logger.warning(f"Failed to get stats: {e}")
        return _empty_stats(period_days)


def _empty_stats(period_days: int) -> dict:
    end_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    start_date = (datetime.now(timezone.utc) - timedelta(days=period_days)).strftime("%Y-%m-%d")
    return {
        "period": {"start": start_date, "end": end_date},
        "totals": {"runs": 0, "passed": 0, "failed": 0, "needs_review": 0, "avg_processing_ms": 0},
        "by_day": [],
        "by_use_case": [],
    }


# ---------------------------------------------------------------------------
# Golden Copies — user-verified extraction ground truth
# ---------------------------------------------------------------------------


def save_golden_copy(
    user_id: str,
    filename: str,
    extracted_fields: dict,
    line_items: list,
    use_case: str = "invoice",
    document_id: Optional[str] = None,
    run_id: Optional[str] = None,
    org_id: Optional[str] = None,
    notes: str = "",
) -> Optional[str]:
    """Save a user-verified extraction as a golden copy. Returns the golden copy id."""
    client = get_client()
    if not client:
        return None
    try:
        row = {
            "user_id": user_id,
            "filename": filename,
            "extracted_fields": extracted_fields,
            "line_items": line_items,
            "use_case": use_case,
            "notes": notes,
        }
        if document_id:
            row["document_id"] = document_id
        if run_id:
            row["run_id"] = run_id
        if org_id:
            row["org_id"] = org_id

        resp = client.table("golden_copies").insert(row).execute()
        return resp.data[0]["id"] if resp.data else None
    except Exception as e:
        logger.warning(f"Failed to save golden copy: {e}")
        return None


def list_golden_copies(
    user_id: str,
    use_case: Optional[str] = None,
    page: int = 1,
    per_page: int = 20,
) -> dict:
    """List golden copies for a user."""
    client = get_client()
    if not client:
        return {"golden_copies": [], "total": 0}
    try:
        query = (
            client.table("golden_copies")
            .select("id, filename, use_case, notes, created_at, updated_at", count="exact")
            .eq("user_id", user_id)
            .order("created_at", desc=True)
            .range((page - 1) * per_page, page * per_page - 1)
        )
        if use_case:
            query = query.eq("use_case", use_case)
        resp = query.execute()
        return {
            "golden_copies": resp.data or [],
            "total": resp.count or 0,
        }
    except Exception as e:
        logger.warning(f"Failed to list golden copies: {e}")
        return {"golden_copies": [], "total": 0}


def get_golden_copy(golden_id: str) -> Optional[dict]:
    """Retrieve a single golden copy by id."""
    client = get_client()
    if not client:
        return None
    try:
        resp = client.table("golden_copies").select("*").eq("id", golden_id).single().execute()
        return resp.data
    except Exception as e:
        logger.warning(f"Failed to get golden copy: {e}")
        return None


# ---------------------------------------------------------------------------
# Prompt Versions — CRUD for configurable prompts
# ---------------------------------------------------------------------------

def list_prompt_versions(
    use_case: str,
    stage: str,
    active_only: bool = False,
) -> list[dict]:
    """List prompt versions for a (use_case, stage) pair, newest first."""
    client = get_client()
    if not client:
        return []

    try:
        query = client.table("prompt_versions").select("*") \
            .eq("use_case", use_case) \
            .eq("stage", stage)
        if active_only:
            query = query.eq("is_active", True)
        query = query.order("version", desc=True)

        response = query.execute()
        return response.data or []
    except Exception as e:
        logger.warning(f"Failed to list prompt versions: {e}")
        return []


def get_active_prompt(use_case: str, stage: str) -> Optional[str]:
    """Get the currently active prompt text for a (use_case, stage) pair.

    Returns None if no active prompt is stored (caller should fall back
    to the hardcoded PROMPT_LIBRARY).
    """
    client = get_client()
    if not client:
        return None

    try:
        response = client.table("prompt_versions").select("prompt_text") \
            .eq("use_case", use_case) \
            .eq("stage", stage) \
            .eq("is_active", True) \
            .order("version", desc=True) \
            .limit(1) \
            .execute()
        rows = response.data or []
        if rows:
            return rows[0]["prompt_text"]
        return None
    except Exception as e:
        logger.warning(f"Failed to get active prompt: {e}")
        return None


def save_prompt_version(
    use_case: str,
    stage: str,
    prompt_text: str,
    created_by: Optional[str] = None,
) -> Optional[dict]:
    """Create a new prompt version, making it the active one.

    Deactivates any previously active version for the same (use_case, stage).
    Returns the new row or None on failure.
    """
    client = get_client()
    if not client:
        return None

    try:
        # Find current max version number
        existing = client.table("prompt_versions").select("version") \
            .eq("use_case", use_case) \
            .eq("stage", stage) \
            .order("version", desc=True) \
            .limit(1) \
            .execute()
        next_version = 1
        if existing.data:
            next_version = existing.data[0]["version"] + 1

        # Deactivate previous active version(s)
        client.table("prompt_versions").update({"is_active": False}) \
            .eq("use_case", use_case) \
            .eq("stage", stage) \
            .eq("is_active", True) \
            .execute()

        # Insert new version
        row = {
            "id": str(uuid4()),
            "use_case": use_case,
            "stage": stage,
            "version": next_version,
            "prompt_text": prompt_text,
            "is_active": True,
            "created_by": created_by,
            "created_at": datetime.now(timezone.utc).isoformat(),
        }
        client.table("prompt_versions").insert(row).execute()
        return row
    except Exception as e:
        logger.warning(f"Failed to save prompt version: {e}")
        return None


def activate_prompt_version(
    use_case: str,
    stage: str,
    version: int,
) -> bool:
    """Set a specific version as the active prompt."""
    client = get_client()
    if not client:
        return False

    try:
        # Deactivate all for this (use_case, stage)
        client.table("prompt_versions").update({"is_active": False}) \
            .eq("use_case", use_case) \
            .eq("stage", stage) \
            .execute()
        # Activate the requested version
        client.table("prompt_versions").update({"is_active": True}) \
            .eq("use_case", use_case) \
            .eq("stage", stage) \
            .eq("version", version) \
            .execute()
        return True
    except Exception as e:
        logger.warning(f"Failed to activate prompt version: {e}")
        return False


def get_all_active_prompts() -> list[dict]:
    """Get all active prompts across all use cases and stages."""
    client = get_client()
    if not client:
        return []

    try:
        response = client.table("prompt_versions").select("*") \
            .eq("is_active", True) \
            .order("use_case") \
            .order("stage") \
            .execute()
        return response.data or []
    except Exception as e:
        logger.warning(f"Failed to get all active prompts: {e}")
        return []