import base64
import io
import os
import re
import shutil
import tempfile
import tarfile
import zipfile
import gzip
import urllib.request
import urllib.error
from urllib.parse import urlparse
from datetime import datetime
from typing import List
from mcp.server.fastmcp import Context, FastMCP
from pdf2image import convert_from_path, pdfinfo_from_path

# Initialize MCP Server
mcp = FastMCP("PDF-Vision-Inspector")

def _extract_arxiv_id(arxiv_url: str) -> str:
    text = (arxiv_url or "").strip()
    if not text:
        return ""

    new_style = r"^\d{4}\.\d{4,5}(v\d+)?$"
    old_style = r"^[a-zA-Z-]+(\.[A-Z]{2})?/\d{7}(v\d+)?$"
    if re.match(new_style, text) or re.match(old_style, text):
        return text

    try:
        parsed = urlparse(text)
    except Exception:
        return ""

    path = parsed.path.strip("/")
    if not path:
        return ""

    parts = path.split("/")
    if parts[0] in ("abs", "pdf", "e-print", "src", "format"):
        arxiv_id = "/".join(parts[1:])
    else:
        arxiv_id = parts[0]

    arxiv_id = re.sub(r"\.(pdf|ps|gz|zip)$", "", arxiv_id, flags=re.IGNORECASE)
    if re.match(new_style, arxiv_id) or re.match(old_style, arxiv_id):
        return arxiv_id
    return ""

def _filename_from_content_disposition(header_value: str) -> str:
    if not header_value:
        return ""
    match = re.search(r'filename="?([^";]+)"?', header_value)
    return match.group(1) if match else ""

def _is_gzip_file(file_path: str) -> bool:
    try:
        with open(file_path, "rb") as f:
            return f.read(2) == b"\x1f\x8b"
    except Exception:
        return False

def _safe_extract_tar(tar: tarfile.TarFile, target_dir: str) -> None:
    target_dir_real = os.path.realpath(target_dir)
    for member in tar.getmembers():
        member_path = os.path.realpath(os.path.join(target_dir, member.name))
        if not member_path.startswith(target_dir_real + os.sep):
            raise ValueError(f"Unsafe path in tar file: {member.name}")
    tar.extractall(target_dir, filter="data")

def _safe_extract_zip(zip_file: zipfile.ZipFile, target_dir: str) -> None:
    target_dir_real = os.path.realpath(target_dir)
    for member in zip_file.namelist():
        member_path = os.path.realpath(os.path.join(target_dir, member))
        if not member_path.startswith(target_dir_real + os.sep):
            raise ValueError(f"Unsafe path in zip file: {member}")
    zip_file.extractall(target_dir)

def _ensure_unique_dir(base_dir: str) -> str:
    if not os.path.exists(base_dir):
        os.makedirs(base_dir, exist_ok=True)
        return base_dir
    index = 1
    while True:
        candidate = f"{base_dir}_{index}"
        if not os.path.exists(candidate):
            os.makedirs(candidate, exist_ok=True)
            return candidate
        index += 1

def _ensure_writable_dir(path: str) -> bool:
    try:
        os.makedirs(path, exist_ok=True)
        with tempfile.NamedTemporaryFile(dir=path, prefix=".write_test_", delete=True):
            pass
        return True
    except Exception:
        return False

def _file_uri_to_path(uri: str) -> str:
    parsed = urlparse(uri)
    if parsed.scheme != "file":
        raise ValueError(f"Unsupported root URI scheme: {parsed.scheme}")
    path = urllib.request.url2pathname(parsed.path)
    if parsed.netloc and parsed.netloc != "localhost":
        path = f"//{parsed.netloc}{path}"
    return os.path.abspath(path)

def _append_candidate(candidates: list[tuple[str, str]], path: str, source: str) -> None:
    candidate = os.path.abspath(os.path.expanduser((path or "").strip()))
    if not candidate:
        return
    if any(existing == candidate for existing, _ in candidates):
        return
    candidates.append((candidate, source))

async def _resolve_output_root(ctx: Context | None, env_var: str = "MCP_OUTPUT_DIR") -> str:
    candidates: list[tuple[str, str]] = []

    if ctx is not None:
        try:
            roots_result = await ctx.session.list_roots()
            for root in roots_result.roots:
                root_path = _file_uri_to_path(str(root.uri))
                candidate = os.path.dirname(root_path) if os.path.isfile(root_path) else root_path
                _append_candidate(candidates, candidate, "roots/list")
        except Exception:
            pass

    env_vars = [env_var]
    if env_var != "MCP_OUTPUT_DIR":
        env_vars.append("MCP_OUTPUT_DIR")

    for name in env_vars:
        configured_dir = os.getenv(name, "").strip()
        if configured_dir:
            _append_candidate(candidates, configured_dir, f"${name}")

    # Many stdio MCP clients do not implement roots/list, but they do spawn the server
    # with cwd set to the active workspace. Use that as a practical fallback.
    pwd = os.getenv("PWD", "").strip()
    if pwd:
        _append_candidate(candidates, pwd, "$PWD")

    try:
        _append_candidate(candidates, os.getcwd(), "cwd")
    except Exception:
        pass

    for candidate, _source in candidates:
        if _ensure_writable_dir(candidate):
            return candidate

    tried = ", ".join(f"{source}={candidate}" for candidate, source in candidates) if candidates else "(none)"
    raise OSError(
        "Unable to determine a writable client workspace directory. "
        f"Tried {tried}. Ensure the MCP client exposes roots/list, "
        f"or set {env_var}."
    )

@mcp.tool()
def get_pdf_metadata(file_path: str) -> dict:
    """
    Step 1: Get basic PDF information, including total page count.
    Must be called before starting visual inspection to determine the number of iterations.
    """
    if not os.path.exists(file_path):
        return {"error": f"File not found at {file_path}"}
    
    try:
        info = pdfinfo_from_path(file_path)
        return {
            "file_name": os.path.basename(file_path),
            "total_pages": info["Pages"],
            "page_size": info.get("Page size")
        }
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
async def download_arxiv_source(arxiv_url: str, ctx: Context | None = None) -> dict:
    """
    Given an arXiv URL or ID, download the corresponding TeX source into the client workspace and extract it.
    """
    arxiv_id = _extract_arxiv_id(arxiv_url)
    if not arxiv_id:
        return {"error": "Invalid arXiv URL or ID"}

    safe_id = arxiv_id.replace("/", "_")
    download_url = f"https://arxiv.org/e-print/{arxiv_id}"

    try:
        output_root = await _resolve_output_root(ctx)
        base_dir = os.path.join(output_root, f"arxiv_{safe_id}")
        target_dir = _ensure_unique_dir(base_dir)
    except Exception as e:
        return {"error": f"Unable to prepare output directory: {str(e)}"}

    try:
        request = urllib.request.Request(
            download_url,
            headers={"User-Agent": "Mozilla/5.0"}
        )
        with urllib.request.urlopen(request) as response:
            header = response.headers.get("Content-Disposition", "")
            filename = _filename_from_content_disposition(header)
            if not filename:
                filename = f"{safe_id}.source"
            filename = os.path.basename(filename)
            if not filename:
                filename = f"{safe_id}.source"
            download_path = os.path.join(target_dir, filename)
            with open(download_path, "wb") as out_file:
                shutil.copyfileobj(response, out_file)
    except urllib.error.HTTPError as e:
        return {"error": f"Download failed ({e.code}): {e.reason}"}
    except Exception as e:
        return {"error": f"Download failed: {str(e)}"}

    try:
        if tarfile.is_tarfile(download_path):
            with tarfile.open(download_path, "r:*") as tar:
                _safe_extract_tar(tar, target_dir)
            extracted_type = "tar"
            extracted_path = target_dir
        elif zipfile.is_zipfile(download_path):
            with zipfile.ZipFile(download_path) as zip_file:
                _safe_extract_zip(zip_file, target_dir)
            extracted_type = "zip"
            extracted_path = target_dir
        elif _is_gzip_file(download_path):
            output_name = os.path.splitext(os.path.basename(download_path))[0]
            if not output_name or output_name == os.path.basename(download_path):
                output_name = f"{safe_id}.tex"
            output_path = os.path.join(target_dir, output_name)
            with gzip.open(download_path, "rb") as gz_file, open(output_path, "wb") as out_file:
                shutil.copyfileobj(gz_file, out_file)
            extracted_type = "gzip"
            extracted_path = output_path
        else:
            extracted_type = "unknown"
            extracted_path = download_path
    except Exception as e:
        return {
            "error": f"Downloaded but failed to extract: {str(e)}",
            "download_path": download_path,
            "target_dir": target_dir
        }

    return {
        "arxiv_id": arxiv_id,
        "download_url": download_url,
        "download_path": download_path,
        "extracted_type": extracted_type,
        "extracted_path": extracted_path,
        "target_dir": target_dir
    }

@mcp.tool()
async def inspect_pdf_visually(file_path: str, pages: List[int] = None, ctx: Context | None = None):
    """
    Step 2: Visually read the PDF file.
    Used to inspect layout, formatting, image placement, and other visual elements.

    Args:
        file_path: Path to the PDF file.
        pages: List of page numbers to inspect (1-indexed). Recommend requesting only 3-5 pages at a time.
        save_debug: Whether to export debug images. Defaults to False; can be enabled via the environment variable MCP_SAVE_DEBUG_IMAGES=1.
    """
    if not os.path.exists(file_path):
        return [f"Error: File not found at {file_path}"]

    try:
        # dpi=100 balances clarity and token consumption
        images = convert_from_path(file_path, dpi=100)
        total_pages = len(images)
        results = []
        
        # Determine target page indices
        if pages:
            target_indices = [p-1 for p in pages if 1 <= p <= total_pages]
        else:
            target_indices = range(min(4, total_pages))  # Default to first 4 pages

        # Check if debug export is enabled via environment variable
        save_debug = os.getenv("MCP_SAVE_DEBUG_IMAGES", "").lower() in ("1", "true", "yes", "on")
        if save_debug:
            debug_root = await _resolve_output_root(ctx, "MCP_DEBUG_IMAGE_DIR")
            timestamp = datetime.utcnow().strftime("%Y%m%dT%H%M%S%f")
            base_name = os.path.splitext(os.path.basename(file_path))[0]
            debug_dir = os.path.join(debug_root, "debug_images", f"{base_name}_{timestamp}")
            os.makedirs(debug_dir, exist_ok=True)

        for i in target_indices:
            img = images[i]
            buffered = io.BytesIO()
            img.save(buffered, format="JPEG", quality=85)
            img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

            # Save debug images for reviewing what the MCP server actually sees
            if save_debug:
                debug_name = f"p{i+1}.jpg"
                debug_path = os.path.join(debug_dir, debug_name)
                img.save(debug_path, format="JPEG", quality=85)
            
            results.append({
                "type": "image",
                "data": img_str,
                "mimeType": "image/jpeg",
                "annotations": f"Page {i+1}"
            })
            
        return results

    except Exception as e:
        return [f"Error processing PDF visually: {str(e)}"]

def main():
    mcp.run()

if __name__ == "__main__":
    main()
