import sys
import os
import pandas as pd
import gradio as gr
from typing import List, Tuple, Dict, Any, Union
import shutil
import re
from datetime import datetime
import time
import asyncio
import logging
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

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 = 131072  # TxAgent's max token limit
MAX_CHUNK_TOKENS = 32768   # Larger chunks to reduce number of chunks
MAX_NEW_TOKENS = 512       # Optimized for fast generation
PROMPT_OVERHEAD = 500      # Estimated tokens for prompt template
MAX_CONCURRENT = 4         # Reduced concurrency to avoid vLLM issues

# Setup logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

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  # Conservative estimate

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:
        logger.error(f"Error extracting Excel: {str(e)}")
        raise ValueError(f"Failed to process Excel file: {str(e)}")
    return "\n".join(all_text)

def split_text_into_chunks(text: str) -> List[str]:
    """Split text into chunks respecting MAX_CHUNK_TOKENS and PROMPT_OVERHEAD"""
    effective_max = MAX_CHUNK_TOKENS - PROMPT_OVERHEAD
    if effective_max <= 0:
        raise ValueError("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:
            if current_chunk:
                chunks.append("\n".join(current_chunk))
            current_chunk = [line]
            current_tokens = line_tokens
        else:
            current_chunk.append(line)
            current_tokens += line_tokens

    if current_chunk:
        chunks.append("\n".join(current_chunk))
    
    logger.info(f"Split text into {len(chunks)} chunks")
    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 concise responses (max {MAX_NEW_TOKENS} tokens):
- Diagnostic Patterns
- Medication Issues
- Missed Opportunities
- Inconsistencies
- Follow-up Recommendations
"""

def init_agent():
    """Initialize TxAgent with conservative settings to avoid vLLM issues"""
    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_chunk_sync(agent, chunk: str, chunk_idx: int) -> Tuple[int, str]:
    """Synchronous wrapper for chunk processing"""
    try:
        prompt = build_prompt_from_text(chunk)
        prompt_tokens = estimate_tokens(prompt)
        
        if prompt_tokens > MAX_MODEL_TOKENS:
            logger.warning(f"Chunk {chunk_idx} prompt too long ({prompt_tokens} tokens)")
            return chunk_idx, ""
        
        response = ""
        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
        
        return chunk_idx, clean_response(response)
    except Exception as e:
        logger.error(f"Error processing chunk {chunk_idx}: {str(e)}")
        return chunk_idx, ""

async def process_file(agent: TxAgent, file_path: str) -> Generator[Tuple[List[Dict[str, str]], Union[str, None]], None, None]:
    """Process the file with improved error handling and vLLM stability"""
    messages = []
    report_path = None
    
    try:
        # Initial messages
        messages.append({"role": "user", "content": f"Processing file: {os.path.basename(file_path)}"})
        messages.append({"role": "assistant", "content": "⏳ Extracting data from Excel..."})
        yield messages, None
        
        # Extract and chunk text
        start_time = time.time()
        text = extract_text_from_excel(file_path)
        chunks = split_text_into_chunks(text)
        messages.append({"role": "assistant", "content": f"✅ Extracted {len(chunks)} chunks in {time.time()-start_time:.1f}s"})
        yield messages, None
        
        # Process chunks sequentially to avoid vLLM socket issues
        chunk_responses = []
        for idx, chunk in enumerate(chunks):
            messages.append({"role": "assistant", "content": f"🔍 Processing chunk {idx+1}/{len(chunks)}..."})
            yield messages, None
            
            _, response = process_chunk_sync(agent, chunk, idx)
            chunk_responses.append(response)
            
            messages.append({"role": "assistant", "content": f"✅ Chunk {idx+1} processed"})
            yield messages, None
        
        # Combine and summarize
        combined = "\n\n".join([r for r in chunk_responses if r])
        messages.append({"role": "assistant", "content": "📊 Generating final report..."})
        yield messages, None
        
        final_response = ""
        for result in agent.run_gradio_chat(
            message=f"Summarize these clinical findings:\n\n{combined}",
            history=[],
            temperature=0.2,
            max_new_tokens=MAX_NEW_TOKENS*2,
            max_token=MAX_MODEL_TOKENS,
            call_agent=False,
            conversation=[],
        ):
            if isinstance(result, str):
                final_response += result
            elif hasattr(result, "content"):
                final_response += result.content
            elif isinstance(result, list):
                for r in result:
                    if hasattr(r, "content"):
                        final_response += r.content
            
            messages[-1]["content"] = f"📊 Generating final report...\n\n{clean_response(final_response)}"
            yield messages, None
        
        # Save report
        final_report = f"# Final Clinical Report\n\n{clean_response(final_response)}"
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        report_path = os.path.join(report_dir, f"report_{timestamp}.md")
        
        with open(report_path, 'w') as f:
            f.write(final_report)
        
        messages.append({"role": "assistant", "content": f"✅ Report saved: report_{timestamp}.md"})
        yield messages, report_path
    
    except Exception as e:
        logger.error(f"Processing failed: {str(e)}")
        messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
        yield messages, None

def create_ui(agent: TxAgent):
    """Create the Gradio interface with simplified interaction"""
    with gr.Blocks(title="Clinical Analysis", css=".gradio-container {max-width: 900px}") as demo:
        gr.Markdown("## 🏥 Clinical Data Analysis (TxAgent)")
        
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label="Analysis Progress",
                    show_copy_button=True,
                    height=600,
                    type="messages"
                )
            with gr.Column(scale=1):
                file_input = gr.File(
                    label="Upload Excel File",
                    file_types=[".xlsx"],
                    height=100
                )
                analyze_btn = gr.Button(
                    "🧠 Analyze Data",
                    variant="primary"
                )
                report_output = gr.File(
                    label="Download Report",
                    visible=False
                )
        
        analyze_btn.click(
            fn=lambda file: process_file(agent, file.name) if file else ([{"role": "assistant", "content": "❌ Please upload a file"}], None),
            inputs=[file_input],
            outputs=[chatbot, report_output],
            concurrency_limit=1  # Ensure sequential processing
        )
    
    return demo

if __name__ == "__main__":
    try:
        # Initialize with conservative settings
        agent = init_agent()
        demo = create_ui(agent)
        
        # Launch with stability optimizations
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            show_error=True,
            allowed_paths=[report_dir],
            share=False,
            max_threads=4  # Reduced thread count for stability
        )
    except Exception as e:
        logger.error(f"Application failed: {str(e)}")
        sys.exit(1)