import os
import re
import requests
import gradio as gr
from huggingface_hub import InferenceClient

# ========= CONFIG =========
CHAT_MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
ROUTER_URL = "https://router.huggingface.co/v1/chat/completions"
ASR_MODEL_ID = "openai/whisper-large-v3"
DEFAULT_SYSTEM_PROMPT = "You are a helpful AI assistant."
MAX_FILE_CHARS = 40_000

HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise RuntimeError("HF_TOKEN is missing. Add it in Space Settings → Secrets as HF_TOKEN.")

HEADERS = {
    "Authorization": f"Bearer {HF_TOKEN}",
    "Content-Type": "application/json",
}

asr_client = InferenceClient(api_key=HF_TOKEN)

# ========= HELPERS =========
def normalize_whitespace(s: str) -> str:
    return re.sub(r"\s+", " ", (s or "")).strip()

def safe_read_text_file(path: str) -> str:
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        return f.read()

def extract_text_from_pdf(path: str) -> str:
    try:
        from pypdf import PdfReader
    except Exception:
        return "PDF support not installed. Add 'pypdf' to requirements.txt."
    reader = PdfReader(path)
    parts = []
    for page in reader.pages[:50]:
        parts.append(page.extract_text() or "")
    return "\n".join(parts)

def extract_file_text(file_path: str) -> str:
    ext = os.path.splitext(file_path)[1].lower()
    if ext in [".txt", ".md", ".csv", ".json", ".py", ".js", ".html", ".css", ".log", ".yaml", ".yml"]:
        return safe_read_text_file(file_path)
    if ext == ".pdf":
        return extract_text_from_pdf(file_path)
    return ""

def call_llm(messages, temperature=0.7, max_tokens=400):
    payload = {
        "model": CHAT_MODEL_ID,
        "messages": messages,
        "temperature": float(temperature),
        "max_tokens": int(max_tokens),
    }
    resp = requests.post(ROUTER_URL, headers=HEADERS, json=payload, timeout=120)
    if resp.status_code != 200:
        try:
            err = resp.json()
        except Exception:
            err = resp.text
        raise RuntimeError(f"HTTP {resp.status_code}: {err}")
    data = resp.json()
    return data["choices"][0]["message"]["content"]

def chat_to_llm_messages(chat_messages, system_prompt: str, file_context: str, user_text: str):
    msgs = [{"role": "system", "content": (system_prompt.strip() or DEFAULT_SYSTEM_PROMPT)}]

    if (file_context or "").strip():
        msgs.append({
            "role": "system",
            "content": "User uploaded file content (use this as reference):\n" + file_context
        })

    for m in (chat_messages or []):
        role = (m.get("role") or "").lower()
        content = m.get("content") or ""
        if role in ("user", "assistant") and content:
            msgs.append({"role": role, "content": content})

    msgs.append({"role": "user", "content": user_text})
    return msgs

# ========= ACTIONS =========
def on_upload_file(chat_messages, file_obj, file_context):
    if file_obj is None:
        return chat_messages, file_context, "No file selected"

    file_path = getattr(file_obj, "name", None) or str(file_obj)
    filename = os.path.basename(file_path)
    file_size = os.path.getsize(file_path) if os.path.exists(file_path) else 0
    
    extracted = extract_file_text(file_path)
    if extracted:
        extracted = extracted[:MAX_FILE_CHARS]
        file_context = f"[{filename}]\n{extracted}\n\n" + (file_context or "")
        chat_messages = (chat_messages or []) + [
            {"role": "user", "content": f"📎 Uploaded file: {filename}"},
            {"role": "assistant", "content": f"✅ **{filename}** loaded ({file_size:,} bytes). Ask me anything about this document!"},
        ]
        file_info = f"✅ Loaded: {filename} ({file_size:,} bytes)"
    else:
        chat_messages = (chat_messages or []) + [
            {"role": "user", "content": f"📎 Uploaded file: {filename}"},
            {"role": "assistant", "content": f"📄 **{filename}** received. I'll answer based on its content."},
        ]
        file_info = f"📄 Uploaded: {filename}"

    return chat_messages, file_context, file_info

def on_record_audio(audio_path):
    if not audio_path:
        return ""
    try:
        out = asr_client.automatic_speech_recognition(audio_path, model=ASR_MODEL_ID)
        text = getattr(out, "text", "") if out is not None else ""
        return normalize_whitespace(text)
    except Exception as e:
        return f"[Voice recognition error: {str(e)[:50]}]"

def on_send(chat_messages, user_text, system_prompt, temperature, max_tokens, file_context):
    user_text = (user_text or "").strip()
    if not user_text:
        return chat_messages, ""

    chat_messages = chat_messages or []
    chat_messages.append({"role": "user", "content": user_text})

    try:
        llm_messages = chat_to_llm_messages(chat_messages, system_prompt, file_context or "", user_text)
        assistant = call_llm(llm_messages, temperature=temperature, max_tokens=max_tokens)
        chat_messages.append({"role": "assistant", "content": assistant})
        return chat_messages, ""
    except Exception as e:
        chat_messages.append({"role": "assistant", "content": f"⚠️ Error: {type(e).__name__}: {str(e)[:100]}"})
        return chat_messages, ""

def on_clear():
    return [], "", "No file selected"

# ========= SIMPLE UI =========
with gr.Blocks(title="AI Chat Assistant") as demo:
    file_context_state = gr.State("")
    chat_state = gr.State([])
    
    # Header
    with gr.Column():
        gr.Markdown("# 🤖 AI Chat Assistant")
        gr.Markdown("### Your intelligent assistant powered by Llama 3.1")
        gr.Markdown(f"**Model:** `{CHAT_MODEL_ID.split('/')[-1]}`")
    
    # Chat display
    chatbot = gr.Chatbot(height=400, label="Chat History")
    
    # File upload section
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📁 Upload Document")
            upload = gr.File(
                label="",
                file_count="single",
                file_types=[".txt", ".pdf", ".md", ".csv", ".json", ".py"]
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### 🎤 Voice Input")
            mic = gr.Audio(
                sources=["microphone"],
                type="filepath",
                label=""
            )
        
        with gr.Column(scale=2):
            file_info_display = gr.Textbox(
                label="📊 File Status",
                value="No file selected",
                interactive=False
            )
    
    # Settings in accordion
    with gr.Accordion("⚙️ Settings", open=False):
        system_prompt = gr.Textbox(
            value=DEFAULT_SYSTEM_PROMPT,
            label="System Prompt",
            lines=3
        )
        
        with gr.Row():
            temperature = gr.Slider(
                minimum=0.0,
                maximum=1.5,
                value=0.7,
                step=0.05,
                label="Temperature (Creativity)"
            )
            
            max_tokens = gr.Slider(
                minimum=64,
                maximum=1024,
                value=400,
                step=64,
                label="Max Response Length"
            )
    
    # Input area
    with gr.Row():
        user_input = gr.Textbox(
            placeholder="Type your message here...",
            label="Your Message",
            lines=3,
            scale=4
        )
        
        with gr.Column(scale=1):
            send_btn = gr.Button("🚀 Send", variant="primary", size="lg")
            clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")
    
    # ========= EVENT HANDLERS =========
    # File upload
    upload.change(
        fn=on_upload_file,
        inputs=[chat_state, upload, file_context_state],
        outputs=[chat_state, file_context_state, file_info_display]
    ).then(
        fn=lambda x: x,
        inputs=[chat_state],
        outputs=[chatbot]
    )
    
    # Voice input
    mic.change(
        fn=on_record_audio,
        inputs=[mic],
        outputs=[user_input]
    )
    
    # Send message (button)
    send_btn.click(
        fn=on_send,
        inputs=[chat_state, user_input, system_prompt, temperature, max_tokens, file_context_state],
        outputs=[chat_state, user_input]
    ).then(
        fn=lambda x: x,
        inputs=[chat_state],
        outputs=[chatbot]
    )
    
    # Send message (enter)
    user_input.submit(
        fn=on_send,
        inputs=[chat_state, user_input, system_prompt, temperature, max_tokens, file_context_state],
        outputs=[chat_state, user_input]
    ).then(
        fn=lambda x: x,
        inputs=[chat_state],
        outputs=[chatbot]
    )
    
    # Clear chat
    clear_btn.click(
        fn=on_clear,
        inputs=[],
        outputs=[chat_state, file_context_state, file_info_display]
    ).then(
        fn=lambda: [],
        inputs=[],
        outputs=[chatbot]
    )

# Launch without CSS to avoid issues
demo.launch(share=False, server_name="0.0.0.0", server_port=7860)