import os, json, time, uuid, subprocess, shutil, base64
from pathlib import Path
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from fastapi import Request
from cryptography.fernet import Fernet
from huggingface_hub import HfApi

from router import PredictiveContextFilter, InferenceRouterCircuitBreaker, ParallelEngine

# ---------- globals ----------
STATE_DIR = Path("/tmp/altamira-state")
STATE_DIR.mkdir(parents=True, exist_ok=True)
WORKSPACE_DIR = Path("/tmp/altamira-workspace")
WORKSPACE_DIR.mkdir(parents=True, exist_ok=True)
CACHE_DIR = Path("/tmp/altamira-cache")
CACHE_DIR.mkdir(parents=True, exist_ok=True)
SANDBOX_DIR = Path("/tmp/altamira-sandboxes")
SANDBOX_DIR.mkdir(parents=True, exist_ok=True)
PROJECTS_FILE = STATE_DIR / "projects.json"
ENCRYPTION_KEY_FILE = STATE_DIR / ".encryption_key"
STAGE_DIRS = [CACHE_DIR / f"stage_{i}" for i in range(3)]

templates = Jinja2Templates(directory="templates")
hf_api = HfApi()
filter_ctx = PredictiveContextFilter()
circuit_breaker = InferenceRouterCircuitBreaker()
engine = ParallelEngine()

# ---------- encryption ----------
if ENCRYPTION_KEY_FILE.exists():
    _key = ENCRYPTION_KEY_FILE.read_bytes()
else:
    _key = Fernet.generate_key()
    ENCRYPTION_KEY_FILE.write_bytes(_key)
cipher = Fernet(_key)

def encrypt(data: dict) -> str:
    return cipher.encrypt(json.dumps(data).encode()).decode()

def decrypt(token: str) -> dict:
    return json.loads(cipher.decrypt(token.encode()).decode())

# ---------- project store ----------
def _load_projects() -> dict:
    if PROJECTS_FILE.exists():
        return json.loads(PROJECTS_FILE.read_text())
    return {}

def _save_projects(projects: dict):
    PROJECTS_FILE.write_text(json.dumps(projects, indent=2))

# ---------- lifecycle ----------
@asynccontextmanager
async def lifespan(_app: FastAPI):
    for d in [STATE_DIR, WORKSPACE_DIR, CACHE_DIR, SANDBOX_DIR]:
        d.mkdir(parents=True, exist_ok=True)
    for d in STAGE_DIRS:
        d.mkdir(parents=True, exist_ok=True)
    yield

app = FastAPI(title="Altamira Orchestrator", lifespan=lifespan)

# ==================== UI ====================
@app.get("/", response_class=HTMLResponse)
async def index():
    return templates.TemplateResponse("index.html", {"request": {}})

@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "app": "altamira-orchestrator",
        "version": "1.2.0",
        "circuit_breaker": circuit_breaker.state,
        "failures": circuit_breaker.failure_count,
    }

# ==================== GATEKEEPER ====================
@app.post("/api/gatekeeper/validate")
async def gatekeeper_validate(body: dict):
    hf_token = body.get("hf_token", "")
    gh_token = body.get("gh_token", "")
    if not hf_token:
        raise HTTPException(400, "HF token required")
    try:
        import huggingface_hub as hh
        hh.login(token=hf_token, add_to_git_credential=False)
        who = hf_api.whoami()
        username = who.get("name", who.get("login", "unknown"))
    except Exception as e:
        raise HTTPException(401, f"HF token invalid: {e}")
    gh_user = None
    if gh_token:
        import httpx
        r = httpx.get("https://api.github.com/user", headers={"Authorization": f"Bearer {gh_token}"})
        if r.status_code != 200:
            raise HTTPException(401, "GitHub PAT invalid")
        gh_user = r.json().get("login")
    sandbox = f"sandbox-{username}-{int(time.time())}"
    (SANDBOX_DIR / sandbox).mkdir(parents=True, exist_ok=True)
    return {"status": "allocated", "sandbox": sandbox, "user": username, "gh_user": gh_user}

@app.post("/api/gatekeeper/check")
async def gatekeeper_check():
    hf_token = os.environ.get("HF_TOKEN")
    if not hf_token:
        return {"hf_authenticated": False}
    try:
        who = hf_api.whoami()
        return {"hf_authenticated": True, "user": who.get("name")}
    except:
        return {"hf_authenticated": False}

# ==================== PROJECTS ====================
@app.get("/api/projects")
async def list_projects():
    return _load_projects()

@app.post("/api/projects")
async def create_project(body: dict):
    name = body.get("name", "").strip()
    repo_url = body.get("repo_url", "").strip()
    if not name:
        raise HTTPException(400, "Project name required")
    projects = _load_projects()
    if name in projects:
        raise HTTPException(409, "Project already exists")
    projects[name] = {
        "name": name,
        "repo_url": repo_url,
        "created": time.time(),
        "active": False,
    }
    _save_projects(projects)
    return projects[name]

@app.post("/api/projects/{name}/activate")
async def activate_project(name: str):
    projects = _load_projects()
    if name not in projects:
        raise HTTPException(404, "Project not found")
    for p in projects.values():
        p["active"] = False
    projects[name]["active"] = True
    _save_projects(projects)
    # workspace switch
    target = WORKSPACE_DIR / name
    target.mkdir(parents=True, exist_ok=True)
    repo_url = projects[name].get("repo_url", "")
    if repo_url:
        subprocess.run(["git", "clone", repo_url, str(target)],
                       capture_output=True, text=True, timeout=60)
    return {"active": name, "workspace": str(target)}

@app.delete("/api/projects/{name}")
async def delete_project(name: str):
    projects = _load_projects()
    if name not in projects:
        raise HTTPException(404)
    del projects[name]
    _save_projects(projects)
    target = WORKSPACE_DIR / name
    if target.exists():
        shutil.rmtree(target)
    return {"deleted": name}

# ==================== AGENT CONSOLE ====================
@app.post("/api/console/exec")
async def console_exec(body: dict):
    command = body.get("command", "").strip()
    project = body.get("project", "default")
    if not command:
        raise HTTPException(400, "Command required")
    cwd = WORKSPACE_DIR / project
    cwd.mkdir(parents=True, exist_ok=True)
    try:
        result = subprocess.run(command, shell=True, capture_output=True,
                                text=True, timeout=int(body.get("timeout", 30)),
                                cwd=str(cwd))
        return {
            "stdout": result.stdout,
            "stderr": result.stderr,
            "returncode": result.returncode,
        }
    except subprocess.TimeoutExpired:
        raise HTTPException(408, "Command timed out")
    except Exception as e:
        raise HTTPException(500, str(e))

@app.websocket("/api/console/stream")
async def console_stream(websocket: WebSocket):
    await websocket.accept()
    try:
        data = await websocket.receive_json()
        command = data.get("command", "")
        project = data.get("project", "default")
        cwd = WORKSPACE_DIR / project
        cwd.mkdir(parents=True, exist_ok=True)
        process = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(cwd),
            shell=True,
        )
        async def stream_output(stream, label):
            while True:
                line = await stream.readline()
                if not line:
                    break
                await websocket.send_json({"type": label, "data": line.decode().rstrip()})
        import asyncio
        await asyncio.gather(
            stream_output(process.stdout, "stdout"),
            stream_output(process.stderr, "stderr"),
        )
        await process.wait()
        await websocket.send_json({"type": "exit", "code": process.returncode})
    except WebSocketDisconnect:
        pass
    except Exception as e:
        await websocket.send_json({"type": "error", "data": str(e)})
    finally:
        try:
            await websocket.close()
        except:
            pass

# ==================== STATE ====================
@app.get("/api/state")
async def get_state():
    state_file = STATE_DIR / "state.json"
    if state_file.exists():
        data = json.loads(state_file.read_text())
    else:
        data = {}
    return {
        "state": encrypt(data),
        "last_sync": data.get("last_sync"),
        "buffer": "double_buffered",
        "stages": [str(d) for d in STAGE_DIRS if d.exists()],
    }

@app.post("/api/state/sync")
async def sync_state():
    data = {"last_sync": time.time(), "buffer": "double_buffered"}
    (STATE_DIR / "state.json").write_text(json.dumps(data))
    # rotate stages
    for i in range(len(STAGE_DIRS) - 1, 0, -1):
        src, dst = STAGE_DIRS[i - 1], STAGE_DIRS[i]
        if src.exists():
            if dst.exists():
                shutil.rmtree(dst)
            shutil.copytree(src, dst, dirs_exist_ok=True)
    return {"status": "synced", "path": str(STATE_DIR / "state.json")}

@app.post("/api/state/encrypt")
async def encrypt_state(body: dict):
    return {"encrypted": encrypt(body.get("data", {}))}

@app.post("/api/state/decrypt")
async def decrypt_state(body: dict):
    try:
        return {"data": decrypt(body.get("token", ""))}
    except Exception as e:
        raise HTTPException(400, f"Decryption failed: {e}")

# ==================== ROUTER ====================
@app.get("/api/router")
async def router_status():
    return {
        "circuit_breaker": circuit_breaker.state,
        "failures": circuit_breaker.failure_count,
        "recovery_timeout": circuit_breaker.recovery_timeout,
        "filter_capacity": filter_ctx.capacity,
        "filter_threshold": filter_ctx.threshold,
        "parallel_max": engine.semaphore._value,
    }

@app.post("/api/router/filter")
async def router_filter(body: dict):
    text = body.get("text", "")
    compacted = await filter_ctx.monitor(text)
    return {"original_length": len(text), "compacted_length": len(compacted), "compacted": compacted}

@app.post("/api/router/circuit/reset")
async def reset_circuit():
    circuit_breaker.failure_count = 0
    circuit_breaker.state = "closed"
    return {"status": "reset"}

# ==================== SANDBOX / PREVIEW ====================
@app.get("/api/sandbox/{name}")
async def get_sandbox(name: str):
    path = SANDBOX_DIR / name
    if not path.exists():
        raise HTTPException(404, "Sandbox not found")
    files = []
    for f in path.rglob("*"):
        if f.is_file():
            files.append({"name": str(f.relative_to(path)),
                          "size": f.stat().st_size,
                          "modified": f.stat().st_mtime})
    return {"name": name, "files": sorted(files, key=lambda x: x["name"])}

@app.get("/api/sandbox/{name}/read")
async def read_sandbox_file(name: str, file: str):
    path = SANDBOX_DIR / name / file
    if not path.exists() or not path.is_file():
        raise HTTPException(404)
    return {"content": path.read_text()}

@app.get("/api/sandboxes")
async def list_sandboxes():
    return [d.name for d in SANDBOX_DIR.iterdir() if d.is_dir()]

# ==================== WORKER INTEGRATION ====================
# mount worker app as sub-app
from worker import app as worker_app
app.mount("/worker", worker_app)

# ==================== WEBHOOK RESULTS ====================
@app.get("/api/results")
async def list_results():
    results_file = Path("/tmp/altamira-results.json")
    if not results_file.exists():
        return {"results": []}
    lines = results_file.read_text().strip().split("\n")
    return {"results": [json.loads(l) for l in lines if l]}

# ==================== SYSTEM ====================
@app.get("/api/system")
async def system_info():
    import platform as _platform
    return {
        "platform": _platform.platform(),
        "python": _platform.python_version(),
        "hostname": os.uname().nodename,
        "cpus": os.cpu_count(),
        "sandbox_count": len(list(SANDBOX_DIR.iterdir())),
        "workspace_count": len(list(WORKSPACE_DIR.iterdir())),
        "disk_tmp": _disk_usage("/tmp"),
    }

def _disk_usage(path: str) -> dict:
    s = os.statvfs(path)
    return {
        "total_gb": round(s.f_frsize * s.f_blocks / 1e9, 2),
        "free_gb": round(s.f_frsize * s.f_bfree / 1e9, 2),
    }