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
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

# Configuration and setup
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 directory in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir]:
    os.makedirs(directory, exist_ok=True)

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

current_dir = os.path.dirname(os.path.abspath(__file__))
src_path = os.path.abspath(os.path.join(current_dir, "src"))
sys.path.insert(0, src_path)

from txagent.txagent import TxAgent

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

def clean_response(text: str) -> str:
    try:
        text = text.encode('utf-8', 'surrogatepass').decode('utf-8')
    except UnicodeError:
        text = text.encode('utf-8', 'replace').decode('utf-8')
    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 = []
    try:
        xls = pd.ExcelFile(file_path)
        for sheet_name in xls.sheet_names:
            df = xls.parse(sheet_name)
            df = df.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)
    except Exception as e:
        raise ValueError(f"Failed to extract text from Excel file: {str(e)}")
    return "\n".join(all_text)

def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> List[str]:
    effective_max_tokens = max_tokens - PROMPT_OVERHEAD
    if effective_max_tokens <= 0:
        raise ValueError(f"Effective max tokens ({effective_max_tokens}) must be positive.")
    lines = text.split("\n")
    chunks, current_chunk, current_tokens = [], [], 0
    for line in lines:
        line_tokens = estimate_tokens(line)
        if current_tokens + line_tokens > effective_max_tokens:
            if current_chunk:
                chunks.append("\n".join(current_chunk))
            current_chunk, current_tokens = [line], line_tokens
        else:
            current_chunk.append(line)
            current_tokens += line_tokens
    if current_chunk:
        chunks.append("\n".join(current_chunk))
    return chunks

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

You are reviewing unstructured, mixed-format clinical documentation from various forms, tables, and sheets.

**Objective:** Identify patterns, missed diagnoses, inconsistencies, and follow-up gaps.

Here is the extracted content chunk:

{chunk}

Please analyze the above and provide:
- Diagnostic Patterns
- Medication Issues
- Missed Opportunities
- Inconsistencies
- Follow-up Recommendations
"""

def init_agent():
    default_tool_path = os.path.abspath("data/new_tool.json")
    target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
    if not os.path.exists(target_tool_path):
        shutil.copy(default_tool_path, target_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": target_tool_path},
        force_finish=True,
        enable_checker=True,
        step_rag_num=4,
        seed=100,
        additional_default_tools=[]
    )
    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 []
    report_path = None

    if file is None or not hasattr(file, "name"):
        messages.append({"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."})
        return messages, report_path

    try:
        messages.append({"role": "user", "content": f"📄 Processing Excel file: {os.path.basename(file.name)}"})
        messages.append({"role": "assistant", "content": "🔍 Analyzing clinical data... This may take a moment."})
        
        extracted_text = extract_text_from_excel(file.name)
        chunks = split_text_into_chunks(extracted_text)
        chunk_responses = [None] * len(chunks)

        def analyze_chunk(index: int, chunk: str) -> Tuple[int, str]:
            prompt = build_prompt_from_text(chunk)
            prompt_tokens = estimate_tokens(prompt)
            if prompt_tokens > MAX_MODEL_TOKENS:
                return index, f"❌ Chunk {index+1} prompt too long ({prompt_tokens} tokens). Skipping..."
            response = ""
            try:
                for result 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(result, str):
                        response += result
                    elif hasattr(result, "content"):
                        response += result.content
                    elif isinstance(result, list):
                        for r in result:
                            if hasattr(r, "content"):
                                response += r.content
            except Exception as e:
                return index, f"❌ Error analyzing chunk {index+1}: {str(e)}"
            return index, clean_response(response)

        # Process chunks silently without displaying progress
        with ThreadPoolExecutor(max_workers=1) as executor:
            futures = [executor.submit(analyze_chunk, i, chunk) for i, chunk in enumerate(chunks)]
            for future in as_completed(futures):
                i, result = future.result()
                chunk_responses[i] = result

        valid_responses = [res for res in chunk_responses if not res.startswith("❌")]
        if not valid_responses:
            messages.append({"role": "assistant", "content": "❌ No valid analysis results to summarize."})
            return messages, report_path

        summary = "\n\n".join(valid_responses)
        final_prompt = f"""Please synthesize the following clinical analyses into a concise, well-structured report:
        
{summary}

Structure your response with clear sections:
1. Key Diagnostic Patterns
2. Medication Concerns
3. Potential Missed Opportunities
4. Notable Inconsistencies
5. Recommended Follow-ups

Use bullet points for clarity and professional medical terminology."""
        
        final_report_text = ""
        try:
            for result in agent.run_gradio_chat(
                message=final_prompt,
                history=[],
                temperature=0.2,
                max_new_tokens=MAX_NEW_TOKENS,
                max_token=MAX_MODEL_TOKENS,
                call_agent=False,
                conversation=[],
            ):
                if isinstance(result, str):
                    final_report_text += result
                elif hasattr(result, "content"):
                    final_report_text += result.content
                elif isinstance(result, list):
                    for r in result:
                        if hasattr(r, "content"):
                            final_report_text += r.content
        except Exception as e:
            messages.append({"role": "assistant", "content": f"❌ Error generating final report: {str(e)}"})
            return messages, report_path

        final_report = f"# 🧠 Clinical Analysis Report\n\n{clean_response(final_report_text)}"
        
        # Update the last message with the final report
        messages[-1]["content"] = f"## 📋 Clinical Analysis Report\n\n{clean_response(final_report_text)}"

        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        report_path = os.path.join(report_dir, f"clinical_report_{timestamp}.md")

        with open(report_path, 'w') as f:
            f.write(final_report)

        messages.append({"role": "assistant", "content": f"✅ Report generated successfully. You can download it below."})

    except Exception as e:
        messages.append({"role": "assistant", "content": f"❌ Error processing file: {str(e)}"})

    return messages, report_path

def create_ui(agent):
    with gr.Blocks(
        title="Clinical Analysis Tool",
        css="""
        .gradio-container {
            max-width: 900px !important;
            margin: auto;
            font-family: 'Inter', sans-serif;
            background-color: #f9fafb;
        }
        .gr-button.primary {
            background: linear-gradient(to right, #4f46e5, #7c3aed);
            color: white;
            border: none;
            border-radius: 8px;
            padding: 12px 24px;
            font-weight: 500;
            transition: all 0.2s;
        }
        .gr-button.primary:hover {
            background: linear-gradient(to right, #4338ca, #6d28d9);
            transform: translateY(-1px);
            box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        }
        .gr-file-upload, .gr-chatbot, .gr-markdown {
            background-color: white;
            border-radius: 12px;
            box-shadow: 0 1px 3px rgba(0,0,0,0.05);
            padding: 1.5rem;
            border: 1px solid #e5e7eb;
        }
        .gr-chatbot {
            min-height: 600px;
            border-left: none;
        }
        .chat-message-user {
            background-color: #f3f4f6;
            border-radius: 12px;
            padding: 12px 16px;
            margin: 8px 0;
        }
        .chat-message-assistant {
            background-color: white;
            border-radius: 12px;
            padding: 12px 16px;
            margin: 8px 0;
            border: 1px solid #e5e7eb;
        }
        .chat-message-content ul, .chat-message-content ol {
            padding-left: 1.5em;
            margin: 0.5em 0;
        }
        .chat-message-content li {
            margin: 0.3em 0;
        }
        h1, h2, h3 {
            color: #111827;
        }
        .gr-markdown h1 {
            font-size: 1.8rem;
            margin-bottom: 1rem;
            font-weight: 600;
        }
        .gr-markdown p {
            color: #4b5563;
            line-height: 1.6;
        }
        .progress-bar {
            height: 4px;
            background: #e5e7eb;
            border-radius: 2px;
            margin: 12px 0;
            overflow: hidden;
        }
        .progress-bar-fill {
            height: 100%;
            background: linear-gradient(to right, #4f46e5, #7c3aed);
            transition: width 0.3s ease;
        }
        """
    ) as demo:
        gr.Markdown("""
        <div style='text-align: center; margin-bottom: 1.5rem;'>
            <h1 style='margin-bottom: 0.5rem; color: #111827;'>Clinical Documentation Analyzer</h1>
            <p style='color: #6b7280; margin-top: 0;'>Upload patient records in Excel format for comprehensive clinical analysis</p>
        </div>
        """)

        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label="Analysis Results",
                    show_copy_button=True,
                    height=600,
                    bubble_full_width=False,
                    avatar_images=(None, "https://i.imgur.com/6wX7Zb4.png"),
                    render_markdown=True
                )
            with gr.Column(scale=1):
                file_upload = gr.File(
                    label="Upload Patient Records",
                    file_types=[".xlsx", ".xls"],
                    height=100,
                    interactive=True
                )
                analyze_btn = gr.Button(
                    "Analyze Clinical Data",
                    variant="primary",
                    elem_classes="primary"
                )
                report_output = gr.File(
                    label="Download Report",
                    visible=False,
                    interactive=False
                )
                gr.Markdown("""
                <div style='margin-top: 1rem; padding: 1rem; background-color: #f8fafc; border-radius: 8px;'>
                    <h3 style='margin-top: 0; margin-bottom: 0.5rem; font-size: 1rem;'>About this tool</h3>
                    <p style='margin: 0; font-size: 0.9rem; color: #64748b;'>
                        This tool analyzes clinical documentation to identify patterns, inconsistencies, and opportunities for improved patient care.
                    </p>
                </div>
                """)

        chatbot_state = gr.State(value=[])

        def update_ui(file, current_state):
            messages, report_path = process_final_report(agent, file, current_state)
            formatted_messages = []
            for msg in messages:
                role = msg.get("role")
                content = msg.get("content", "")
                if role == "assistant":
                    # Format lists and sections for better readability
                    content = content.replace("- ", "• ")
                    content = re.sub(r"(\d+\.\s)", r"\n\1", content)
                    content = f"<div class='chat-message-assistant'>{content}</div>"
                else:
                    content = f"<div class='chat-message-user'>{content}</div>"
                formatted_messages.append({"role": role, "content": content})
            
            report_update = gr.update(visible=report_path is not None, value=report_path)
            return formatted_messages, report_update, formatted_messages

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

    return demo

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