# New file: api/results.py
from fastapi import APIRouter, Depends, HTTPException, Response
from fastapi.responses import JSONResponse
from typing import Dict, Any, Optional
import pandas as pd
import io
import os
import logging
from uuid import UUID as _UUID
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select

from ..auth import get_current_account
from ..storage import get_storage_backend
from ..models import Job, JobStatus
from ..db import get_session

router = APIRouter(prefix="/api/v1/jobs", tags=["Results"])

log = logging.getLogger(__name__)


async def _get_job_by_id(db: AsyncSession, job_id: str) -> Optional[Job]:
    """Lookup a job by string id first, then UUID fallback."""
    # Prefer direct string match (supports non-UUID ids like "api_...")
    job = await db.scalar(select(Job).where(Job.id == job_id))
    if job:
        return job
    # Legacy/UUID fallback
    try:
        return await db.get(Job, _UUID(job_id))
    except Exception:
        return None


@router.get("/{job_id}/result", summary="Get result download URL")
async def get_result_url(
    job_id: str,
    format: Optional[str] = "parquet",
    account_id: _UUID = Depends(get_current_account),
    db: AsyncSession = Depends(get_session),
) -> Dict[str, Any]:
    """
    Get pre-signed URL for downloading results.
    Supports multiple formats: parquet, csv, json
    """
    # Validate job ownership
    job = await _get_job_by_id(db, job_id)

    if not job:
        raise HTTPException(404, detail={"error": "job_not_found", "job_id": job_id})

    if str(job.account_id) != str(account_id):
        raise HTTPException(403, "Access denied")

    if job.status != JobStatus.COMPLETED:
        raise HTTPException(
            400, f"Results not available. Job status: {job.status.value}"
        )

    if not job.results_path:
        raise HTTPException(404, "Results file not found")

    # Get storage backend
    storage = get_storage_backend()
    # Initialize results_key for later size checks
    results_key = job.results_path

    # For S3, generate pre-signed URL
    if hasattr(storage, "generate_presigned_url"):
        # Extract the S3 key from results_path
        # Handle both s3:// URLs and plain keys
        if results_key.startswith("s3://"):
            # Parse s3://bucket/key format
            parts = results_key.replace("s3://", "").split("/", 1)
            if len(parts) > 1:
                results_key = parts[1]

        # Generate different formats if requested
        if format != "parquet":
            # Convert on-demand (cached for performance)
            converted_key = await convert_results_format(
                storage, results_key, format, job_id
            )
            results_key = converted_key

        download_url = await storage.generate_presigned_url(
            results_key,
            expires_in=86400,  # 24 hours
        )
    else:
        # Local storage - use direct endpoint
        download_url = f"/api/v1/jobs/{job_id}/download?format={format}"

    out = {
        "job_id": job_id,
        "download_url": download_url,
        "format": format,
        "expires_in_seconds": 86400,
        "size_bytes": await storage.get_file_size(results_key)
        if hasattr(storage, "get_file_size")
        else None,
        "stats": {
            "matches_found": job.matches_found,
            "processing_time_seconds": (
                (job.updated_at - job.created_at).total_seconds()
                if job.updated_at and job.created_at
                else None
            ),
            "completed_at": job.updated_at.isoformat() if job.updated_at else None,
        },
    }
    return JSONResponse(out, headers={"Cache-Control": "no-store"})


@router.get("/{job_id}/download", include_in_schema=False)
async def download_results_direct(
    job_id: str,
    format: str = "parquet",
    account_id: _UUID = Depends(get_current_account),
    db: AsyncSession = Depends(get_session),
) -> Response:
    """
    Direct download endpoint for local storage.
    In production, this redirects to pre-signed S3 URL.
    """
    # Validate job
    job = await _get_job_by_id(db, job_id)

    if not job or str(job.account_id) != str(account_id):
        raise HTTPException(403, "Access denied")

    if not job.results_path:
        raise HTTPException(404, "Results not found")

    storage = get_storage_backend()

    # Read the parquet file
    file_content = await storage.read(job.results_path)
    df = pd.read_parquet(io.BytesIO(file_content))

    # Convert to requested format
    if format == "csv":
        output = io.StringIO()
        df.to_csv(output, index=False)
        content = output.getvalue().encode()
        media_type = "text/csv"
        filename = f"results_{job_id}.csv"
    elif format == "json":
        content = df.to_json(orient="records").encode()
        media_type = "application/json"
        filename = f"results_{job_id}.json"
    else:  # parquet
        content = file_content
        media_type = "application/octet-stream"
        filename = f"results_{job_id}.parquet"

    return Response(
        content=content,
        media_type=media_type,
        headers={
            "Content-Disposition": f"attachment; filename={filename}",
            "Cache-Control": "no-store",
        },
    )


@router.get("/{job_id}/results", include_in_schema=False)
async def get_result_url_alias(
    job_id: str,
    format: Optional[str] = "parquet",
    account_id: _UUID = Depends(get_current_account),
    db: AsyncSession = Depends(get_session),
):
    # Alias to maintain compatibility with callers using "/results"
    return await get_result_url(job_id, format, account_id, db)


@router.get("/{job_id:path}/preview", summary="Preview results")
async def preview_results(
    job_id: str,
    limit: int = 100,
    offset: int = 0,
    account_id: _UUID = Depends(get_current_account),
    db: AsyncSession = Depends(get_session),
) -> Dict[str, Any]:
    """
    Preview results without downloading full file.
    Returns paginated JSON response.
    """
    try:
        # Validate job
        job = await _get_job_by_id(db, job_id)
        if not job:
            raise HTTPException(
                status_code=404, detail={"error": "job_not_found", "job_id": job_id}
            )

        if str(job.account_id) != str(account_id):
            # Dev guard: allow preview when job has no account and DEV_ACCOUNT_ID matches
            if os.getenv("ENV", "development") == "development":
                dev_id = os.getenv("DEV_ACCOUNT_ID")
                if (not job.account_id) and dev_id and str(account_id) == str(dev_id):
                    pass
                else:
                    log.warning(
                        "preview ownership mismatch job=%s job_acct=%s req_acct=%s",
                        job_id,
                        job.account_id,
                        account_id,
                    )
                    raise HTTPException(403, "Access denied")
            else:
                log.warning(
                    "preview ownership mismatch job=%s job_acct=%s req_acct=%s",
                    job_id,
                    job.account_id,
                    account_id,
                )
                raise HTTPException(403, "Access denied")

        # If not completed yet, return empty preview to keep clients happy
        status_str = (
            job.status.value if hasattr(job.status, "value") else str(job.status)
        ).lower()
        if status_str != "completed":
            return JSONResponse(
                {
                    "job_id": job.id,
                    "status": status_str,
                    "preview": {"columns": [], "rows": []},
                },
                headers={"Cache-Control": "no-store"},
            )

        # Clamp pagination
        limit = max(1, min(int(limit or 1), 1_000_000))
        offset = max(int(offset or 0), 0)

        # If no artifact or explicitly zero matches, return empty preview
        if not getattr(job, "results_path", None) or (
            getattr(job, "matches_found", None) == 0
        ):
            cols, rows = [], []
        else:
            storage = get_storage_backend()
            cols, rows = [], []
            try:
                # Read artifact (parquet or csv)
                path = job.results_path
                content = await storage.read(path)
                ext = os.path.splitext(path)[1].lower()
                if ext == ".parquet":
                    df = pd.read_parquet(io.BytesIO(content))
                elif ext == ".csv":
                    df = pd.read_csv(io.BytesIO(content))
                else:
                    # Try parquet by default
                    df = pd.read_parquet(io.BytesIO(content))

                # Apply pagination
                total_rows = len(df)
                if offset < 0:
                    offset = 0
                if limit <= 0:
                    limit = 100
                df_page = df.iloc[offset : offset + limit]
                cols = list(df.columns)
                rows = df_page.astype(object).values.tolist()
            except FileNotFoundError:
                cols, rows = [], []

            # Fallback to inline JSON results if present and usable
            if (
                not cols
                and not rows
                and isinstance(getattr(job, "results", None), dict)
            ):
                res = job.results or {}
                cols = list(res.get("columns") or [])
                raw_rows = res.get("rows") or []
                if isinstance(raw_rows, list):
                    if raw_rows and not isinstance(raw_rows[0], list):
                        # Normalize list of dicts to columns/rows
                        cols = cols or sorted(
                            {
                                k
                                for r in raw_rows
                                for k in (r.keys() if isinstance(r, dict) else [])
                            }
                        )
                        rows = [
                            [(r.get(c) if isinstance(r, dict) else None) for c in cols]
                            for r in raw_rows
                        ][:limit]
                    else:
                        rows = raw_rows[:limit]

        return JSONResponse(
            {"job_id": job.id, "preview": {"columns": cols, "rows": rows}},
            headers={"Cache-Control": "no-store"},
        )
    except HTTPException:
        raise
    except Exception as e:
        # Log full traceback and return JSON error
        log.exception("preview failed for job_id=%s", job_id)
        raise HTTPException(
            status_code=500, detail={"error": "preview_failed", "message": str(e)}
        )


async def convert_results_format(
    storage, source_key: str, target_format: str, job_id: str
) -> str:
    """
    Convert results to different format and cache.
    """
    # Generate cache key
    cache_key = source_key.replace(".parquet", f".{target_format}")

    # Check if already converted
    if await storage.exists(cache_key):
        return cache_key

    # Read and convert
    content = await storage.read(source_key)
    df = pd.read_parquet(io.BytesIO(content))

    # Convert based on format
    if target_format == "csv":
        output = io.StringIO()
        df.to_csv(output, index=False)
        converted_content = output.getvalue().encode()
    elif target_format == "json":
        converted_content = df.to_json(orient="records", indent=2).encode()
    else:
        raise ValueError(f"Unsupported format: {target_format}")

    # Save converted file
    await storage.save_bytes(converted_content, cache_key)

    return cache_key
