import sys
import os
import pandas as pd
import json
import gradio as gr
from typing import List, Tuple, Dict, Any, Union
import hashlib
import shutil
import re
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed

# Setup directories
persistent_dir = "/data/hf_cache"
os.makedirs(persistent_dir, exist_ok=True)

model_cache_dir = os.path.join(persistent_dir, "txagent_models")
tool_cache_dir = os.path.join(persistent_dir, "tool_cache")
file_cache_dir = os.path.join(persistent_dir, "cache")
report_dir = os.path.join(persistent_dir, "reports")

for d in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir]:
    os.makedirs(d, exist_ok=True)

os.environ["HF_HOME"] = model_cache_dir
os.environ["TRANSFORMERS_CACHE"] = model_cache_dir

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "src")))
from txagent.txagent import TxAgent

MAX_MODEL_TOKENS = 32768
MAX_CHUNK_TOKENS = 8192
MAX_NEW_TOKENS = 2048
PROMPT_OVERHEAD = 500

def clean_response(text: str) -> str:
    text = re.sub(r"\[.*?\]|\bNone\b", "", text, flags=re.DOTALL)
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = re.sub(r"[^\n#\-\*\w\s\.,:\(\)]+", "", text)
    return text.strip()

def estimate_tokens(text: str) -> int:
    return len(text) // 3.5 + 1

def extract_text_from_excel(file_path: str) -> str:
    all_text = []
    xls = pd.ExcelFile(file_path)
    for sheet_name in xls.sheet_names:
        df = xls.parse(sheet_name).astype(str).fillna("")
        rows = df.apply(lambda row: " | ".join(row), axis=1)
        sheet_text = [f"[{sheet_name}] {line}" for line in rows]
        all_text.extend(sheet_text)
    return "\n".join(all_text)

def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> List[str]:
    effective_max = max_tokens - PROMPT_OVERHEAD
    lines, chunks, curr_chunk, curr_tokens = text.split("\n"), [], [], 0
    for line in lines:
        t = estimate_tokens(line)
        if curr_tokens + t > effective_max:
            if curr_chunk:
                chunks.append("\n".join(curr_chunk))
            curr_chunk, curr_tokens = [line], t
        else:
            curr_chunk.append(line)
            curr_tokens += t
    if curr_chunk:
        chunks.append("\n".join(curr_chunk))
    return chunks

def build_prompt_from_text(chunk: str) -> str:
    return f"""
### Unstructured Clinical Records

Analyze the following clinical notes and provide a detailed, concise summary focusing on:
- Diagnostic Patterns
- Medication Issues
- Missed Opportunities
- Inconsistencies
- Follow-up Recommendations

---

{chunk}

---
Respond in well-structured bullet points with medical reasoning.
"""

def init_agent():
    tool_path = os.path.join(tool_cache_dir, "new_tool.json")
    if not os.path.exists(tool_path):
        shutil.copy(os.path.abspath("data/new_tool.json"), tool_path)
    agent = TxAgent(
        model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
        rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
        tool_files_dict={"new_tool": tool_path},
        force_finish=True,
        enable_checker=True,
        step_rag_num=4,
        seed=100
    )
    agent.init_model()
    return agent

def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Union[str, None]]:
    messages = chatbot_state if chatbot_state else []
    if file is None or not hasattr(file, "name"):
        return messages + [{"role": "assistant", "content": "❌ Please upload a valid Excel file."}], None

    messages.append({"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"})
    text = extract_text_from_excel(file.name)
    chunks = split_text_into_chunks(text)
    chunk_responses = [None] * len(chunks)

    def analyze_chunk(i, chunk):
        prompt = build_prompt_from_text(chunk)
        response = ""
        for res in agent.run_gradio_chat(message=prompt, history=[], temperature=0.2, max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS, call_agent=False, conversation=[]):
            if isinstance(res, str):
                response += res
            elif hasattr(res, "content"):
                response += res.content
            elif isinstance(res, list):
                for r in res:
                    if hasattr(r, "content"):
                        response += r.content
        return i, clean_response(response)

    with ThreadPoolExecutor(max_workers=1) as executor:
        futures = [executor.submit(analyze_chunk, i, c) for i, c in enumerate(chunks)]
        for f in as_completed(futures):
            i, result = f.result()
            chunk_responses[i] = result

    valid = [r for r in chunk_responses if r and not r.startswith("❌")]
    if not valid:
        return messages + [{"role": "assistant", "content": "❌ No valid chunk results."}], None

    summary_prompt = f"Summarize this analysis in a final structured report:\n\n" + "\n\n".join(valid)
    messages.append({"role": "assistant", "content": "📊 Generating final report..."})

    final_report = ""
    for res in agent.run_gradio_chat(message=summary_prompt, history=[], temperature=0.2, max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS, call_agent=False, conversation=[]):
        if isinstance(res, str):
            final_report += res
        elif hasattr(res, "content"):
            final_report += res.content

    cleaned = clean_response(final_report)
    report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md")
    with open(report_path, 'w') as f:
        f.write(f"# 🧠 Final Patient Report\n\n{cleaned}")

    messages.append({"role": "assistant", "content": f"📊 Final Report:\n\n{cleaned}"})
    messages.append({"role": "assistant", "content": f"✅ Report generated and saved: {os.path.basename(report_path)}"})
    return messages, report_path

def create_ui(agent):
    with gr.Blocks(css="""
        :root {
            --primary-color: #2563eb;
            --secondary-color: #1e40af;
            --bg-color: #f8fafc;
            --text-color: #1e293b;
            --user-bubble: #ffffff;
            --bot-bubble: #f1f5f9;
            --border-color: #e2e8f0;
        }
        body {
            font-family: 'Inter', sans-serif;
            background-color: var(--bg-color);
            color: var(--text-color);
        }
        .gradio-container {
            max-width: 800px;
            margin: 0 auto;
            padding: 20px;
            background-color: var(--bg-color);
        }
        .chat-container {
            display: flex;
            flex-direction: column;
            height: 80vh;
            border: 1px solid var(--border-color);
            border-radius: 12px;
            overflow: hidden;
            background-color: white;
        }
        .chat-header {
            padding: 16px;
            background-color: white;
            border-bottom: 1px solid var(--border-color);
            font-weight: 600;
            font-size: 18px;
        }
        .chat-messages {
            flex: 1;
            padding: 16px;
            overflow-y: auto;
            background-color: var(--bg-color);
        }
        .message {
            display: flex;
            margin-bottom: 16px;
            align-items: flex-start;
        }
        .message-avatar {
            width: 36px;
            height: 36px;
            border-radius: 50%;
            margin-right: 12px;
            background-color: var(--primary-color);
            color: white;
            display: flex;
            align-items: center;
            justify-content: center;
            font-weight: bold;
        }
        .message-content {
            max-width: 80%;
        }
        .user-message .message-content {
            margin-left: auto;
            background-color: var(--user-bubble);
            padding: 12px 16px;
            border-radius: 18px 18px 0 18px;
            box-shadow: 0 1px 2px rgba(0,0,0,0.1);
        }
        .bot-message .message-content {
            background-color: var(--bot-bubble);
            padding: 12px 16px;
            border-radius: 18px 18px 18px 0;
            box-shadow: 0 1px 2px rgba(0,0,0,0.1);
        }
        .message-time {
            font-size: 12px;
            color: #64748b;
            margin-top: 4px;
        }
        .chat-input-container {
            padding: 16px;
            background-color: white;
            border-top: 1px solid var(--border-color);
        }
        .file-upload {
            display: flex;
            gap: 8px;
            margin-bottom: 12px;
        }
        .upload-btn {
            flex: 1;
        }
        .analyze-btn {
            background-color: var(--primary-color);
            color: white;
            border: none;
            padding: 10px 16px;
            border-radius: 8px;
            cursor: pointer;
            font-weight: 500;
        }
        .analyze-btn:hover {
            background-color: var(--secondary-color);
        }
        .report-download {
            margin-top: 12px;
            padding: 12px;
            background-color: var(--bot-bubble);
            border-radius: 8px;
            display: none;
        }
    """) as demo:
        with gr.Column(elem_classes="chat-container"):
            gr.Markdown("""<div class="chat-header">Patient History AI Assistant</div>""")
            
            chatbot = gr.Chatbot(
                elem_classes="chat-messages",
                label=None,
                show_label=False,
                bubble_full_width=False,
                avatar_images=[
                    "https://ui-avatars.com/api/?name=AI&background=2563eb&color=fff&size=128",  # Bot avatar
                    None  # User avatar (default)
                ]
            )
            
            with gr.Column(elem_classes="chat-input-container"):
                with gr.Row(elem_classes="file-upload"):
                    file_upload = gr.File(
                        label="Upload Excel File",
                        file_types=[".xlsx"],
                        elem_classes="upload-btn",
                        scale=4
                    )
                    analyze_btn = gr.Button(
                        "Analyze",
                        elem_classes="analyze-btn",
                        scale=1
                    )
                
                report_output = gr.File(
                    label="Download Report",
                    visible=False,
                    interactive=False,
                    elem_classes="report-download"
                )

        chatbot_state = gr.State(value=[])

        def update_ui(file, current_state):
            messages, report_path = process_final_report(agent, file, current_state)
            return messages, gr.update(visible=report_path is not None, value=report_path), messages

        analyze_btn.click(
            fn=update_ui,
            inputs=[file_upload, chatbot_state],
            outputs=[chatbot, report_output, chatbot_state]
        )

    return demo

if __name__ == "__main__":
    try:
        agent = init_agent()
        demo = create_ui(agent)
        demo.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=False)
    except Exception as e:
        print(f"Error: {str(e)}")
        sys.exit(1)