from __future__ import annotations

import os
import re
import time
import uuid
import logging
import threading
from dataclasses import dataclass, asdict
from typing import Dict, Optional, List
from urllib.parse import urlparse

import requests
import nbformat
from nbclient import NotebookClient

from fastapi import FastAPI, HTTPException, Header
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn

# =========================
# Hybrid SaaS Agent ACE
# - Secure local execution bridge for Cloud Control Plane
# - Auth (Bearer token)
# - Origin allowlist via env
# - URL + domain allowlist
# - Notebook execution in background thread
# - Status + logs endpoints
# =========================

logger = logging.getLogger("cifer.agent_ace")
logging.basicConfig(
    level=os.getenv("CIFER_AGENT_LOGLEVEL", "INFO").upper(),
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
)

app = FastAPI(title="Cifer Agent ACE", version="1.1.0")

# =========================
# Config (ENV-driven)
# =========================
CIFER_KERNEL_NAME = os.getenv("CIFER_KERNEL_NAME", "cifer-kernel")
CIFER_DISPLAY_NAME = os.getenv("CIFER_DISPLAY_NAME", "Cifer Kernel")

# Agent auth token for SaaS bridge (set on client machines)
# Example: export CIFER_AGENT_TOKEN="..."
CIFER_AGENT_TOKEN = os.getenv("CIFER_AGENT_TOKEN", "")

# CORS allowlist (comma-separated)
# Example: export CIFER_ALLOW_ORIGINS="http://localhost:5173,https://workspace.cifer.ai"
_allow_origins_raw = os.getenv("CIFER_ALLOW_ORIGINS", "http://localhost:5173")
ALLOW_ORIGINS = [o.strip() for o in _allow_origins_raw.split(",") if o.strip()]

# Notebook download allowlist domains (comma-separated)
# Example: export CIFER_NOTEBOOK_ALLOWLIST_DOMAINS="workspace.cifer.ai,assets.cifer.ai"
_allow_domains_raw = os.getenv("CIFER_NOTEBOOK_ALLOWLIST_DOMAINS", "")
ALLOW_NOTEBOOK_DOMAINS = [d.strip().lower() for d in _allow_domains_raw.split(",") if d.strip()]

# Downloads + execution
DOWNLOAD_TIMEOUT_SEC = int(os.getenv("CIFER_DOWNLOAD_TIMEOUT_SEC", "15"))
EXECUTION_TIMEOUT_SEC = int(os.getenv("CIFER_EXECUTION_TIMEOUT_SEC", "900"))  # nbclient cell timeout
MAX_NOTEBOOK_BYTES = int(os.getenv("CIFER_MAX_NOTEBOOK_BYTES", str(20 * 1024 * 1024)))  # 20MB

# Jupyter open URL (optional)
# If you want to open in an existing Jupyter server UI.
# Example: http://localhost:8888
JUPYTER_BASE_URL = os.getenv("CIFER_JUPYTER_BASE_URL", "")

# Storage directory for downloaded notebooks and outputs
AGENT_WORKDIR = os.getenv("CIFER_AGENT_WORKDIR", os.path.expanduser("~/.cifer/agent_ace"))

# =========================
# CORS
# =========================
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOW_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"] ,
    allow_headers=["*"] ,
)


# =========================
# Models
# =========================
class NotebookRequest(BaseModel):
    url: str = Field(..., description="Notebook URL (.ipynb)")
    project_id: str = Field(..., description="SaaS project id")
    execution_id: Optional[str] = Field(None, description="Optional client-provided execution id")
    open_in_browser: bool = Field(False, description="If true, open notebook in Jupyter UI (if configured)")


@dataclass
class ExecState:
    execution_id: str
    project_id: str
    url: str
    filename: str
    save_path: str
    status: str  # pending|running|success|failed
    started_at: Optional[float] = None
    finished_at: Optional[float] = None
    error: Optional[str] = None
    logs: Optional[List[str]] = None


# In-memory state (good enough for local agent). If you want persistence, write to sqlite/json.
_EXEC: Dict[str, ExecState] = {}
_EXEC_LOCK = threading.Lock()


# =========================
# Helpers
# =========================
_IPYNB_RE = re.compile(r"^[^\\/\\\\]+\\.ipynb$", re.IGNORECASE)


def _require_auth(authorization: Optional[str]) -> None:
    # If token not set, allow local/dev usage.
    if not CIFER_AGENT_TOKEN:
        return
    if not authorization or authorization.strip() != f"Bearer {CIFER_AGENT_TOKEN}":
        raise HTTPException(status_code=401, detail="Unauthorized")


def _validate_url(url: str) -> None:
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        raise HTTPException(status_code=400, detail="Invalid URL scheme")

    # Require .ipynb file
    filename = os.path.basename(parsed.path)
    if not _IPYNB_RE.match(filename):
        raise HTTPException(status_code=400, detail="URL must point to a .ipynb file")

    # Optional domain allowlist
    if ALLOW_NOTEBOOK_DOMAINS:
        host = (parsed.hostname or "").lower()
        if host not in ALLOW_NOTEBOOK_DOMAINS:
            raise HTTPException(status_code=403, detail="Notebook domain not allowed")


def _safe_filename_from_url(url: str) -> str:
    parsed = urlparse(url)
    filename = os.path.basename(parsed.path)
    # strip weird characters just in case
    filename = re.sub(r"[^a-zA-Z0-9._-]", "_", filename)
    if not _IPYNB_RE.match(filename):
        raise HTTPException(status_code=400, detail="Invalid notebook filename")
    return filename


def _download_notebook(url: str, dest_path: str) -> None:
    os.makedirs(os.path.dirname(dest_path), exist_ok=True)

    with requests.get(url, stream=True, timeout=DOWNLOAD_TIMEOUT_SEC) as r:
        if r.status_code != 200:
            raise HTTPException(status_code=502, detail=f"Failed to download notebook (HTTP {r.status_code})")

        total = 0
        with open(dest_path, "wb") as f:
            for chunk in r.iter_content(chunk_size=1024 * 64):
                if not chunk:
                    continue
                total += len(chunk)
                if total > MAX_NOTEBOOK_BYTES:
                    raise HTTPException(status_code=413, detail="Notebook too large")
                f.write(chunk)


def _open_in_jupyter(filename: str) -> None:
    if not JUPYTER_BASE_URL:
        return
    # Best-effort open
    try:
        import webbrowser

        url = f"{JUPYTER_BASE_URL.rstrip('/')}/notebooks/{filename}"
        webbrowser.open(url)
    except Exception:
        logger.exception("Failed to open notebook in browser")


def _execute_notebook(execution_id: str) -> None:
    # Background execution
    with _EXEC_LOCK:
        state = _EXEC.get(execution_id)
        if not state:
            return
        state.status = "running"
        state.started_at = time.time()
        state.logs = []

    logger.info("exec.start execution_id=%s project_id=%s file=%s", execution_id, state.project_id, state.filename)

    try:
        nb = nbformat.read(state.save_path, as_version=4)

        # NotebookClient timeout is per-cell; we also implement a hard guard by time tracking.
        client = NotebookClient(
            nb,
            kernel_name=CIFER_KERNEL_NAME,
            timeout=EXECUTION_TIMEOUT_SEC,
            allow_errors=False,
        )

        client.execute()

        with _EXEC_LOCK:
            st = _EXEC[execution_id]
            st.status = "success"
            st.finished_at = time.time()
            st.error = None

        logger.info("exec.success execution_id=%s duration=%.2fs", execution_id, (_EXEC[execution_id].finished_at - _EXEC[execution_id].started_at))

    except Exception as e:
        with _EXEC_LOCK:
            st = _EXEC.get(execution_id)
            if st:
                st.status = "failed"
                st.finished_at = time.time()
                st.error = str(e)

        logger.exception("exec.failed execution_id=%s", execution_id)


# =========================
# API
# =========================
@app.get("/health")
def health():
    return {
        "status": "ok",
        "kernel": CIFER_KERNEL_NAME,
        "workdir": AGENT_WORKDIR,
        "allow_origins": ALLOW_ORIGINS,
        "allow_domains": ALLOW_NOTEBOOK_DOMAINS,
    }


@app.post("/run_notebook")
async def run_notebook(
    data: NotebookRequest,
    authorization: Optional[str] = Header(None),
):
    _require_auth(authorization)
    _validate_url(data.url)

    execution_id = data.execution_id or str(uuid.uuid4())
    filename = _safe_filename_from_url(data.url)

    # Per-project workspace
    project_dir = os.path.join(AGENT_WORKDIR, "projects", data.project_id)
    save_path = os.path.join(project_dir, "notebooks", execution_id, filename)

    # Download first (sync) so we fail fast
    _download_notebook(data.url, save_path)

    with _EXEC_LOCK:
        _EXEC[execution_id] = ExecState(
            execution_id=execution_id,
            project_id=data.project_id,
            url=data.url,
            filename=filename,
            save_path=save_path,
            status="pending",
        )

    if data.open_in_browser:
        _open_in_jupyter(filename)

    # Start background execution (thread)
    t = threading.Thread(target=_execute_notebook, args=(execution_id,), daemon=True)
    t.start()

    return {
        "status": "accepted",
        "execution_id": execution_id,
        "project_id": data.project_id,
        "file": filename,
    }


@app.get("/status/{execution_id}")
def status(execution_id: str, authorization: Optional[str] = Header(None)):
    _require_auth(authorization)

    with _EXEC_LOCK:
        st = _EXEC.get(execution_id)
        if not st:
            raise HTTPException(status_code=404, detail="execution_id not found")
        return asdict(st)


@app.get("/executions")
def list_executions(authorization: Optional[str] = Header(None)):
    _require_auth(authorization)

    with _EXEC_LOCK:
        # return a compact list
        items = []
        for v in _EXEC.values():
            items.append(
                {
                    "execution_id": v.execution_id,
                    "project_id": v.project_id,
                    "status": v.status,
                    "filename": v.filename,
                    "started_at": v.started_at,
                    "finished_at": v.finished_at,
                    "error": v.error,
                }
            )
        return {"items": items}


# ▶️ Function to start the FastAPI server
def run_agent_ace(port: int = 9999):
    logger.info("Starting agent-ace FastAPI server at http://localhost:%s", port)
    logger.info("Using kernel: %s (%s)", CIFER_DISPLAY_NAME, CIFER_KERNEL_NAME)

    dev_mode = os.getenv("CIFER_DEV_MODE", "false").lower() == "true"

    uvicorn.run(
        "cifer.agent_ace:app",
        host="0.0.0.0",
        port=port,
        reload=dev_mode,
        log_level=os.getenv("CIFER_AGENT_LOGLEVEL", "info").lower(),
    )


if __name__ == "__main__":
    run_agent_ace()
