# VidContext MCP Server — wraps the VidContext API for use in Claude Desktop, Cursor, and Claude Code.

import asyncio
import os
import tempfile
from pathlib import Path

import httpx
from mcp.server.fastmcp import FastMCP, Context

# Track submitted jobs in-memory for session recovery
_recent_jobs: list[dict] = []
_MAX_TRACKED = 20

mcp = FastMCP(
    "vidcontext",
    instructions=(
        "Video Intelligence API — analyze any video and get detailed text output. "
        "This server runs LOCALLY and can read any file on the user's computer.\n\n"
        "IMPORTANT: The analyze_video tool uses a guided step-by-step workflow. "
        "Follow the tool's responses exactly — it will walk you through each step.\n\n"
        "Step 1: Call analyze_video() with no arguments to start\n"
        "Step 2: The tool will ask for the video file path — get it from the user\n"
        "        (accepts a single file, a URL, or a folder of up to 10 videos)\n"
        "Step 3: Call again with the file_path — the tool will show available modes\n"
        "        (for folders, it also lists all found videos with sizes)\n"
        "Step 4: Ask the user which mode they want, then call with both file_path and mode\n"
        "Step 5: Present the COMPLETE output to the user without summarizing\n"
        "        (for folders, show the summary table, then use check_job_status for each)\n\n"
        "Do NOT skip steps. Do NOT pick a mode for the user. Do NOT summarize the output.\n"
        "The detailed analysis is the product — always show it in full."
    ),
)

SUPPORTED_EXTENSIONS = {".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v", ".flv", ".wmv"}
MAX_POLL_SECONDS = 600
MAX_FILE_BYTES = 500 * 1024 * 1024
R2_THRESHOLD_BYTES = 95 * 1024 * 1024
MAX_UPLOAD_RETRIES = 2
RETRY_DELAY_SECONDS = 15
MAX_BATCH_SIZE = 10

VALID_MODES = ("context", "editor", "analysis", "ad", "ecommerce", "training", "ugc", "competitor")

MODE_MENU = (
    "Please ask the user which analysis mode they'd like to use:\n\n"
    "1. context    — Complete scene-by-scene description of the video (visuals, audio, text, pacing)\n"
    "2. editor     — Frame-by-frame breakdown with sub-second timestamps for video editing\n"
    "3. analysis   — Creator/social media performance analysis with scoring\n"
    "4. ad         — Ad effectiveness evaluation with creative scoring\n"
    "5. ecommerce  — Product video analysis for e-commerce optimization\n"
    "6. training   — Educational video effectiveness evaluation\n"
    "7. ugc        — Creator vetting and authenticity assessment for brand partnerships\n"
    "8. competitor  — Competitive intelligence and market positioning analysis\n\n"
    "Once the user picks a mode, call analyze_video again with both file_path and output_format."
)


def _get_api_url() -> str:
    return os.getenv("VIDCONTEXT_API_URL", "https://api.vidcontext.com")


def _headers() -> dict[str, str]:
    key = os.getenv("VIDCONTEXT_API_KEY", "")
    if not key:
        raise ValueError(
            "VIDCONTEXT_API_KEY not set. "
            "Get your API key at https://www.vidcontext.com/app/developer"
        )
    return {"X-API-Key": key}


async def _send_progress(ctx: Context, progress: int) -> None:
    """Send a progress heartbeat if context is available."""
    if ctx:
        try:
            await ctx.report_progress(progress=progress, total=MAX_POLL_SECONDS)
        except Exception:
            pass


async def _api_get(endpoint: str) -> dict:
    async with httpx.AsyncClient(timeout=30.0) as client:
        resp = await client.get(f"{_get_api_url()}{endpoint}", headers=_headers())
        resp.raise_for_status()
        return resp.json()


def _handle_error_status(resp: httpx.Response) -> dict | None:
    """Map HTTP error codes to user-friendly messages. Returns None on success."""
    if resp.status_code == 400:
        return {"error": resp.json().get("detail", "Bad request")}
    if resp.status_code == 402:
        return {"error": "Insufficient credits. Buy more at https://www.vidcontext.com"}
    if resp.status_code == 413:
        return {"error": "Video file too large (max 500MB)."}
    if resp.status_code == 429:
        return {"error": "Rate limited. Please wait a moment and try again.", "_retryable": True}
    if resp.status_code == 503:
        return {"error": "System is busy. Try again in a few minutes.", "_retryable": True}
    return None


async def _upload_video_direct(
    file_path: str, filename: str, file_size: int, output_format: str
) -> dict:
    """Direct file upload for files under 95MB."""
    write_timeout = max(60.0, 30.0 + file_size / (1024 * 1024))
    upload_timeout = httpx.Timeout(connect=10.0, read=300.0, write=write_timeout, pool=30.0)

    async with httpx.AsyncClient(timeout=upload_timeout) as client:
        with open(file_path, "rb") as f:
            resp = await client.post(
                f"{_get_api_url()}/v1/analyze",
                headers=_headers(),
                files={"file": (filename, f)},
                data={"output_format": output_format},
            )

        error = _handle_error_status(resp)
        if error:
            return error

        resp.raise_for_status()
        return resp.json()


async def _upload_video_r2(
    file_path: str, filename: str, file_size: int, output_format: str
) -> dict:
    """Upload large files via R2 presigned URL, then trigger analysis."""
    headers = _headers()

    async with httpx.AsyncClient(timeout=30.0) as client:
        presign_resp = await client.post(
            f"{_get_api_url()}/v1/upload/presign",
            headers={**headers, "Content-Type": "application/json"},
            json={"filename": filename, "file_size": file_size},
        )

        if presign_resp.status_code != 200:
            detail = presign_resp.json().get("detail", "Failed to prepare upload")
            return {"error": detail}

    presign_data = presign_resp.json()
    upload_url = presign_data["upload_url"]
    r2_key = presign_data["r2_key"]

    write_timeout = max(120.0, file_size / (1024 * 1024))
    r2_timeout = httpx.Timeout(connect=10.0, read=60.0, write=write_timeout, pool=30.0)

    # Read file into memory for async upload (file objects can't be sent via AsyncClient)
    file_data = Path(file_path).read_bytes()

    async with httpx.AsyncClient(timeout=r2_timeout) as client:
        r2_resp = await client.put(upload_url, content=file_data)

    del file_data  # free memory immediately

    if r2_resp.status_code not in (200, 201):
        return {"error": f"Upload to storage failed ({r2_resp.status_code})"}

    async with httpx.AsyncClient(timeout=30.0) as client:
        analyze_resp = await client.post(
            f"{_get_api_url()}/v1/analyze",
            headers=headers,
            data={"r2_key": r2_key, "output_format": output_format},
        )

        error = _handle_error_status(analyze_resp)
        if error:
            return error

        analyze_resp.raise_for_status()
        return analyze_resp.json()


async def _upload_video(
    file_path: str, output_format: str, ctx: Context = None
) -> dict:
    """Upload video with automatic retry on rate limits."""
    path = Path(file_path)
    filename = path.name
    file_size = path.stat().st_size

    if file_size > MAX_FILE_BYTES:
        return {"error": f"File too large ({file_size // (1024 * 1024)}MB). Max 500MB."}

    upload_fn = _upload_video_r2 if file_size > R2_THRESHOLD_BYTES else _upload_video_direct
    args = (file_path, filename, file_size, output_format)

    for attempt in range(MAX_UPLOAD_RETRIES + 1):
        # Send heartbeat before upload to reset client timeout
        await _send_progress(ctx, 1 + attempt)

        try:
            result = await upload_fn(*args)
        except httpx.HTTPStatusError as exc:
            return {"error": f"Upload failed (HTTP {exc.response.status_code})"}
        except httpx.HTTPError as exc:
            return {"error": f"Upload failed: {exc}"}

        # Success or non-retryable error — return immediately
        if not result.get("_retryable") or attempt >= MAX_UPLOAD_RETRIES:
            result.pop("_retryable", None)
            return result

        # Retryable error — wait before next attempt with periodic heartbeats
        delay = RETRY_DELAY_SECONDS * (attempt + 1)
        for i in range(0, delay, 5):
            await asyncio.sleep(min(5, delay - i))
            await _send_progress(ctx, 5 + i)

    result.pop("_retryable", None)
    return result


async def _poll_job(job_id: str, ctx: Context = None) -> dict:
    """Poll job status until complete, failed, or timeout."""
    elapsed = 0
    interval = 3

    while elapsed < MAX_POLL_SECONDS:
        await asyncio.sleep(interval)
        elapsed += interval

        # Send heartbeat to prevent client timeout
        await _send_progress(ctx, elapsed)

        try:
            result = await _api_get(f"/v1/status/{job_id}")
        except httpx.HTTPStatusError as exc:
            if exc.response.status_code in (401, 403):
                return {"status": "failed", "error": "Authentication failed. Check your API key."}
            if exc.response.status_code >= 500:
                continue
            return {"status": "failed", "error": f"HTTP {exc.response.status_code}"}
        except httpx.HTTPError:
            continue

        if result.get("status") in ("complete", "failed"):
            return result

        # Gradually back off: 3s → 5s → 7s → 10s max
        interval = min(interval + 2, 10)

    return {"status": "timeout", "job_id": job_id}


async def _download_url(url: str) -> str:
    """Download a video from URL to a temp file."""
    url_path = Path(url.split("?")[0])
    ext = url_path.suffix.lower() if url_path.suffix else ".mp4"
    if ext not in SUPPORTED_EXTENSIONS:
        ext = ".mp4"

    tmp = tempfile.NamedTemporaryFile(suffix=ext, delete=False)
    try:
        async with httpx.AsyncClient(
            timeout=httpx.Timeout(10.0, read=120.0),
            follow_redirects=True,
        ) as client:
            async with client.stream("GET", url) as resp:
                resp.raise_for_status()

                content_length = resp.headers.get("content-length")
                if content_length and int(content_length) > MAX_FILE_BYTES:
                    raise ValueError(
                        f"Remote file too large ({int(content_length) // (1024 * 1024)}MB). Max 500MB."
                    )

                downloaded = 0
                async for chunk in resp.aiter_bytes(chunk_size=65536):
                    downloaded += len(chunk)
                    if downloaded > MAX_FILE_BYTES:
                        raise ValueError("Remote file exceeds 500MB limit.")
                    tmp.write(chunk)

        tmp.close()
        return tmp.name
    except Exception:
        tmp.close()
        Path(tmp.name).unlink(missing_ok=True)
        raise


def _find_videos(folder: Path) -> list[Path]:
    """Find all video files in a folder (non-recursive, sorted by name)."""
    return sorted(
        p for p in folder.iterdir()
        if p.is_file() and p.suffix.lower() in SUPPORTED_EXTENSIONS
    )


async def _process_single_video(
    file_path: str, output_format: str, ctx: Context = None
) -> dict:
    """Process one video end-to-end. Returns dict with status, result text, timing."""
    temp_file = None
    try:
        if file_path.startswith(("http://", "https://")):
            await _send_progress(ctx, 1)
            try:
                temp_file = await _download_url(file_path)
            except httpx.HTTPError as exc:
                return {"status": "failed", "error": f"Download error: {exc}"}
            except ValueError as exc:
                return {"status": "failed", "error": str(exc)}
            upload_path = temp_file
        else:
            upload_path = str(Path(file_path).resolve())

        result = await _upload_video(upload_path, output_format, ctx)

        if "error" in result:
            return {"status": "failed", "error": result["error"]}

        job_id = result["job_id"]
        _track_job(job_id, Path(upload_path).name, output_format)

        final = await _poll_job(job_id, ctx)

        if final["status"] == "complete":
            _update_job_status(job_id, "complete")
            return {
                "status": "complete",
                "job_id": job_id,
                "result": final.get("result", ""),
                "time_s": final.get("processing_time_seconds"),
            }

        if final["status"] == "failed":
            _update_job_status(job_id, "failed")
            return {"status": "failed", "job_id": job_id, "error": final.get("error", "Unknown")}

        _update_job_status(job_id, "processing")
        return {"status": "timeout", "job_id": job_id}

    finally:
        if temp_file:
            Path(temp_file).unlink(missing_ok=True)


def _track_job(job_id: str, filename: str, output_format: str) -> None:
    """Track a submitted job for session recovery."""
    _recent_jobs.insert(0, {
        "job_id": job_id,
        "file": filename,
        "mode": output_format,
        "status": "submitted",
    })
    while len(_recent_jobs) > _MAX_TRACKED:
        _recent_jobs.pop()


def _update_job_status(job_id: str, status: str) -> None:
    """Update the tracked status of a job."""
    for job in _recent_jobs:
        if job["job_id"] == job_id:
            job["status"] = status
            break


@mcp.tool()
async def analyze_video(
    file_path: str = "",
    output_format: str = "",
    ctx: Context = None,
) -> str:
    """Analyze a video file (or folder of videos) using VidContext.

    This tool uses a guided workflow — call it step by step:

    Step 1: Call with no arguments to start.
    Step 2: Call with file_path once the user provides the video location.
            Accepts a single file, a URL, or a folder containing videos.
    Step 3: Call with file_path AND output_format after the user picks a mode.

    The tool will guide you through each step. Do not skip steps or pick
    a mode on behalf of the user — always ask them.

    Folders: If the user points to a folder, the tool will find all videos,
    show the list, and process them one at a time (max 10 per batch).

    Processing takes 30-180 seconds per video. Progress is tracked automatically.
    Cost: 1 credit per minute of video (rounded up).

    Args:
        file_path: Path to a video file, a URL, or a folder of videos (ask the user)
        output_format: Analysis mode (ask the user — do not pick for them)
    """

    # --- Step 1: No file path provided — ask the user for the video location ---
    if not file_path:
        return (
            "To analyze a video, I need the file location.\n\n"
            "Please ask the user for:\n"
            "- The full path to a video file (e.g. /Users/name/Downloads/video.mp4)\n"
            "- A video URL\n"
            "- A folder path to analyze multiple videos at once (max 10)\n\n"
            "Supported formats: "
            f"{', '.join(sorted(SUPPORTED_EXTENSIONS))}.\n"
            "Max file size: 500MB per video. Max duration: 15 minutes per video.\n\n"
            "Once you have the path, call analyze_video(file_path='...') again."
        )

    # --- Detect if it's a URL, file, or folder ---
    is_url = file_path.startswith(("http://", "https://"))
    is_folder = False
    videos_in_folder: list[Path] = []

    if not is_url:
        path = Path(file_path).resolve()
        if not path.exists():
            return (
                f"File not found: {file_path}\n\n"
                "Please ask the user to double-check the path and try again."
            )

        if path.is_dir():
            is_folder = True
            videos_in_folder = _find_videos(path)

            if not videos_in_folder:
                return (
                    f"No video files found in: {file_path}\n\n"
                    f"Supported formats: {', '.join(sorted(SUPPORTED_EXTENSIONS))}\n\n"
                    "Please ask the user for a folder that contains video files."
                )

            if len(videos_in_folder) > MAX_BATCH_SIZE:
                return (
                    f"Found {len(videos_in_folder)} videos in that folder, "
                    f"but the max per batch is {MAX_BATCH_SIZE}.\n\n"
                    "Please ask the user to select a smaller set, or move up to "
                    f"{MAX_BATCH_SIZE} videos into a separate folder."
                )

        elif path.is_file():
            if path.suffix.lower() not in SUPPORTED_EXTENSIONS:
                return (
                    f"Unsupported format '{path.suffix}'.\n"
                    f"Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}\n\n"
                    "Please ask the user for a video in one of these formats."
                )
        else:
            return (
                f"Not a file or folder: {file_path}\n\n"
                "Please ask the user for a valid path."
            )

    # --- Step 2: Path given but no mode — show what we found + mode menu ---
    if not output_format:
        if is_folder:
            file_list = "\n".join(
                f"  {i+1}. {v.name} ({v.stat().st_size // (1024*1024)}MB)"
                for i, v in enumerate(videos_in_folder)
            )
            return (
                f"Found {len(videos_in_folder)} video(s) in that folder:\n\n"
                f"{file_list}\n\n"
                f"All {len(videos_in_folder)} will be analyzed one at a time using the same mode.\n"
                "Cost: 1 credit per minute of video (rounded up).\n\n"
                f"{MODE_MENU}"
            )
        return MODE_MENU

    # --- Validate the mode ---
    if output_format not in VALID_MODES:
        return f"Invalid mode '{output_format}'.\n\n{MODE_MENU}"

    # --- Step 3: Both path and mode provided — run the analysis ---

    # === FOLDER BATCH MODE ===
    if is_folder:
        completed = 0
        failed_list: list[str] = []
        results_summary: list[str] = []

        for i, video_path in enumerate(videos_in_folder):
            await _send_progress(ctx, i * 100)

            outcome = await _process_single_video(str(video_path), output_format, ctx)

            if outcome["status"] == "complete":
                completed += 1
                job_id = outcome["job_id"]
                time_s = outcome.get("time_s")
                time_str = f" ({time_s:.0f}s)" if time_s else ""
                results_summary.append(
                    f"  {i+1}. {video_path.name} — done{time_str} — "
                    f"job_id: {job_id}"
                )
            elif outcome["status"] == "timeout":
                job_id = outcome.get("job_id", "unknown")
                results_summary.append(
                    f"  {i+1}. {video_path.name} — still processing — "
                    f"job_id: {job_id}"
                )
            else:
                failed_list.append(video_path.name)
                error = outcome.get("error", "Unknown error")
                results_summary.append(
                    f"  {i+1}. {video_path.name} — FAILED: {error}"
                )

        summary_text = "\n".join(results_summary)
        header = (
            f"[Batch complete: {completed}/{len(videos_in_folder)} succeeded"
        )
        if failed_list:
            header += f", {len(failed_list)} failed"
        header += "]\n\n"

        footer = (
            "\n\nUse check_job_status(job_id) to view the full analysis "
            "for any completed video."
        )

        return f"{header}{summary_text}{footer}"

    # === SINGLE FILE MODE ===
    outcome = await _process_single_video(file_path, output_format, ctx)

    if outcome["status"] == "complete":
        text = outcome.get("result", "")
        time_s = outcome.get("time_s")
        header = (
            "[IMPORTANT: Present the COMPLETE output below to the user. "
            "Do not summarize, truncate, or rephrase — show it exactly as-is.]\n\n"
        )
        footer = "\n\n---\n"
        if time_s:
            footer += f"Processed in {time_s:.1f}s. "
        footer += "1 credit per minute of video (rounded up)."
        return f"{header}{text}{footer}"

    if outcome["status"] == "failed":
        return f"Analysis failed: {outcome.get('error', 'Unknown error')}"

    job_id = outcome.get("job_id", "unknown")
    return (
        f"Still processing after {MAX_POLL_SECONDS}s. "
        f"Use check_job_status('{job_id}') to get the result when ready."
    )


@mcp.tool()
async def check_job_status(job_id: str) -> str:
    """Check the status of a video analysis job or retrieve its results.

    Args:
        job_id: The job ID from analyze_video or list_recent_jobs
    """
    try:
        result = await _api_get(f"/v1/status/{job_id}")
    except httpx.HTTPStatusError as exc:
        if exc.response.status_code == 404:
            return f"Job '{job_id}' not found. Jobs expire after 24 hours."
        return f"Error: HTTP {exc.response.status_code}"
    except httpx.HTTPError as exc:
        return f"Error: {exc}"
    except ValueError as exc:
        return str(exc)

    status = result.get("status", "unknown")

    if status == "complete":
        _update_job_status(job_id, "complete")
        text = result.get("result", "No output.")
        time_s = result.get("processing_time_seconds")
        fmt = result.get("output_format", "unknown")
        header = f"[Complete — {fmt} mode"
        if time_s:
            header += f" | {time_s:.1f}s"
        header += "]\n\n"
        return f"{header}{text}"

    if status == "failed":
        _update_job_status(job_id, "failed")
        return f"Failed: {result.get('error', 'Unknown error')}"

    if status == "processing":
        return "Still processing. Check again in 15-30 seconds."

    if status == "queued":
        return "Queued — waiting to start. Check again in 15-30 seconds."

    return f"Status: {status}"


@mcp.tool()
async def list_recent_jobs() -> str:
    """List video analysis jobs submitted during this session.

    Shows job IDs, modes, and filenames. Use check_job_status(job_id)
    to get results or check progress.
    """
    if not _recent_jobs:
        return "No jobs submitted yet this session."

    lines = []
    for job in _recent_jobs:
        status = job.get("status", "submitted")
        lines.append(f"  {job['job_id']} | {job['mode']} | {job['file']} | {status}")

    return "Recent jobs:\n" + "\n".join(lines) + "\n\nUse check_job_status(job_id) to get full results."


@mcp.tool()
async def check_credits() -> str:
    """Check your VidContext credit balance and usage limits."""
    try:
        usage = await _api_get("/v1/usage")
    except httpx.HTTPStatusError as exc:
        return f"Error: HTTP {exc.response.status_code}"
    except httpx.HTTPError as exc:
        return f"Error: {exc}"
    except ValueError as exc:
        return str(exc)

    tier = usage.get("tier", "unknown")
    credits = usage.get("credits", 0)

    if tier == "credit":
        return (
            f"Credits: {credits}\n"
            f"Tier: Paid\n"
            f"Max file size: {usage.get('max_file_size_mb', 500)}MB\n"
            f"Max duration: {usage.get('max_duration_seconds', 900) // 60} minutes\n"
            f"Cost: 1 credit per minute of video (rounded up)"
        )

    remaining = usage.get("remaining", 0)
    return (
        f"Free uses remaining: {remaining} of {usage.get('limit', 5)}\n"
        f"Max file size: {usage.get('max_file_size_mb', 100)}MB\n"
        f"Max duration: {usage.get('max_duration_seconds', 60)}s\n"
        f"Get more credits at https://www.vidcontext.com"
    )


@mcp.tool()
async def get_account() -> str:
    """Get your VidContext account information."""
    try:
        me = await _api_get("/v1/me")
    except httpx.HTTPStatusError as exc:
        if exc.response.status_code == 401:
            return "Invalid API key. Check your VIDCONTEXT_API_KEY."
        return f"Error: HTTP {exc.response.status_code}"
    except httpx.HTTPError as exc:
        return f"Error: {exc}"
    except ValueError as exc:
        return str(exc)

    if not me.get("authenticated"):
        return "Not authenticated. Set VIDCONTEXT_API_KEY in your MCP config."

    lines = [
        f"Email: {me.get('email', 'N/A')}",
        f"Name: {me.get('display_name', 'N/A')}",
        f"Credits: {me.get('credits', 0)}",
        f"Tier: {me.get('tier', 'N/A')}",
    ]

    sub = me.get("subscription")
    if sub:
        lines.extend([
            f"\nSubscription: {sub.get('plan', 'N/A').title()} ({sub.get('billing_interval', 'N/A')})",
            f"Credits/period: {sub.get('credits_per_period', 'N/A')}",
            f"Status: {sub.get('status', 'N/A')}",
            f"Renews: {sub.get('current_period_end', 'N/A')}",
        ])
    else:
        lines.append("\nNo active subscription")

    return "\n".join(lines)


def main():
    """Entry point for the VidContext MCP server."""
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
