import torch
import torch.nn.functional as F
import pickle
import os
import re
import sys
import csv
import time
import json as json_module
import uuid
import asyncio
from typing import List, Optional
from fastapi import FastAPI, HTTPException, Path
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from contextlib import asynccontextmanager
import sqlite3

if sys.stdout.encoding != 'utf-8':
    sys.stdout.reconfigure(encoding='utf-8')

# ── Paths ──
VOCAB_PATH      = 'data/processed/vocab.pkl'
MODEL_PATH      = 'models/best/model.pt'
TOKENIZER_PATH  = 'data/processed/bpe_tokenizer.json'
RAW_CSV_PATH    = 'data/raw/output.csv'
DB_PATH         = 'data/chatbot.sqlite'

# ── Model globals ──
model          = None
word2idx       = None
idx2word       = None
device         = None
model_info     = {}
bpe_tokenizer  = None

qa_database       = []
qa_embeddings     = None
SIMILARITY_THRESHOLD = 0.92


# ═══════════════════════════════════════
# DATABASE – Sessions + Messages
# ═══════════════════════════════════════

def init_db():
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute("PRAGMA foreign_keys = ON")
    c.execute('''
        CREATE TABLE IF NOT EXISTS sessions (
            id TEXT PRIMARY KEY,
            title TEXT NOT NULL DEFAULT 'Chat mới',
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    c.execute('''
        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            role TEXT NOT NULL,
            content TEXT NOT NULL,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
        )
    ''')
    # Migrate old chat_history table if it exists
    c.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='chat_history'")
    if c.fetchone():
        c.execute("SELECT COUNT(*) FROM chat_history")
        count = c.fetchone()[0]
        if count > 0:
            sid = str(uuid.uuid4())
            c.execute("INSERT INTO sessions (id, title) VALUES (?, ?)", (sid, "Chat cũ"))
            c.execute("""
                INSERT INTO messages (session_id, role, content, timestamp)
                SELECT ?, role, content, timestamp FROM chat_history
            """, (sid,))
            print(f"[DB] Migrated {count} old messages to session '{sid}'")
        c.execute("DROP TABLE chat_history")
    conn.commit()
    conn.close()


def save_message(session_id, role, content):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute("INSERT INTO messages (session_id, role, content) VALUES (?, ?, ?)",
              (session_id, role, content))
    c.execute("UPDATE sessions SET updated_at = CURRENT_TIMESTAMP WHERE id = ?",
              (session_id,))
    conn.commit()
    conn.close()


def create_session(title="Chat mới"):
    sid = str(uuid.uuid4())
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute("INSERT INTO sessions (id, title) VALUES (?, ?)", (sid, title))
    conn.commit()
    conn.close()
    return sid


def update_session_title(session_id, title):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute("UPDATE sessions SET title = ? WHERE id = ?", (title, session_id))
    conn.commit()
    conn.close()


# ═══════════════════════════════════════
# MODEL LOADING
# ═══════════════════════════════════════

def load_bpe_tokenizer():
    global bpe_tokenizer
    if os.path.exists(TOKENIZER_PATH):
        from tokenizers import Tokenizer
        bpe_tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
        print(f"[OK] BPE Tokenizer loaded | Vocab: {bpe_tokenizer.get_vocab_size()}")
        return True
    else:
        print(f"[WARN] Không tìm thấy BPE tokenizer: {TOKENIZER_PATH}")
        return False


def tokenize_bpe(text):
    if bpe_tokenizer is not None:
        enc = bpe_tokenizer.encode(str(text).strip())
        return enc.ids
    else:
        text = str(text).lower()
        text = re.sub(r'([.,!?()])', r' \\1 ', text)
        text = re.sub(r'\\s{2,}', ' ', text)
        tokens = text.strip().split()
        return [word2idx.get(w, word2idx.get('<unk>', 3)) for w in tokens]


def clean_response(text):
    text = text.strip()
    if not text:
        return ""
    text = text[0].upper() + text[1:]
    import unicodedata
    text = unicodedata.normalize('NFC', text)
    text = re.sub(r'\\\s+([a-zA-Z]+)', r'\\\1', text)
    text = re.sub(r'\s*_\s*', '_', text)
    text = re.sub(r'\s*\^\s*', '^', text)
    text = re.sub(r'\s*\{\s*', '{', text)
    text = re.sub(r'\s*\}\s*', '}', text)
    text = re.sub(r'\s+([.,!?():;])', r'\1', text)
    text = re.sub(r'\(\s+', r'(', text)
    text = re.sub(r'\s{2,}', ' ', text)
    return text.strip()


def generate_response(message, temperature=0.5, top_k=3, top_p=0.9,
                      repetition_penalty=1.3, max_tokens=100):
    """Core inference function. Returns (response_text, tokens_count, elapsed_ms)."""
    start = time.perf_counter()

    # Greeting check
    msg_lower = message.lower()
    if len(msg_lower) < 25:
        if any(w in msg_lower for w in ["chào", "hello", "hi ", "xin chào", "alo"]):
            return "Chào bạn! Mình là trợ lý môn Toán Rời Rạc. Bạn cần hỏi gì nào?", 0, 1.0
        if any(w in msg_lower for w in ["cảm ơn", "thank", "cám ơn", "tks"]):
            return "Không có gì! Chúc bạn học tốt môn Toán Rời Rạc nhé!", 0, 1.0

    tokens = tokenize_bpe(message)
    if not tokens:
        return "Mình không hiểu câu hỏi. Bạn thử hỏi lại nhé!", 0, 1.0

    sos_id = word2idx.get('<sos>', 1)
    sep_id = word2idx.get('<sep>', 4)

    prompt = [sos_id] + tokens + [sep_id]
    x = torch.tensor([prompt]).to(device)

    with torch.no_grad():
        out_ids = model.generate(x, max_new_tokens=max_tokens,
                                 temperature=temperature,
                                 top_k=top_k, top_p=top_p,
                                 repetition_penalty=repetition_penalty)

    elapsed_ms = (time.perf_counter() - start) * 1000
    resp_ids = out_ids[0][len(prompt):]

    if bpe_tokenizer is not None:
        valid_ids = []
        for i in resp_ids:
            tid = i.item()
            if tid == word2idx.get('<eos>', 2):
                break
            if tid not in [word2idx.get('<pad>', 0), word2idx.get('<sos>', 1),
                           word2idx.get('<unk>', 3), word2idx.get('<sep>', 4)]:
                valid_ids.append(tid)
        response = bpe_tokenizer.decode(valid_ids)
    else:
        resp_tokens = []
        for i in resp_ids:
            sym = idx2word.get(i.item(), '<unk>')
            if sym == '<eos>':
                break
            if sym not in ['<pad>', '<sos>', '<unk>', '<sep>']:
                resp_tokens.append(sym)
        response = " ".join(resp_tokens)

    response = clean_response(response)
    if not response.strip():
        response = "Mình chưa học tới phần này. Bạn hỏi câu khác nhé!"

    return response, len(resp_ids), round(elapsed_ms, 2)


async def generate_response_stream(message, temperature=0.5, top_k=3, top_p=0.9,
                                   repetition_penalty=1.3, max_tokens=100):
    """Streaming inference function. Yields SSE data chunks containing the accumulated full text."""
    msg_lower = message.lower()
    if len(msg_lower) < 25:
        if any(w in msg_lower for w in ["chào", "hello", "hi ", "xin chào", "alo"]):
            yield "data: " + json_module.dumps({"text": "Chào bạn! Mình là trợ lý môn Toán Rời Rạc. Bạn cần hỏi gì nào?", "done": True}, ensure_ascii=False) + "\n\n"
            return
        if any(w in msg_lower for w in ["cảm ơn", "thank", "cám ơn", "tks"]):
            yield "data: " + json_module.dumps({"text": "Không có gì! Chúc bạn học tốt môn Toán Rời Rạc nhé!", "done": True}, ensure_ascii=False) + "\n\n"
            return

    tokens = tokenize_bpe(message)
    if not tokens:
        yield "data: " + json_module.dumps({"text": "Mình không hiểu câu hỏi. Bạn thử hỏi lại nhé!", "done": True}, ensure_ascii=False) + "\n\n"
        return

    sos_id = word2idx.get('<sos>', 1)
    sep_id = word2idx.get('<sep>', 4)
    eos_id = word2idx.get('<eos>', 2)
    pad_id = word2idx.get('<pad>', 0)
    unk_id = word2idx.get('<unk>', 3)

    prompt = [sos_id] + tokens + [sep_id]
    x = torch.tensor([prompt]).to(device)

    valid_ids = []
    
    with torch.no_grad():
        for tid in model.generate_stream(x, max_new_tokens=max_tokens,
                                         temperature=temperature, top_k=top_k,
                                         top_p=top_p, repetition_penalty=repetition_penalty):
            if tid == eos_id:
                break
            if tid not in [pad_id, sos_id, unk_id, sep_id]:
                valid_ids.append(tid)
                
                if bpe_tokenizer is not None:
                    current_text = bpe_tokenizer.decode(valid_ids)
                else:
                    current_text = " ".join([idx2word.get(i, '<unk>') for i in valid_ids])
                
                current_text = clean_response(current_text)
                yield "data: " + json_module.dumps({"text": current_text, "done": False}, ensure_ascii=False) + "\n\n"
            
            await asyncio.sleep(0.04)

    if not valid_ids:
        yield "data: " + json_module.dumps({"text": "Mình chưa học tới phần này. Bạn hỏi câu khác nhé!", "done": True}, ensure_ascii=False) + "\n\n"
    else:
        yield "data: " + json_module.dumps({"done": True}, ensure_ascii=False) + "\n\n"


def load_model():
    global model, word2idx, idx2word, device, model_info
    sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
    from model import make_model

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if not os.path.exists(VOCAB_PATH):
        print(f"[WARN] Không tìm thấy vocab: {VOCAB_PATH}")
        return False
    if not os.path.exists(MODEL_PATH):
        print(f"[WARN] Không tìm thấy model: {MODEL_PATH}")
        return False

    with open(VOCAB_PATH, 'rb') as f:
        vocab = pickle.load(f)
    word2idx = vocab['word2idx']
    idx2word = vocab['idx2word']

    ckpt = torch.load(MODEL_PATH, map_location=device)
    d_m = ckpt.get('d_m', 768)
    n_l = ckpt.get('n_l', 12)
    n_h = ckpt.get('n_h', 12)
    dr  = ckpt.get('dr', 0.1)

    model = make_model(len(word2idx), n_layer=n_l, d_model=d_m, n_head=n_h, dropout=dr)
    model.load_state_dict(ckpt['model_state_dict'])
    model.to(device)
    model.eval()

    model_info = {
        "vocab_size": len(word2idx),
        "d_model": d_m,
        "n_layers": n_l,
        "n_heads": n_h,
        "device": str(device),
    }
    print(f"[OK] Model loaded on {device} | d_model={d_m} | layers={n_l}")
    return True


# ═══════════════════════════════════════
# FASTAPI APP
# ═══════════════════════════════════════

@asynccontextmanager
async def lifespan(app: FastAPI):
    success = load_model()
    load_bpe_tokenizer()
    init_db()
    if success:
        print("[INFO] Server ready.")
    else:
        print("[WARN] Server khởi động KHÔNG có model.")
    yield


app = FastAPI(title="Toán Rời Rạc Chatbot", version="3.0.0", lifespan=lifespan)

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


# ── Health ──
@app.get("/api/health")
async def health_check():
    return {
        "status": "ok",
        "model_loaded": model is not None,
        "model_info": model_info,
        "bpe_tokenizer": bpe_tokenizer is not None,
    }


# ── Sessions ──
@app.get("/api/sessions")
async def list_sessions():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    c = conn.cursor()
    c.execute("""
        SELECT s.id, s.title, s.created_at, s.updated_at,
               (SELECT COUNT(*) FROM messages m WHERE m.session_id = s.id) as msg_count
        FROM sessions s ORDER BY s.updated_at DESC
    """)
    rows = c.fetchall()
    conn.close()
    return [dict(r) for r in rows]


@app.post("/api/sessions")
async def create_session_endpoint():
    sid = create_session()
    return {"id": sid, "title": "Chat mới"}


@app.delete("/api/sessions/{session_id}")
async def delete_session(session_id: str = Path(...)):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute("PRAGMA foreign_keys = ON")
    c.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
    c.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
    conn.commit()
    conn.close()
    return {"status": "ok"}


@app.get("/api/sessions/{session_id}/messages")
async def get_session_messages(session_id: str = Path(...)):
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    c = conn.cursor()
    c.execute("SELECT role, content, timestamp FROM messages WHERE session_id = ? ORDER BY id ASC",
              (session_id,))
    rows = c.fetchall()
    conn.close()
    return [dict(r) for r in rows]


# ── Streaming Chat ──
class StreamChatRequest(BaseModel):
    message: str
    session_id: Optional[str] = None
    temperature: float = 0.5
    top_k: int = 3
    top_p: float = 0.9
    repetition_penalty: float = 1.3
    max_tokens: int = 100


@app.post("/api/chat/stream")
async def chat_stream(req: StreamChatRequest):
    if model is None:
        raise HTTPException(status_code=503, detail="Model chưa được load.")

    message = req.message.strip()
    if not message:
        raise HTTPException(status_code=400, detail="Tin nhắn không được để trống.")

    # Create or use existing session
    session_id = req.session_id
    is_new_session = False
    if not session_id:
        session_id = create_session(message[:40])
        is_new_session = True

    # Save user message
    save_message(session_id, "user", message)

    # If new session, update title with first message
    if is_new_session:
        title = message[:40] + ("..." if len(message) > 40 else "")
        update_session_title(session_id, title)

    async def event_stream():
        # Send a 1024-byte padding comment to bypass browser buffering
        yield f": {' ' * 1024}\n\n"

        # Send session info
        yield f"data: {json_module.dumps({'type': 'session', 'session_id': session_id, 'is_new': is_new_session})}\n\n"

        # Send thinking indicator
        yield f"data: {json_module.dumps({'type': 'thinking'})}\n\n"
        await asyncio.sleep(0.05)

        # Generate response using real streaming
        token_count = 0
        start_time = time.perf_counter()
        final_response = ""
        
        async for chunk in generate_response_stream(
            message,
            temperature=req.temperature,
            top_k=req.top_k,
            top_p=req.top_p,
            repetition_penalty=req.repetition_penalty,
            max_tokens=req.max_tokens,
        ):
            # Parse chunk JSON
            data = json_module.loads(chunk[6:]) # strip 'data: '
            if not data.get("done"):
                final_response = data["text"]
                yield f"data: {json_module.dumps({'type': 'token', 'content': final_response})}\n\n"
            else:
                if "text" in data:
                    final_response = data["text"]
                    yield f"data: {json_module.dumps({'type': 'token', 'content': final_response})}\n\n"
            token_count += 1
            
        elapsed_ms = (time.perf_counter() - start_time) * 1000

        # Save bot response
        save_message(session_id, "bot", final_response)

        # Send done
        yield f"data: {json_module.dumps({'type': 'done', 'tokens': token_count, 'time_ms': elapsed_ms, 'session_id': session_id})}\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache, no-transform",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
            "X-Content-Type-Options": "nosniff"
        }
    )


# ── Legacy non-streaming (backward compat) ──
class ChatRequest(BaseModel):
    message: str
    session_id: Optional[str] = None
    history: Optional[list] = None
    temperature: float = 0.5
    top_k: int = 50
    top_p: float = 0.9
    repetition_penalty: float = 1.3
    max_tokens: int = 100


@app.post("/api/chat")
async def chat(req: ChatRequest):
    if model is None:
        raise HTTPException(status_code=503, detail="Model chưa được load.")

    message = req.message.strip()
    if not message:
        raise HTTPException(status_code=400, detail="Tin nhắn không được để trống.")

    response, token_count, elapsed_ms = generate_response(
        message, req.temperature, req.top_k, req.top_p,
        req.repetition_penalty, req.max_tokens
    )

    return {"response": response, "tokens_generated": token_count, "inference_time_ms": elapsed_ms}


# ── Static files ──
web_dir = os.path.join(os.path.dirname(__file__), "web")
if os.path.isdir(web_dir):
    app.mount("/assets", StaticFiles(directory=os.path.join(web_dir, "assets")), name="assets")

    @app.get("/")
    async def serve_frontend():
        return FileResponse(
            os.path.join(web_dir, "index.html"),
            headers={"Cache-Control": "no-cache, no-store, must-revalidate", "Pragma": "no-cache", "Expires": "0"}
        )


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("api:app", host="127.0.0.1", port=8000, reload=True)