import os, json, asyncio, subprocess, uuid
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, BackgroundTasks, UploadFile
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
from pathlib import Path

from model_router import route_chat, get_available_models_sync, refresh_model_cache
from agent import CodingAgent
from autocomplete import generate_autocomplete

# Lazy codebase index
CodebaseIndex = None
def get_codebase_index():
    global CodebaseIndex
    if CodebaseIndex is None:
        try:
            from codebase_index import CodebaseIndex as CBI
            CodebaseIndex = CBI
        except ImportError:
            raise HTTPException(status_code=503, detail="Codebase indexing unavailable")
    return CodebaseIndex

async def startup_index():
    try:
        from codebase_index import CodebaseIndex
        idx = CodebaseIndex("/app/sandbox")
        idx.build_index()
        indexers["/app/sandbox"] = idx
    except Exception as e:
        print(f"Startup indexing failed: {e}")

@asynccontextmanager
async def lifespan(app: FastAPI):
    await refresh_model_cache()
    await startup_index()
    yield

app = FastAPI(title="Codeki API", lifespan=lifespan)

indexers = {}
agents = {}
STATIC_DIR = "static"

class ChatRequest(BaseModel):
    model: str = "groq-llama-3.3-70b-versatile"
    messages: List[Dict[str, str]]
    context_files: Optional[List[str]] = None
    project_path: str = "/app/sandbox"

class AgentTaskRequest(BaseModel):
    prompt: str
    model: str = "groq-llama-3.3-70b-versatile"
    project_path: str = "/app/sandbox"
    background: bool = False

class AutocompleteRequest(BaseModel):
    prefix: str
    suffix: str = ""
    filepath: str
    language: str
    max_tokens: int = 64
    model: str = "groq-llama-3.1-8b-instant"

class IndexRequest(BaseModel):
    project_path: str

class BugFixRequest(BaseModel):
    file_path: str
    error_output: str
    project_path: str

class ReviewRequest(BaseModel):
    pr_diff: str
    project_path: str

class ApplyRequest(BaseModel):
    task_id: str
    approved_indices: List[int]

# ---------- API Routes ----------
@app.get("/api/models")
async def get_models():
    return get_available_models_sync()

@app.post("/api/refresh-models")
async def refresh_models():
    await refresh_model_cache()
    return {"status": "ok", "count": len(get_available_models_sync())}

@app.post("/api/chat")
async def chat(req: ChatRequest):
    context = ""
    if req.context_files:
        for f in req.context_files:
            if f == "__git__":
                import git
                try:
                    repo = git.Repo(req.project_path)
                    diff = repo.git.diff(repo.head.commit.tree) if repo.head.is_valid() else ""
                    context += f"\n--- Git Diff ---\n{diff}\n"
                except:
                    pass
            else:
                p = Path(req.project_path) / f
                if p.exists():
                    context += f"\n--- {f} ---\n{p.read_text()}\n"
        req.messages.insert(0, {"role": "system", "content": f"Relevant codebase context:\n{context}"})
    async def event_stream():
        try:
            async for token in route_chat(req.model, req.messages):
                yield f"data: {json.dumps({'token': token})}\n\n"
            yield "data: [DONE]\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'token': f'ERROR: {str(e)}'})}\n\n"
            yield "data: [DONE]\n\n"
    return StreamingResponse(event_stream(), media_type="text/event-stream")

# ---------- Agent endpoints with error handling ----------
@app.post("/api/agent/run")
async def run_agent(req: AgentTaskRequest, background_tasks: BackgroundTasks):
    try:
        agent = CodingAgent(project_path=req.project_path, model=req.model)
        task_id = str(uuid.uuid4())
        agents[task_id] = agent
        if req.background:
            background_tasks.add_task(agent.run_task, req.prompt)
            return {"task_id": task_id, "status": "started"}
        else:
            result = await agent.run_task(req.prompt)
            return {
                "task_id": task_id,
                "status": agent.status,
                "logs": agent.logs,
                "pending_changes": result["pending_changes"]
            }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/agent/approve/{task_id}")
async def approve_agent_changes(task_id: str):
    agent = agents.get(task_id)
    if not agent:
        raise HTTPException(404, "Agent not found")
    if agent.status != "awaiting_approval":
        raise HTTPException(400, "Agent not awaiting approval")
    try:
        await agent.approve_and_apply()
        return {"status": agent.status, "logs": agent.logs}
    except Exception as e:
        raise HTTPException(500, detail=str(e))

@app.get("/api/agent/task/{task_id}")
async def get_agent_status(task_id: str):
    agent = agents.get(task_id)
    if not agent:
        raise HTTPException(404, "Task not found")
    return {
        "status": agent.status,
        "logs": agent.logs,
        "pending_changes": agent.pending_changes
    }

@app.get("/api/agent/tasks")
async def list_agent_tasks():
    return {"tasks": list(agents.keys())}

# ---------- Other endpoints ----------
@app.post("/api/autocomplete")
async def autocomplete(req: AutocompleteRequest):
    suggestions = await generate_autocomplete(req)
    return {"completions": suggestions}

@app.post("/api/index")
async def index_project(req: IndexRequest):
    CBI = get_codebase_index()
    idx = CBI(req.project_path)
    idx.build_index()
    indexers[req.project_path] = idx
    return {"status": "indexed", "files": len(idx.file_list)}

@app.post("/api/search")
async def search_codebase(query: str, project_path: str = "/app/sandbox", top_k: int = 5):
    if project_path not in indexers:
        CBI = get_codebase_index()
        idx = CBI(project_path)
        try:
            idx.build_index()
        except Exception as e:
            raise HTTPException(503, f"Indexing failed: {str(e)}")
        indexers[project_path] = idx
    idx = indexers[project_path]
    results = idx.search(query, top_k)
    return {"results": results}

@app.post("/api/bugfix")
async def bugfix(req: BugFixRequest):
    agent = CodingAgent(project_path=req.project_path, model="groq-llama-3.3-70b-versatile")
    result = await agent.fix_error(req.file_path, req.error_output)
    return result

@app.post("/api/review")
async def review_code(req: ReviewRequest):
    agent = CodingAgent(project_path=req.project_path, model="groq-llama-3.3-70b-versatile")
    review = await agent.review_pr(req.pr_diff)
    return {"review": review}

@app.post("/api/git")
async def git_operation(action: str, project_path: str, args: Optional[dict] = None):
    import git
    try:
        repo = git.Repo(project_path)
        if action == "commit":
            repo.git.add(A=True)
            repo.git.commit(m=args.get("message", "auto commit"))
        elif action == "log":
            logs = list(repo.iter_commits(max_count=10))
            return {"log": [{"message": c.message, "hash": c.hexsha} for c in logs]}
        elif action == "branch":
            branches = [b.name for b in repo.branches]
            return {"branches": branches}
        elif action == "diff":
            diff_output = repo.git.diff(repo.head.commit.tree) if repo.head.is_valid() else ""
            return {"diff": diff_output}
        else:
            raise HTTPException(400, "Unknown action")
    except Exception as e:
        raise HTTPException(400, str(e))

@app.websocket("/ws/terminal")
async def terminal_ws(websocket: WebSocket):
    await websocket.accept()
    proc = None
    try:
        while True:
            cmd = await websocket.receive_text()
            if cmd.startswith("exec:"):
                command = cmd[5:]
                proc = await asyncio.create_subprocess_shell(
                    command,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                    cwd="/app/sandbox"
                )
                stdout, stderr = await proc.communicate()
                output = stdout.decode() + stderr.decode()
                await websocket.send_text(output)
            elif cmd == "ctrl-c" and proc:
                proc.terminate()
    except WebSocketDisconnect:
        pass

@app.get("/api/files")
async def list_files(project_path: str = "/app/sandbox"):
    def build_tree(path):
        tree = []
        for item in sorted(os.listdir(path)):
            if item.startswith('.') or item == '__pycache__':
                continue
            full = os.path.join(path, item)
            node = {
                "name": item,
                "path": full.replace(project_path, "").lstrip("/"),
                "type": "folder" if os.path.isdir(full) else "file"
            }
            if os.path.isdir(full):
                node["children"] = build_tree(full)
            tree.append(node)
        return tree
    try:
        return build_tree(project_path)
    except Exception as e:
        raise HTTPException(400, str(e))

@app.get("/api/file")
async def get_file(path: str, project_path: str = "/app/sandbox"):
    full_path = Path(project_path) / path
    if not full_path.exists():
        raise HTTPException(404, "File not found")
    return {"content": full_path.read_text()}

@app.post("/api/upload")
async def upload_file(file: UploadFile, project_path: str = "/app/sandbox"):
    dest = Path(project_path) / file.filename
    dest.parent.mkdir(parents=True, exist_ok=True)
    content = await file.read()
    dest.write_bytes(content)
    return {"filename": file.filename, "size": len(content)}

# ---------- Serve frontend ----------
@app.get("/{full_path:path}")
async def serve_frontend(full_path: str):
    if full_path.startswith("api"):
        raise HTTPException(404)
    file_path = os.path.join(STATIC_DIR, full_path)
    if os.path.isfile(file_path):
        return FileResponse(file_path)
    index_path = os.path.join(STATIC_DIR, "index.html")
    if os.path.isfile(index_path):
        return FileResponse(index_path)
    raise HTTPException(404, detail="Not found")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)