import os
import shutil
import json
import urllib.parse
from typing import Optional
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
import whisper
from llama_cpp import Llama
from duckduckgo_search import DDGS
from huggingface_hub import InferenceClient

app = FastAPI()

# Cấu hình thư mục lưu ảnh tạo ra
os.makedirs("static/images", exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")

print("Đang tải Whisper-tiny (Audio)...")
whisper_model = whisper.load_model("tiny")

# Khởi tạo API tạo ảnh miễn phí (Sử dụng Flux hoặc Stable Diffusion)
image_client = InferenceClient("black-forest-labs/FLUX.1-schnell")

ALLOWED_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".mp4", ".mov", ".mp3", ".wav", ".txt", ".md", ".csv"}

MODEL_PATHS = {
    "flash": {"repo": "Qwen/Qwen2.5-VL-3B-Instruct-GGUF", "filename": "qwen2.5-vl-3b-instruct-q4_k_m.gguf"},
    "pro": {"repo": "Qwen/Qwen2.5-VL-7B-Instruct-GGUF", "filename": "qwen2.5-vl-7b-instruct-q4_k_m.gguf"}
}

current_model = None
current_mode = None

def get_model(mode: str):
    global current_model, current_mode
    if current_mode != mode:
        print(f"🔄 Đang chuyển lõi AI sang: {mode.upper()}...")
        current_model = None
        info = MODEL_PATHS[mode]
        current_model = Llama.from_pretrained(
            repo_id=info["repo"],
            filename=info["filename"],
            n_ctx=4096,
            n_threads=2,
            chat_format="qwen2-vl"
        )
        current_mode = mode
    return current_model

def search_web(query: str) -> str:
    try:
        results = [f"- {r['title']}: {r['body']}" for r in DDGS().text(query, max_results=3)]
        return "\n".join(results) if results else ""
    except:
        return ""

@app.post("/api/chat")
async def chat_api(
    prompt: str = Form(""),
    mode: str = Form("flash"),
    web_search: bool = Form(False),
    history: str = Form("[]"), # Trí nhớ được Frontend gửi lên
    file: Optional[UploadFile] = File(None)
):
    try:
        # 1. KIỂM TRA LỆNH TẠO ẢNH (Bỏ qua LLM, chạy thẳng Image API để tiết kiệm Token/Thời gian)
        if prompt.lower().startswith(("tạo ảnh", "vẽ ảnh", "vẽ hình", "tạo hình")):
            # Dịch ngầm hoặc truyền thẳng prompt tiếng Việt (Flux hiểu khá tốt)
            image_prompt = prompt.replace("tạo ảnh", "").replace("vẽ ảnh", "").strip()
            image = image_client.text_to_image(image_prompt)
            image_name = f"gen_{os.urandom(4).hex()}.png"
            image_path = f"static/images/{image_name}"
            image.save(image_path)
            
            return JSONResponse({
                "status": "success",
                "response": f"🎨 Tôi đã tạo ảnh theo yêu cầu của bạn:\n\n![{image_prompt}](/{image_path})"
            })

        # 2. XỬ LÝ FILE ĐÍNH KÈM
        extracted_text = ""
        image_url_path = None
        if file and file.filename:
            ext = os.path.splitext(file.filename)[1].lower()
            if ext not in ALLOWED_EXTS:
                raise HTTPException(status_code=400, detail="Định dạng file không được hỗ trợ.")
            
            temp_path = f"temp_{file.filename}"
            with open(temp_path, "wb") as f:
                shutil.copyfileobj(file.file, f)

            if ext in {".jpg", ".jpeg", ".png", ".webp"}:
                image_url_path = temp_path
            elif ext in {".mp3", ".wav"}:
                extracted_text += f"\n[Âm thanh]: {whisper_model.transcribe(temp_path)['text']}"
                os.remove(temp_path)
            elif ext in {".txt", ".md", ".csv"}:
                with open(temp_path, "r", encoding="utf-8", errors="ignore") as doc:
                    extracted_text += f"\n[Tài liệu]:\n{doc.read()}"
                os.remove(temp_path)

        # 3. XỬ LÝ TRA WEB
        search_ctx = ""
        if web_search and prompt.strip():
            web_data = search_web(prompt)
            if web_data:
                search_ctx = f"\n[Dữ liệu Web cập nhật]:\n{web_data}\n"

        # 4. XỬ LÝ TRÍ NHỚ (LIGHT MEMORY ~ 100 TOKENS)
        # Lấy lịch sử dạng JSON, chỉ giữ lại 2 đoạn hội thoại gần nhất để làm tiền ngữ cảnh
        chat_history = json.loads(history)
        memory_context = ""
        if chat_history:
            # Lấy 2 tin nhắn cuối, giới hạn độ dài để không tràn token
            recent_msgs = chat_history[-2:]
            memory_context = "\n[Lịch sử ngắn]:\n" + "\n".join([f"{m['role']}: {m['content'][:200]}" for m in recent_msgs]) + "\n"

        # 5. GỘP PROMPT
        system_msg = "Bạn là AI thông minh. Năm nay là 2026. Hãy trả lời ngắn gọn, chính xác."
        final_prompt = f"{system_msg}{memory_context}{search_ctx}{extracted_text}\nNgười dùng: {prompt}"

        # 6. GỌI LLM
        messages = []
        if image_url_path:
            messages.append({
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": image_url_path}},
                    {"type": "text", "text": final_prompt}
                ]
            })
        else:
            messages.append({"role": "user", "content": final_prompt})

        llm = get_model(mode)
        response = llm.create_chat_completion(messages=messages, max_tokens=1024, temperature=0.7)

        if image_url_path and os.path.exists(image_url_path):
            os.remove(image_url_path)

        return JSONResponse({"status": "success", "response": response["choices"][0]["message"]["content"]})

    except Exception as e:
        return JSONResponse({"status": "error", "message": str(e)}, status_code=500)

@app.get("/", response_class=HTMLResponse)
async def read_index():
    if os.path.exists("index.html"):
        with open("index.html", "r", encoding="utf-8") as f:
            return f.read()
    return "<h1>Chưa tìm thấy index.html</h1>"