import sys
import os
import json
import shutil
import re
import gc
import time
from datetime import datetime
from typing import List, Tuple, Dict, Union
import pandas as pd
import pdfplumber
import gradio as gr
import torch
import matplotlib.pyplot as plt
from fpdf import FPDF
import unicodedata

# === Configuration ===
persistent_dir = "/data/hf_cache"
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

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

MAX_MODEL_TOKENS = 131072
MAX_NEW_TOKENS = 4096
MAX_CHUNK_TOKENS = 8192
BATCH_SIZE = 2
PROMPT_OVERHEAD = 300
SAFE_SLEEP = 0.5

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

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)
    return text.strip()

def remove_duplicate_paragraphs(text: str) -> str:
    paragraphs = text.strip().split("\n\n")
    seen = set()
    unique_paragraphs = []
    for p in paragraphs:
        clean_p = p.strip()
        if clean_p and clean_p not in seen:
            unique_paragraphs.append(clean_p)
            seen.add(clean_p)
    return "\n\n".join(unique_paragraphs)

def extract_text_from_excel(path: str) -> str:
    all_text = []
    xls = pd.ExcelFile(path)
    for sheet_name in xls.sheet_names:
        try:
            df = xls.parse(sheet_name).astype(str).fillna("")
        except Exception:
            continue
        for _, row in df.iterrows():
            non_empty = [cell.strip() for cell in row if cell.strip()]
            if len(non_empty) >= 2:
                text_line = " | ".join(non_empty)
                if len(text_line) > 15:
                    all_text.append(f"[{sheet_name}] {text_line}")
    return "\n".join(all_text)

def extract_text_from_csv(path: str) -> str:
    all_text = []
    try:
        df = pd.read_csv(path).astype(str).fillna("")
    except Exception:
        return ""
    for _, row in df.iterrows():
        non_empty = [cell.strip() for cell in row if cell.strip()]
        if len(non_empty) >= 2:
            text_line = " | ".join(non_empty)
            if len(text_line) > 15:
                all_text.append(text_line)
    return "\n".join(all_text)

def extract_text_from_pdf(path: str) -> str:
    import logging
    logging.getLogger("pdfminer").setLevel(logging.ERROR)
    all_text = []
    try:
        with pdfplumber.open(path) as pdf:
            for page in pdf.pages:
                text = page.extract_text()
                if text:
                    all_text.append(text.strip())
    except Exception:
        return ""
    return "\n".join(all_text)

def extract_text(file_path: str) -> str:
    if file_path.endswith(".xlsx"):
        return extract_text_from_excel(file_path)
    elif file_path.endswith(".csv"):
        return extract_text_from_csv(file_path)
    elif file_path.endswith(".pdf"):
        return extract_text_from_pdf(file_path)
    else:
        return ""

def split_text(text: str, max_tokens=MAX_CHUNK_TOKENS) -> List[str]:
    effective_limit = max_tokens - PROMPT_OVERHEAD
    chunks, current, current_tokens = [], [], 0
    for line in text.split("\n"):
        tokens = estimate_tokens(line)
        if current_tokens + tokens > effective_limit:
            if current:
                chunks.append("\n".join(current))
            current, current_tokens = [line], tokens
        else:
            current.append(line)
            current_tokens += tokens
    if current:
        chunks.append("\n".join(current))
    return chunks

def batch_chunks(chunks: List[str], batch_size: int = BATCH_SIZE) -> List[List[str]]:
    return [chunks[i:i+batch_size] for i in range(0, len(chunks), batch_size)]

def build_prompt(chunk: str) -> str:
    return f"""### Unstructured Clinical Records\n\nAnalyze the clinical notes below and summarize with:\n- Diagnostic Patterns\n- Medication Issues\n- Missed Opportunities\n- Inconsistencies\n- Follow-up Recommendations\n\n---\n\n{chunk}\n\n---\nRespond concisely in bullet points with clinical reasoning."""

def init_agent() -> TxAgent:
    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 analyze_batches(agent, batches: List[List[str]]) -> List[str]:
    results = []
    for batch in batches:
        prompt = "\n\n".join(build_prompt(chunk) for chunk in batch)
        try:
            batch_response = ""
            for r in agent.run_gradio_chat(
                message=prompt,
                history=[],
                temperature=0.0,
                max_new_tokens=MAX_NEW_TOKENS,
                max_token=MAX_MODEL_TOKENS,
                call_agent=False,
                conversation=[]
            ):
                if isinstance(r, str):
                    batch_response += r
                elif isinstance(r, list):
                    for m in r:
                        if hasattr(m, "content"):
                            batch_response += m.content
                elif hasattr(r, "content"):
                    batch_response += r.content
            results.append(clean_response(batch_response))
            time.sleep(SAFE_SLEEP)
        except Exception as e:
            results.append(f"❌ Batch failed: {str(e)}")
            time.sleep(SAFE_SLEEP * 2)
    torch.cuda.empty_cache()
    gc.collect()
    return results

def generate_final_summary(agent, combined: str) -> str:
    combined = remove_duplicate_paragraphs(combined)
    final_prompt = f"""
You are an expert clinical summarizer. Analyze the following summaries carefully and generate a **single final concise structured medical report**, avoiding any repetition or redundancy.

Summaries:
{combined}

Respond with:
- Diagnostic Patterns
- Medication Issues
- Missed Opportunities
- Inconsistencies
- Follow-up Recommendations
Avoid repeating the same points multiple times.
""".strip()

    final_response = ""
    for r in agent.run_gradio_chat(
        message=final_prompt,
        history=[],
        temperature=0.0,
        max_new_tokens=MAX_NEW_TOKENS,
        max_token=MAX_MODEL_TOKENS,
        call_agent=False,
        conversation=[]
    ):
        if isinstance(r, str):
            final_response += r
        elif isinstance(r, list):
            for m in r:
                if hasattr(m, "content"):
                    final_response += m.content
        elif hasattr(r, "content"):
            final_response += r.content

    final_response = clean_response(final_response)
    final_response = remove_duplicate_paragraphs(final_response)
    return final_response

def remove_non_ascii(text):
    return unicodedata.normalize('NFKD', text).encode('ascii', 'ignore').decode('ascii')

def generate_pdf_report_with_charts(summary: str, report_path: str):
    chart_dir = os.path.join(os.path.dirname(report_path), "charts")
    os.makedirs(chart_dir, exist_ok=True)

    chart_path = os.path.join(chart_dir, "summary_chart.png")
    categories = ['Diagnostics', 'Medications', 'Missed', 'Inconsistencies', 'Follow-up']
    values = [4, 2, 3, 1, 5]
    plt.figure(figsize=(6, 4))
    plt.bar(categories, values)
    plt.title('Clinical Issues Overview')
    plt.tight_layout()
    plt.savefig(chart_path)
    plt.close()

    pdf_path = report_path.replace('.md', '.pdf')
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    pdf.multi_cell(0, 10, txt="Final Medical Report", align="C")
    pdf.ln(5)

    for line in summary.split("\n"):
        clean_line = remove_non_ascii(line)
        pdf.multi_cell(0, 10, txt=clean_line)

    pdf.ln(10)
    pdf.image(chart_path, w=150)
    pdf.output(pdf_path)
    return pdf_path

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

    start_time = time.time()
    messages.append({"role": "user", "content": f"📂 Processing file: {os.path.basename(file.name)}"})

    try:
        extracted = extract_text(file.name)
        if not extracted:
            messages.append({"role": "assistant", "content": "❌ Could not extract text."})
            return messages, None

        chunks = split_text(extracted)
        batches = batch_chunks(chunks, batch_size=BATCH_SIZE)
        messages.append({"role": "assistant", "content": f"🔍 Split into {len(batches)} batches. Analyzing..."})

        batch_results = analyze_batches(agent, batches)
        valid = [res for res in batch_results if not res.startswith("❌")]

        if not valid:
            messages.append({"role": "assistant", "content": "❌ No valid batch outputs."})
            return messages, None

        summary = generate_final_summary(agent, "\n\n".join(valid))

        report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md")
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write(f"# Final Medical Report\n\n{summary}")

        pdf_path = generate_pdf_report_with_charts(summary, report_path)

        end_time = time.time()
        elapsed_time = end_time - start_time

        messages.append({"role": "assistant", "content": f"📊 **Final Report:**\n\n{summary}"})
        messages.append({"role": "assistant", "content": f"✅ Report generated in **{elapsed_time:.2f} seconds**.\n\n📥 PDF report ready: {os.path.basename(pdf_path)}"})

        return messages, pdf_path

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

def create_ui(agent):
    with gr.Blocks(css="""
        html, body, .gradio-container { background: #0e1621; color: #e0e0e0; padding: 16px; }
        button.svelte-1ipelgc { background: linear-gradient(to right, #1e88e5, #0d47a1) !important; border: 1px solid #0d47a1 !important; color: white !important; font-weight: bold !important; padding: 10px 20px !important; border-radius: 8px !important; }
        button.svelte-1ipelgc:hover { background: linear-gradient(to right, #2196f3, #1565c0) !important; border: 1px solid #1565c0 !important; }
        .gr-column { align-items: center !important; gap: 12px; }
        .gr-file, .gr-button { width: 100% !important; max-width: 400px; }
    """) as demo:
        gr.Markdown("""
        <h2 style='text-align:center;'>📄 CPS: Clinical Patient Support System</h2>
        <p style='text-align:center;'>Analyze and summarize unstructured medical files using AI (optimized for A100 GPU).</p>
        """)

        with gr.Column():
            chatbot = gr.Chatbot(label="🧠 CPS Assistant", height=480, type="messages")
            upload = gr.File(label="📂 Upload Medical File", file_types=[".xlsx", ".csv", ".pdf"])
            analyze = gr.Button("🧠 Analyze")
            download = gr.File(label="📥 Download Report", visible=False, interactive=False)

        state = gr.State(value=[])

        def handle_analysis(file, chat):
            messages, report_path = process_report(agent, file, chat)
            return messages, gr.update(visible=bool(report_path), value=report_path), messages

        analyze.click(fn=handle_analysis, inputs=[upload, state], outputs=[chatbot, download, state])

    return demo

if __name__ == "__main__":
    agent = init_agent()
    ui = create_ui(agent)
    ui.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=False)