api.py

import os
import gc
import time
import shutil
import asyncio
import subprocess
import sqlite3
import uvicorn
from pathlib import Path
from typing import List, Optional

from fastapi import (
    FastAPI, Query, HTTPException, WebSocket, WebSocketDisconnect, 
    BackgroundTasks, UploadFile, File
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from contextlib import asynccontextmanager

from core.analyzer import GraphAnalyzer
from core.embeddings import EmbeddingService
from core.parser import CodeParser
from core.db import Database
from core.ignore import GitIgnoreChecker
from core.api_linker import APILinker
from core.config import get_config, save_config
from core.llm import CodeAtlasAI, get_installed_ollama_models
from core.strategies.path_normalizer import normalize_path
from core.git_helper import get_git_authors

MAIN_LOOP: Optional[asyncio.AbstractEventLoop] = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global MAIN_LOOP
    MAIN_LOOP = asyncio.get_running_loop()
    print("🚀 [LOCAL ATLAS] Server started. Local desktop mode active.")
    yield

app = FastAPI(title="Code Atlas Local API", lifespan=lifespan)


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

DATA_DIR = Path("data").resolve()
PROJECTS_DIR = DATA_DIR / "projects"
PROJECTS_DIR.mkdir(parents=True, exist_ok=True)

# WebSocket Manager
class ConnectionManager:
    def __init__(self):
        self.connections: List[WebSocket] = []

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.connections.append(ws)

    def disconnect(self, ws: WebSocket):
        if ws in self.connections:
            self.connections.remove(ws)

    async def broadcast(self, message: dict):
        for conn in self.connections:
            try:
                await conn.send_json(message)
            except Exception:
                pass

ws_manager = ConnectionManager()

def dispatch_progress(status: str, message: str, percent: int):
    global MAIN_LOOP
    payload = {"status": status, "message": message, "percent": percent}
    if MAIN_LOOP and MAIN_LOOP.is_running():
        asyncio.run_coroutine_threadsafe(ws_manager.broadcast(payload), MAIN_LOOP)

@app.websocket("/ws/progress")
async def websocket_endpoint(websocket: WebSocket):
    await ws_manager.connect(websocket)
    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        ws_manager.disconnect(websocket)

# Models
class LocalParseRequest(BaseModel):
    project_name: str
    directory: str

class GithubParseRequest(BaseModel):
    project_name: Optional[str] = None
    github_url: str

class ChatRequest(BaseModel):
    node_id: str
    message: str
    selected_model: str

class LayoutUpdate(BaseModel):
    node_id: str
    fx: Optional[float] = None
    fy: Optional[float] = None


# --- CONFIG & OLLAMA MODEL DISCOVERY ---
@app.get("/api/config")
def read_config():
    return {"config": get_config()}

@app.post("/api/config")
def update_config(new_config: dict):
    save_config(new_config)
    return {"status": "success"}

@app.get("/api/models")
def list_available_models():
    cfg = get_config()
    ollama_host = cfg.get("ollama_host", "http://localhost:11434")
    ollama_models = get_installed_ollama_models(ollama_host)
    
    return {
        "ollama_active": len(ollama_models) > 0,
        "ollama_models": ollama_models,
        "cloud_models": [
            "gemini-3.5-flash-lite",
            "gemini-3.5-flash",
            "gemini-3.6-flash",
            "gemini-3.7-flash",
            "gemini-3.8-flash",
        ]
    }

# --- NATIVE OS FOLDER PICKER ---
@app.post("/api/utils/browse-folder")
def open_native_folder_dialog():
    """Opens native OS folder chooser (Explorer/Finder) using Tkinter."""
    try:
        import tkinter as tk
        from tkinter import filedialog
        root = tk.Tk()
        root.withdraw()
        root.attributes('-topmost', True)
        folder = filedialog.askdirectory(title="Select Codebase Folder for Code Atlas")
        root.destroy()
        return {"path": folder or ""}
    except Exception as e:
        return {"path": "", "error": str(e)}


# --- LOCAL PIPELINE EXECUTION ---
def execute_local_pipeline(source_path: Path, project_name: str):
    project_dir = PROJECTS_DIR / project_name
    project_dir.mkdir(parents=True, exist_ok=True)
    db_path = project_dir / "atlas.db"
    index_path = project_dir / "atlas.index"

    if db_path.exists():
        os.remove(db_path)

    # Save target directory link for live git blame and code inspection
    with open(project_dir / "target_path.txt", "w") as f:
        f.write(str(source_path.resolve()))

    dispatch_progress("Initializing...", "Starting database...", 5)
    db = Database(str(db_path))
    parser = CodeParser()
    ignore_checker = GitIgnoreChecker(str(source_path))

    valid_files = []
    valid_extensions = {'.py', '.js', '.jsx', '.ts', '.tsx'}

    for root, dirs, files in os.walk(source_path):
        root_path = Path(root)
        try:
            rel_root = str(root_path.relative_to(source_path))
        except ValueError:
            rel_root = ""

        for i in range(len(dirs) - 1, -1, -1):
            if ignore_checker.is_ignored(dirs[i], os.path.join(rel_root, dirs[i])):
                del dirs[i]

        for f in files:
            file_path = root_path / f
            if file_path.suffix in valid_extensions and not ignore_checker.is_ignored(f, os.path.join(rel_root, f)):
                valid_files.append(file_path)

    total = len(valid_files)
    if total == 0:
        dispatch_progress("Error", "No supported source files found in directory.", 0)
        db.close()
        return

    BATCH_SIZE = 50
    parsed_batch = []
    last_reported = 0

    try:
        for i, file in enumerate(valid_files, 1):
            try:
                parsed_module = parser.parse_file(str(file))
                rel = file.relative_to(source_path)
                parsed_module.filepath = str(rel).replace("\\", "/")
                parsed_batch.append(parsed_module)
            except Exception:
                pass

            pct = int(10 + (i / total) * 50)
            if i % 10 == 0 or i == total or (pct - last_reported >= 3):
                last_reported = pct
                dispatch_progress("Parsing AST...", f"Parsed {i} of {total} source files...", pct)

            if len(parsed_batch) >= BATCH_SIZE:
                db.save_modules_batch(parsed_batch)
                parsed_batch.clear()

        if parsed_batch:
            db.save_modules_batch(parsed_batch)
            parsed_batch.clear()

    finally:
        db.close()

    dispatch_progress("Linking APIs...", "Connecting cross-language endpoints...", 68)
    linker = APILinker(str(db_path))
    linker.run_linkage()

    dispatch_progress("Vectorizing...", "Generating semantic embeddings (FAISS)...", 82)
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    try:
        embedder = EmbeddingService(str(db_path), str(index_path))
        embedder.generate_embeddings()
    except Exception as e:
        print(f"Embedding error: {e}")

    dispatch_progress("Complete!", "Architecture map is ready.", 100)

@app.post("/api/projects/local")
def parse_local_project(req: LocalParseRequest, background_tasks: BackgroundTasks):
    src = Path(req.directory).resolve()
    if not src.exists() or not src.is_dir():
        raise HTTPException(status_code=400, detail="Invalid local directory path.")

    # Allow spaces so "Nutri Verify" doesn't mismatch the frontend route
    raw_name = req.project_name.strip() if req.project_name else src.name
    clean_name = "".join(c for c in raw_name if c.isalnum() or c in " -_").strip()
    if not clean_name:
        clean_name = "project"

    background_tasks.add_task(execute_local_pipeline, src, clean_name)
    return {"status": "started", "project_name": clean_name}


@app.post("/api/projects/github")
def parse_github_project(req: GithubParseRequest, background_tasks: BackgroundTasks):
    url = req.github_url.strip().rstrip('/')
    raw_name = req.project_name.strip() if req.project_name else url.split('/')[-1].replace('.git', '')
    clean_name = "".join(c for c in raw_name if c.isalnum() or c in " -_").strip() or "repo"
    
    def process_clone():
        clone_dir = Path("data/clones") / clean_name
        if clone_dir.exists():
            shutil.rmtree(clone_dir)
        clone_dir.parent.mkdir(parents=True, exist_ok=True)
        
        dispatch_progress("Cloning...", f"Cloning {clean_name} from GitHub...", 5)
        try:
            subprocess.run(["git", "clone", "--depth", "1", url, str(clone_dir)], check=True)
            execute_local_pipeline(clone_dir, clean_name)
        except Exception as e:
            dispatch_progress("Error", f"Git clone failed: {e}", 0)

    background_tasks.add_task(process_clone)
    return {"status": "started", "project_name": clean_name}

# --- PROJECTS & GRAPH ---
@app.get("/api/projects")
def list_projects():
    return [d.name for d in PROJECTS_DIR.iterdir() if d.is_dir() and (d / "atlas.db").exists()]

@app.delete("/api/projects/{project_name}")
def delete_project(project_name: str):
    p_dir = PROJECTS_DIR / project_name
    if p_dir.exists() and p_dir.is_dir():
        shutil.rmtree(p_dir, ignore_errors=True)
        return {"status": "success"}
    raise HTTPException(status_code=404, detail="Project not found")

@app.get("/api/graph/{project_name}")
def get_graph(project_name: str):
    db_path = PROJECTS_DIR / project_name / "atlas.db"
    if not db_path.exists():
        raise HTTPException(status_code=404, detail="Project DB not found")
        
    analyzer = GraphAnalyzer(str(db_path))
    try:
        analyzer.build_module_graph()
        return {
            "graph": analyzer.export_json(),
            "cycles": [c for c in analyzer.get_cyclic_dependencies() if len(c) > 1]
        }
    finally:
        analyzer.conn.close()

@app.post("/api/graph/{project_name}/layout")
def save_layout(project_name: str, updates: List[LayoutUpdate]):
    db_path = PROJECTS_DIR / project_name / "atlas.db"
    conn = sqlite3.connect(str(db_path))
    cursor = conn.cursor()
    for u in updates:
        if u.fx is None or u.fy is None:
            cursor.execute("DELETE FROM layout WHERE node_id = ?", (u.node_id,))
        else:
            cursor.execute("INSERT OR REPLACE INTO layout (node_id, fx, fy) VALUES (?, ?, ?)", (u.node_id, u.fx, u.fy))
    conn.commit()
    conn.close()
    return {"status": "success"}


# --- SEARCH & NODE DETAILS ---
@app.get("/api/search/{project_name}")
def search_codebase(project_name: str, q: str = Query(...)):
    t0 = time.perf_counter()
    db_path = PROJECTS_DIR / project_name / "atlas.db"
    index_path = PROJECTS_DIR / project_name / "atlas.index"
    embedder = EmbeddingService(str(db_path), str(index_path))
    results = embedder.search(q, top_k=5)
    print(f"🔍 [SEARCH] Query '{q}' returned {len(results)} matches in {(time.perf_counter() - t0)*1000:.1f}ms")
    return {"query": q, "results": results}

@app.get("/api/node/{project_name}/{node_id:path}")
def get_node_details(project_name: str, node_id: str):
    project_dir = PROJECTS_DIR / project_name
    db_path = project_dir / "atlas.db"
    node_name = node_id.split('.')[-1]

    # Resolve local original directory
    target_path = None
    target_path_file = project_dir / "target_path.txt"
    if target_path_file.exists():
        target_path = Path(target_path_file.read_text().strip())

    conn = sqlite3.connect(str(db_path))
    try:
        cursor = conn.cursor()
        cursor.execute("""
            SELECT n.name, n.node_type, n.parent_name, n.start_line, n.end_line, n.code_snippet, f.filepath
            FROM nodes n JOIN files f ON n.file_id = f.id WHERE n.name = ?
        """, (node_name,))
        
        rows = cursor.fetchall()
        for row in rows:
            name, node_type, parent_name, start_line, end_line, snippet, filepath = row
            mod_name = normalize_path(filepath)
            expected_id = f"{mod_name}.{parent_name}.{name}" if parent_name else f"{mod_name}.{name}"
            
            if expected_id == node_id or len(rows) == 1:
                scoped_caller = f"{parent_name}.{name}" if parent_name else name

                cursor.execute("""
                    SELECT DISTINCT callee FROM calls 
                    WHERE caller = ? OR caller = ? OR caller LIKE ?
                """, (name, scoped_caller, f"%.{name}"))
                callees = [r[0] for r in cursor.fetchall() if r[0] and r[0] != "API"]

                callee_targets = [name]
                if parent_name:
                    callee_targets.extend([parent_name, scoped_caller])
                placeholders = ",".join(["?"] * len(callee_targets))
                cursor.execute(f"""
                    SELECT DISTINCT caller FROM calls 
                    WHERE callee IN ({placeholders})
                """, callee_targets)
                callers = [r[0] for r in cursor.fetchall() if r[0] and r[0] not in ["global", name, scoped_caller]]

                # Local Git blame directly from original source on disk
                git_data = None
                if target_path:
                    full_file = target_path / filepath
                    if full_file.exists():
                        git_data = get_git_authors(str(full_file), start_line, end_line)

                return {
                    "id": node_id, "name": name, "type": node_type, "filepath": filepath,
                    "line_start": start_line, "line_end": end_line, "code": snippet,
                    "git": git_data,
                    "callers": sorted(list(set(callers))),
                    "callees": sorted(list(set(callees)))
                }
    finally:
        conn.close()

    return {"id": node_id, "type": "module/package", "message": "No code snippet available."}


@app.post("/api/chat/{project_name}")
async def chat_endpoint(project_name: str, req: ChatRequest):
    db_path = str((PROJECTS_DIR / project_name / "atlas.db").resolve())
    ai = CodeAtlasAI(db_path)
    return StreamingResponse(
        ai.stream_chat(req.node_id, req.message, req.selected_model), 
        media_type="text/plain"
    )

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