import os
import logging
import traceback
from typing import Optional, Dict, Any, Tuple
from huggingface_hub import InferenceClient
from utils.meldrx import MeldRxAPI
from utils.pdfutils import PDFGenerator
from utils.responseparser import PatientDataExtractor
from datetime import datetime

# Set up logging with detailed output
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN environment variable not set.")
client = InferenceClient(api_key=HF_TOKEN)
MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"

def generate_ai_discharge_summary(patient_dict: Dict[str, str]) -> Optional[str]:
    """Generate a discharge summary using AI based on extracted patient data."""
    try:
        patient_info = (
            f"Patient Name: {patient_dict['first_name']} {patient_dict['last_name']}\n"
            f"Gender: {patient_dict['sex']}\n"
            f"Age: {patient_dict['age']}\n"
            f"Date of Birth: {patient_dict['dob']}\n"
            f"Admission Date: {patient_dict['admission_date']}\n"
            f"Discharge Date: {patient_dict['discharge_date']}\n\n"
            f"Diagnosis:\n{patient_dict['diagnosis']}\n\n"
            f"Medications:\n{patient_dict['medications']}\n\n"
            f"Discharge Instructions:\n[Generated based on available data]"
        )

        logger.info("Generating AI discharge summary with patient info: %s", patient_info)

        messages = [
            {
                "role": "assistant",
                "content": (
                    "You are a senior medical practitioner tasked with creating discharge summaries. "
                    "Generate a complete discharge summary based on the provided patient information."
                )
            },
            {"role": "user", "content": patient_info}
        ]

        stream = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            temperature=0.4,
            max_tokens=3584,
            top_p=0.7,
            stream=True
        )

        discharge_summary = ""
        for chunk in stream:
            content = chunk.choices[0].delta.content
            if content:
                discharge_summary += content

        logger.info("AI discharge summary generated successfully")
        return discharge_summary.strip()

    except Exception as e:
        logger.error("Error generating AI discharge summary: %s\n%s", str(e), traceback.format_exc())
        return None

def generate_discharge_paper_one_click(
    meldrx_api: MeldRxAPI,
    patient_id: str = None,
    first_name: str = None,
    last_name: str = None
) -> Tuple[Optional[str], str, Optional[str]]:
    """Generate a discharge paper with AI content in one click."""
    try:
        if not meldrx_api.access_token:
            if not meldrx_api.authenticate():
                return None, "Error: Authentication failed. Please authenticate first.", None

        logger.info("Fetching patient data from MeldRx API")
        patient_data = meldrx_api.get_patients()
        if not patient_data:
            return None, "Error: No patient data returned from MeldRx API.", None
        if "entry" not in patient_data:
            logger.error("Invalid patient data format: %s", patient_data)
            return None, "Error: Patient data is not in expected FHIR Bundle format.", None

        logger.info("Extracting patient data")
        extractor = PatientDataExtractor(patient_data, format_type="json")
        patients = extractor.get_all_patients()

        if not patients:
            return None, "Error: No patients found in the workspace.", None

        patient_dict = None
        if patient_id:
            for p in patients:
                extractor.set_patient_by_index(patients.index(p))
                if extractor.get_id() == patient_id:
                    patient_dict = p
                    break
            if not patient_dict:
                return None, f"Error: Patient with ID {patient_id} not found.", None
        elif first_name and last_name:
            patient_dict = next(
                (p for p in patients if
                 p["first_name"].lower() == first_name.lower() and
                 p["last_name"].lower() == last_name.lower()),
                None
            )
            if not patient_dict:
                return None, f"Error: Patient with name {first_name} {last_name} not found.", None
        else:
            patient_dict = patients[0]

        logger.info("Selected patient: %s %s", patient_dict['first_name'], patient_dict['last_name'])

        ai_content = generate_ai_discharge_summary(patient_dict)
        if not ai_content:
            return None, "Error: Failed to generate AI discharge summary.", None

        display_summary = (
            f"<div style='color:#00FFFF; font-family: monospace;'>"
            f"<strong>Discharge Summary Preview</strong><br>"
            f"- Name: {patient_dict['first_name']} {patient_dict['last_name']}<br>"
            f"- DOB: {patient_dict['dob']}, Age: {patient_dict['age']}, Sex: {patient_dict['sex']}<br>"
            f"- Address: {patient_dict['address']}, {patient_dict['city']}, {patient_dict['state']} {patient_dict['zip_code']}<br>"
            f"- Admission Date: {patient_dict['admission_date']}<br>"
            f"- Discharge Date: {patient_dict['discharge_date']}<br>"
            f"- Diagnosis: {patient_dict['diagnosis']}<br>"
            f"- Medications: {patient_dict['medications']}<br>"
            f"</div>"
        )

        pdf_generator = PDFGenerator()
        pdf_path = pdf_generator.generate_pdf_from_text(
            ai_content,
            f"discharge_summary_{patient_id or 'unknown'}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
        )

        if pdf_path:
            return pdf_path, f"Success: Discharge paper generated for {patient_dict['first_name']} {patient_dict['last_name']}", display_summary
        return None, "Error: Failed to generate PDF.", display_summary

    except Exception as e:
        logger.error("Error in one-click discharge generation: %s\n%s", str(e), traceback.format_exc())
        return None, f"Error: {str(e)}", None