import sys
import os
import pandas as pd
import json
import gradio as gr
from typing import List, Tuple, Dict, Any
import hashlib
import shutil
import re
from datetime import datetime
import time
from collections import defaultdict

# 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 - Updated for 32,768 token limit
MAX_TOKENS = 32768
CHUNK_SIZE = 10000  # Target chunk size (allowing 3 chunks within limit)
MAX_NEW_TOKENS = 2048  # Increased output length
MAX_BOOKINGS_PER_CHUNK = 5  # Process 5 bookings per chunk

def file_hash(path: str) -> str:
    """Generate MD5 hash of file contents"""
    with open(path, "rb") as f:
        return hashlib.md5(f.read()).hexdigest()

def clean_response(text: str) -> str:
    """Clean and normalize text output"""
    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:
    """More accurate token estimation (1 token ~ 3-4 characters)"""
    return len(text) // 3.5  # More conservative estimate

def process_patient_data(df: pd.DataFrame) -> Dict[str, Any]:
    """Enhanced patient data processing with chronology"""
    data = {
        'bookings': defaultdict(list),
        'medications': defaultdict(list),
        'diagnoses': defaultdict(list),
        'tests': defaultdict(list),
        'procedures': defaultdict(list),
        'doctors': set(),
        'timeline': []
    }
    
    # Sort by date and group by booking
    df = df.sort_values('Interview Date')
    for booking, group in df.groupby('Booking Number'):
        for _, row in group.iterrows():
            entry = {
                'booking': booking,
                'date': str(row['Interview Date']),
                'doctor': str(row['Interviewer']),
                'form': str(row['Form Name']),
                'item': str(row['Form Item']),
                'response': str(row['Item Response']),
                'notes': str(row['Description'])
            }
            
            data['bookings'][booking].append(entry)
            data['timeline'].append(entry)
            data['doctors'].add(entry['doctor'])
            
            # Enhanced categorization
            form_lower = entry['form'].lower()
            if 'medication' in form_lower or 'drug' in form_lower:
                data['medications'][entry['item']].append(entry)
            elif 'diagnosis' in form_lower or 'condition' in form_lower:
                data['diagnoses'][entry['item']].append(entry)
            elif 'test' in form_lower or 'lab' in form_lower or 'result' in form_lower:
                data['tests'][entry['item']].append(entry)
            elif 'procedure' in form_lower or 'surgery' in form_lower:
                data['procedures'][entry['item']].append(entry)
    
    return data

def generate_analysis_prompt(patient_data: Dict[str, Any], bookings: List[str]) -> str:
    """Generate comprehensive prompt for multiple bookings"""
    prompt_lines = [
        "**Comprehensive Patient Analysis**",
        f"Analyzing {len(bookings)} bookings spanning {patient_data['timeline'][0]['date']} to {patient_data['timeline'][-1]['date']}",
        "Focus on identifying patterns, inconsistencies, and missed opportunities across the entire treatment history.",
        "",
        "**Key Analysis Points:**",
        "- Chronological progression of symptoms and diagnoses",
        "- Medication changes and potential interactions over time",
        "- Diagnostic consistency across different providers",
        "- Missed diagnostic opportunities based on symptoms and test results",
        "- Gaps in follow-up or incomplete assessments",
        "- Emerging patterns that may indicate chronic conditions",
        "",
        "**Patient Timeline (Condensed):**"
    ]
    
    # Add condensed timeline
    for entry in patient_data['timeline']:
        if entry['booking'] in bookings:
            prompt_lines.append(
                f"- {entry['date']}: {entry['form']} - {entry['item']} = {entry['response']} (by {entry['doctor']})"
            )
    
    # Add current medications
    prompt_lines.extend([
        "",
        "**Medication History:**",
        *[f"- {med}: " + " → ".join(
            f"{e['date']}: {e['response']}" 
            for e in entries if e['booking'] in bookings
        ) for med, entries in patient_data['medications'].items()],
        "",
        "**Diagnostic History:**",
        *[f"- {diag}: " + " → ".join(
            f"{e['date']}: {e['response']}" 
            for e in entries if e['booking'] in bookings
        ) for diag, entries in patient_data['diagnoses'].items()],
        "",
        "**Required Analysis Format:**",
        "### Diagnostic Patterns",
        "[Identify patterns in symptoms and diagnoses over time]",
        "",
        "### Medication Analysis",
        "[Review all medication changes and potential issues]",
        "",
        "### Provider Consistency",
        "[Note any discrepancies between different doctors]",
        "",
        "### Missed Opportunities",
        "[Potential diagnoses or interventions that were missed]",
        "",
        "### Comprehensive Recommendations",
        "[Actionable recommendations for current care]"
    ])
    
    return "\n".join(prompt_lines)

def chunk_bookings(patient_data: Dict[str, Any]) -> List[List[str]]:
    """Split bookings into 3 balanced chunks based on token count"""
    all_bookings = list(patient_data['bookings'].keys())
    
    # Estimate token count for each booking
    booking_sizes = []
    for booking in all_bookings:
        entries = patient_data['bookings'][booking]
        size = sum(estimate_tokens(str(e)) for e in entries)
        booking_sizes.append((booking, size))
    
    # Sort by size (descending) for better chunk balancing
    booking_sizes.sort(key=lambda x: x[1], reverse=True)
    
    # Initialize 3 chunks
    chunks = [[] for _ in range(3)]
    chunk_sizes = [0, 0, 0]
    
    # Distribute bookings to chunks
    for booking, size in booking_sizes:
        # Find the chunk with smallest current size
        min_chunk = chunk_sizes.index(min(chunk_sizes))
        chunks[min_chunk].append(booking)
        chunk_sizes[min_chunk] += size
    
    return chunks

def init_agent():
    """Initialize TxAgent with enhanced configuration"""
    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=[],
        device_map="auto"
    )
    agent.init_model()
    return agent

def analyze_with_agent(agent, prompt: str) -> str:
    """Enhanced analysis with progress tracking"""
    try:
        response = ""
        for result in agent.run_gradio_chat(
            message=prompt,
            history=[],
            temperature=0.2,
            max_new_tokens=MAX_NEW_TOKENS,
            max_token=MAX_TOKENS,
            call_agent=False,
            conversation=[],
        ):
            if isinstance(result, list):
                for r in result:
                    if hasattr(r, 'content') and r.content:
                        response += clean_response(r.content) + "\n"
            elif isinstance(result, str):
                response += clean_response(result) + "\n"
            elif hasattr(result, 'content'):
                response += clean_response(result.content) + "\n"
        
        return response.strip()
    except Exception as e:
        return f"Error in analysis: {str(e)}"

def create_ui(agent):
    with gr.Blocks(theme=gr.themes.Soft(), title="Patient History Analyzer") as demo:
        gr.Markdown("# 🏥 Comprehensive Patient History Analyzer")
        
        with gr.Tabs():
            with gr.TabItem("Analysis"):
                with gr.Row():
                    with gr.Column(scale=1):
                        file_upload = gr.File(
                            label="Upload Patient Excel File",
                            file_types=[".xlsx"],
                            file_count="single"
                        )
                        analysis_btn = gr.Button("Analyze Full History", variant="primary")
                        status = gr.Markdown("Ready for analysis")
                        progress = gr.Slider(
                            minimum=0, 
                            maximum=100, 
                            value=0,
                            label="Analysis Progress",
                            interactive=False
                        )
                    
                    with gr.Column(scale=2):
                        output_display = gr.Markdown(
                            label="Analysis Results",
                            elem_id="results"
                        )
                        report_download = gr.File(
                            label="Download Full Report",
                            interactive=False
                        )
            
            with gr.TabItem("Instructions"):
                gr.Markdown("""
                ## Enhanced Patient History Analysis
                
                This tool processes complete medical histories across multiple visits, identifying:
                - Patterns in symptoms and diagnoses over time
                - Medication safety issues across providers
                - Missed diagnostic opportunities
                - Gaps in follow-up care
                
                **How to Use:**
                1. Upload Excel file with patient history
                2. Click "Analyze Full History"
                3. View progressive results
                4. Download comprehensive report
                
                **File Requirements:**
                - Must contain complete visit history
                - Required columns: Booking Number, Interview Date, Interviewer, 
                  Form Name, Form Item, Item Response, Description
                """)
        
        def analyze_patient(file) -> Tuple[str, str, int]:
            if not file:
                raise gr.Error("Please upload an Excel file first")
            
            full_report = []
            report_path = ""
            
            try:
                # Process Excel file
                df = pd.read_excel(file.name)
                patient_data = process_patient_data(df)
                
                # Split into 3 balanced chunks
                booking_chunks = chunk_bookings(patient_data)
                total_chunks = len(booking_chunks)
                
                for chunk_idx, bookings in enumerate(booking_chunks, 1):
                    # Update progress
                    progress_value = int((chunk_idx/total_chunks)*100)
                    yield "", "", progress_value
                    
                    # Generate and process prompt
                    prompt = generate_analysis_prompt(patient_data, bookings)
                    response = analyze_with_agent(agent, prompt)
                    
                    if "Error in analysis" not in response:
                        full_report.append(
                            f"## Analysis Segment {chunk_idx} (Bookings: {', '.join(bookings)})\n{response}\n"
                        )
                        yield "\n".join(full_report), "", progress_value
                    
                    time.sleep(0.1)  # Prevent UI freezing
                
                # Generate final summary
                if len(booking_chunks) > 1:
                    summary_prompt = f"""
**Final Comprehensive Summary**

Analyze all {len(patient_data['bookings'])} bookings to identify:
1. Overall health trajectory
2. Chronic condition patterns
3. Medication safety across entire treatment
4. Most critical missed opportunities
5. Priority recommendations

**Required Format:**
### Health Trajectory
[Overall progression of health status]

### Chronic Condition Analysis
[Patterns indicating chronic issues]

### Critical Concerns
[Most urgent issues needing attention]

### Priority Recommendations
[Action items ranked by importance]
"""
                    summary = analyze_with_agent(agent, summary_prompt)
                    full_report.append(f"## Final Comprehensive Summary\n{summary}\n")
                
                # Save report
                report_filename = f"patient_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
                report_path = os.path.join(report_dir, report_filename)
                with open(report_path, 'w', encoding='utf-8') as f:
                    f.write("\n".join(full_report))
                
                yield "\n".join(full_report), report_path, 100
                
            except Exception as e:
                raise gr.Error(f"Analysis failed: {str(e)}")
        
        analysis_btn.click(
            analyze_patient,
            inputs=file_upload,
            outputs=[output_display, report_download, progress],
            api_name="analyze"
        )
    
    return demo

if __name__ == "__main__":
    try:
        agent = init_agent()
        demo = create_ui(agent)
        
        demo.queue(
            api_open=False,
            max_size=20
        ).launch(
            server_name="0.0.0.0",
            server_port=7860,
            show_error=True,
            allowed_paths=[report_dir],
            share=False
        )
    except Exception as e:
        print(f"Failed to launch application: {str(e)}")
        sys.exit(1)