backend/treatment.py

"""
Treatment module - Generates disease treatment recommendations
Uses Gemini AI or falls back to predefined treatments
"""

class TreatmentGenerator:
    """Generates treatment recommendations for fish diseases"""
    
    def __init__(self, gemini_model=None):
        """
        Initialize treatment generator
        
        Args:
            gemini_model: Google Gemini model instance (optional)
        """
        self.gemini_model = gemini_model
    
    def get_recommendations(self, disease_name, confidence):
        """
        Get treatment recommendations for detected disease
        
        Args:
            disease_name: Name of detected disease
            confidence: Confidence score (0-100)
            
        Returns:
            str: Formatted treatment recommendations
        """
        if self.gemini_model:
            return self._get_ai_treatment(disease_name, confidence)
        else:
            return self._get_fallback_treatment(disease_name)
    
    def _get_ai_treatment(self, disease_name, confidence):
        """Generate AI-powered treatment using Gemini"""
        try:
            prompt = f"""You are an expert aquaculture veterinarian with 20 years of experience.

**DIAGNOSIS:**
- Disease: {disease_name.replace('_', ' ')}
- Confidence: {confidence:.1f}%

Provide a detailed, practical treatment plan in this EXACT format:

**IMMEDIATE ACTIONS (First 24 hours):**
- [List 3-4 urgent steps with specific timing]

**TREATMENT PROTOCOL (Days 1-10):**
- [Medications with exact dosages, frequency, and duration]
- [Water management with specific parameters]

**PREVENTION MEASURES:**
- [List 4-5 specific preventive measures for future]

**EXPECTED OUTCOME:**
- [Recovery timeline and realistic success rate]

Requirements:
- Be specific with dosages (mg/kg, ppm, etc.)
- Include water parameter targets (pH, DO, temperature)
- Use simple language for fish farmers
- Focus on practical, actionable steps"""

            response = self.gemini_model.generate_content(prompt)
            return response.text
            
        except Exception as e:
            print(f"⚠️ Gemini treatment generation error: {e}")
            return self._get_fallback_treatment(disease_name)
    
    def _get_fallback_treatment(self, disease_name):
        """Fallback treatment database (predefined)"""
        treatments = {
            'Aeromoniasis': """**IMMEDIATE ACTIONS (First 24 hours):**
- Isolate infected fish immediately to prevent spread
- Perform 50% water change within 2 hours
- Increase aeration to maximum capacity
- Stop feeding for 24 hours to reduce waste

**TREATMENT PROTOCOL (Days 1-10):**
- Antibiotics: Oxytetracycline 50-75 mg/kg feed for 10 consecutive days
- Salt bath: 3% NaCl solution for 5-10 minutes daily (days 1-7)
- Maintain water temperature at 25-28°C
- Monitor daily for 14 days post-treatment

**PREVENTION MEASURES:**
- Maintain dissolved oxygen > 5 mg/L at all times
- Reduce stocking density by 30% if currently overcrowded
- Implement 14-21 day quarantine for all new fish
- Weekly health monitoring and immediate isolation of sick fish
- Regular disinfection of equipment and nets

**EXPECTED OUTCOME:**
70-80% recovery rate with prompt treatment. Symptoms should improve within 3-5 days. Full recovery typically takes 10-14 days.""",
            
            'Bacterial_gill_disease': """**IMMEDIATE ACTIONS (First 24 hours):**
- Increase aeration immediately (CRITICAL - survival depends on oxygen)
- Reduce feeding by 50% to minimize ammonia buildup
- Test and adjust water quality (pH, ammonia, nitrite)
- Separate severely affected fish to reduce stress

**TREATMENT PROTOCOL (Days 1-10):**
- Salt bath: 3% NaCl for 5-10 minutes daily for 7 days
- Improve water circulation (add aerators/water pumps)
- Reduce stocking density by 40%
- Antibiotics: Florfenicol 10 mg/kg feed daily for 7-10 days (if severe)

**PREVENTION MEASURES:**
- Maintain dissolved oxygen > 5 mg/L continuously (CRITICAL)
- Regular pond/tank cleaning (bi-weekly minimum)
- Daily ammonia monitoring (keep < 0.02 mg/L)
- Avoid sudden temperature changes (±2°C maximum)
- Proper biofilter maintenance and regular backwashing

**EXPECTED OUTCOME:**
85-90% recovery if caught early (within 2-3 days). Symptoms improve within 48 hours of treatment. Mortality high if delayed beyond 5 days.""",
            
            'EUS': """**IMMEDIATE ACTIONS (First 24 hours):**
- REPORT TO AUTHORITIES (notifiable disease in many regions)
- Isolate all affected fish immediately
- Implement 75% water change
- Remove and properly dispose of dead/dying fish

**TREATMENT PROTOCOL (Days 1-21):**
- Topical disinfectant: Potassium permanganate 2-5 ppm bath (3 times weekly)
- NO specific cure available - focus on supportive care
- Maintain optimal water conditions (see below)
- Consider humane culling of severely affected fish (>50% body lesions)
- Secondary infection prevention: Oxytetracycline 50 mg/kg feed

**PREVENTION MEASURES:**
- NEVER use water from infected sources
- Implement strict biosecurity: footbaths, equipment disinfection
- Quarantine ALL new fish for minimum 21 days
- Regular health surveillance (weekly inspections)
- Maintain year-round optimal water quality

**EXPECTED OUTCOME:**
Variable (40-90% mortality depending on strain and timing). NO cure exists. Early detection critical. Focus on prevention and biosecurity.""",
            
            'Healthy_Fish': """**MAINTENANCE (Continue Good Practices):**
- Continue current care routine (it's working excellently!)
- Daily observations for early disease detection
- Maintain current feeding schedule (adjust seasonally)
- Regular water quality monitoring

**PREVENTION MEASURES:**
- Weekly water parameter testing (pH, ammonia, nitrite, nitrate, DO)
- Quarantine new fish for 14-21 days before introduction
- Clean and disinfect all equipment monthly
- Daily behavioral observations (feeding response, swimming patterns)
- Maintain detailed health records

**EXPECTED OUTCOME:**
Excellent prognosis. Continue preventive care to maintain optimal health status. Regular monitoring ensures early detection of any future issues."""
        }
        
        # Default treatment for diseases not in database
        default = """**IMPORTANT: Consult a fish health expert immediately.**

This disease requires professional veterinary assessment for accurate diagnosis and treatment protocol.

**IMMEDIATE STEPS:**
1. Isolate affected fish to prevent potential spread
2. Improve water quality (50% water change, increase aeration)
3. Document all symptoms with photos and notes
4. Contact your local aquaculture veterinarian or extension service

**EMERGENCY CONTACT:**
Seek professional help from:
- Local fish health specialist
- Aquaculture extension service
- Veterinary diagnostic laboratory

Do NOT attempt unguided treatment as it may worsen the condition."""
        
        return treatments.get(disease_name, default)