app/main_api.py

import base64
import os
import subprocess
import threading
import logging
from typing import Any, Dict, List, Optional, Union
from app.vision import extract_text_from_image, describe_image

import uvicorn
from fastapi import (
    FastAPI, APIRouter, Depends, UploadFile, File, Form, HTTPException, status
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel

from app import chat_memory
from app.auth import (
    USERS_DB, authenticate_user, get_password_hash,
    create_access_token, get_current_user, verify_token
)
from app.agent import LocalLLMAgent
from app.langchain_agent import make_local_agent
from app.email_tool import generate_email
from app.files_api import save_upload
from app.embeddings import DocStore, embed_file, query_file_chunks
# ✅ Option A - import the actual objects
from app.chat_memory import persistent_memory, chat_history, semantic_search

from app.tools import TOOLS, use_tool, get_tools
from app.vision import extract_text_from_image, caption_image
from app.audio_tool import transcribe_audio, text_to_speech

# Initialize app
app = FastAPI(title="🧠 LLaMA Local Agent")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

docs_store = DocStore()
agent = LocalLLMAgent("models/capybarahermes-2.5-mistral-7b.Q5_K_S.gguf", docs_store)

chain_agent = make_local_agent("models/capybarahermes-2.5-mistral-7b.Q5_K_S.gguf")



# === Auth Models ===
class TokenResponse(BaseModel):
    access_token: str
    token_type: str = "bearer"


class RegisterRequest(BaseModel):
    username: str
    password: str
    role: Optional[str] = "user"


# === Routes ===

@app.post("/login", response_model=TokenResponse)
def login(username: str = Form(...), password: str = Form(...)):
    user = authenticate_user(username, password)
    if not user:
        raise HTTPException(status_code=401, detail="Invalid credentials")
    token = create_access_token({"sub": user["username"]})
    return {"access_token": token, "token_type": "bearer"}


@app.post("/register", response_model=TokenResponse)
def register(payload: RegisterRequest):
    if payload.username in USERS_DB:
        raise HTTPException(status_code=400, detail="User exists")
    USERS_DB[payload.username] = {
        "username": payload.username,
        "hashed_password": get_password_hash(payload.password),
        "role": payload.role,
    }
    token = create_access_token({"sub": payload.username})
    return {"access_token": token, "token_type": "bearer"}


@app.get("/me")
def read_current_user(user: dict = Depends(get_current_user)):
    return {"username": user["username"], "role": user.get("role")}


@app.get("/tools", response_model=List[str])
def list_tools():
    return list(TOOLS.keys())


@app.get("/tools/details")
def tool_details():
    return [{"name": t.name, "description": t.description} for t in get_tools()]


class ToolRequest(BaseModel):
    input: Optional[Union[str, Dict, List]] = None


@app.post("/tool/{tool_name}")
def call_tool(tool_name: str, req: ToolRequest, _=Depends(verify_token)):
    data = req.input
    if tool_name not in TOOLS:
        raise HTTPException(404, f"Tool '{tool_name}' not found")
    return {"tool": tool_name, "result": use_tool(tool_name, data)}


@app.post("/chat")
def chat(req: Dict[str, Any], _=Depends(verify_token)):
    prompt = req.get("prompt")
    chat_history.append_message("user", prompt)

    resp = agent.local_llm_chat(prompt)

    chat_history.append_message("ai", resp)
    return {"response": resp}


@app.post("/agent")
def run_chain(req: Dict[str, Any], _=Depends(verify_token)):
    prompt = req.get("prompt")
    return {"result": chain_agent.run(prompt)}


@app.post("/upload")
async def upload_file(file: UploadFile = File(...), _=Depends(verify_token)):
    path = save_upload(file)
    embed_file(path)
    return {"path": path}


@app.get("/docs")
def list_docs():
    return os.listdir("uploaded_files")


@app.delete("/docs/{name}")
def delete_doc(name: str):
    path = os.path.join("uploaded_files", name)
    if os.path.exists(path):
        os.remove(path)
        return {"deleted": name}
    raise HTTPException(404, "Not found")


@app.post("/ask-doc")
def ask_doc(req: Dict[str, Any], _=Depends(verify_token)):
    prompt = req.get("prompt")
    return agent.ask_doc(prompt)


@app.get("/query_file")
def query_file(filename: str, question: str):
    return query_file_chunks(filename, question)


@app.post("/image-caption")
async def img_caption(file: UploadFile = File(...), _=Depends(verify_token)):
    tmp = save_upload(file)
    caption = caption_image(tmp)
    return {"caption": caption}



@app.post("/ocr")
def ocr_image(base64_image: str):
    return {"text": extract_text_from_image(base64_image)}

@app.post("/caption")
def caption_image_api(base64_image: str):
    return {"caption": describe_image(base64.b64decode(base64_image))}

@app.post("/transcribe")
async def transcribe(file: UploadFile = File(...), _=Depends(verify_token)):
    tmp = save_upload(file)
    return {"transcription": transcribe_audio(tmp)}


@app.post("/speak")
def speak(text: str, _=Depends(verify_token)):
    mp3 = text_to_speech(text)
    return FileResponse(mp3, media_type="audio/mpeg")


@app.get("/generate_email")
def email_gen(to: str, product: str, discount: float, _=Depends(verify_token)):
    return {"email": generate_email(to, product, discount)}


@app.get("/history/export")
def export_history():
    return {"text": chat_memory.chat_history.export_history()}

@app.get("/search")
def search_chat(query: str):
    return {"matches": chat_memory.chat_history.search_history(query)}
@app.get("/memory/stats")
def memory_stats(_=Depends(verify_token)):
    return {"size": len(agent.mem.db.all())}


@app.post("/reset")
def reset_memory(_=Depends(verify_token)):
    agent.reset()
    return {"status": "cleared"}


# --- Launch both backend & frontend concurrently (optional) ---
def run_backend():
    uvicorn.run("app.main_api:app", host="0.0.0.0", port=8000, reload=True)


def run_frontend():
    subprocess.run(["streamlit", "run", "frontend/streamlit_app.py"])


if __name__ == "__main__":
    t1 = threading.Thread(target=run_backend)
    t2 = threading.Thread(target=run_frontend)
    t1.start()
    t2.start()
    t1.join()
    t2.join()