"""
Infomaniak AI API Client (Legacy Wrapper)
This file now serves as a wrapper around the Universal AI Client.
"""
import logging
import json
from typing import Optional, Dict, Any, List
from sqlalchemy.orm import Session
from fastapi import Depends
from app.database import get_db
from app.models.setting import Setting
from app.integrations.ai_client import UniversalAIClient, AIError

logger = logging.getLogger(__name__)

# Maintain alias for compatibility
InfomaniakAIClient = UniversalAIClient
InfomaniakAIError = AIError

def get_infomaniak_ai_client(db: Session = Depends(get_db)) -> UniversalAIClient:
    """
    Factory function for dependency injection.
    """
    config_setting = db.query(Setting).filter(Setting.key == "ai_config").first()
    if config_setting and config_setting.value:
        try:
            config = json.loads(config_setting.value)
            return UniversalAIClient.from_config(config)
        except Exception:
            logger.error("Invalid AI config in DB")
    
    return UniversalAIClient()

# Core service layer remains similar but uses the universal client
from app.models.ai_analysis import AIAnalysis
from app.models.vulnerability import Vulnerability
from app.models.asset import Asset
from datetime import datetime
from sqlalchemy import desc

class AIAnalysisService:
    """Service for CVE analysis with caching"""

    def __init__(self, db: Session, ai_client: UniversalAIClient):
        self.db = db
        self.ai_client = ai_client

    def get_or_create_analysis(
        self,
        vulnerability_id: int,
        force_refresh: bool = False
    ) -> AIAnalysis:
        if not force_refresh:
            cached = self.db.query(AIAnalysis).filter(
                AIAnalysis.vulnerability_id == vulnerability_id,
                AIAnalysis.cache_expires_at > datetime.now()
            ).first()

            if cached:
                return cached

        vuln = self.db.query(Vulnerability).filter(Vulnerability.id == vulnerability_id).first()
        if not vuln:
            raise ValueError(f"Vulnerability {vulnerability_id} not found")

        analysis_data = self.ai_client.analyze_cve(
            cve_id=vuln.cve_id,
            package_name=vuln.package_name,
            package_version=vuln.package_version,
            cvss_score=vuln.cvss_score
        )

        existing = self.db.query(AIAnalysis).filter(AIAnalysis.vulnerability_id == vulnerability_id).first()
        if existing:
            for key, value in analysis_data.items():
                setattr(existing, key, value)
            self.db.commit()
            self.db.refresh(existing)
            return existing
        else:
            ai_analysis = AIAnalysis(vulnerability_id=vulnerability_id, **analysis_data)
            self.db.add(ai_analysis)
            self.db.commit()
            self.db.refresh(ai_analysis)
            return ai_analysis

    def get_priority_recommendations(self, limit: int = 20, severity: Optional[str] = None) -> Dict[str, Any]:
        # Implementation similar to original but using universal client
        query = self.db.query(Vulnerability).filter(Vulnerability.status == "open")
        if severity:
            query = query.filter(Vulnerability.severity == severity)

        top_vulns = query.order_by(desc(Vulnerability.cvss_score)).limit(limit).all()
        if not top_vulns:
            return {"recommendations": [], "global_strategy": "No open vulnerabilities found."}

        vuln_data = []
        for v in top_vulns:
            vuln_data.append({
                "cve_id": v.cve_id,
                "hostname": v.asset.hostname if v.asset else "Unknown",
                "cvss_score": v.cvss_score,
                "package_name": v.package_name
            })

        # Prompt for structured prioritization
        prompt = f"""You are a Vulnerability Management Expert. Analyze the following list of vulnerabilites and prioritize the top 10 that need immediate attention.

Vulnerabilities:
{json.dumps(vuln_data)}

You MUST respond with a valid JSON object ONLY, following this exact structure:
{{
  "global_strategy": "Brief strategic summary of the risk landscape and focus areas.",
  "recommendations": [
    {{
      "priority": 1,
      "cve_id": "CVE-XXXX-XXXX",
      "asset_hostname": "hostname",
      "justification": "Why this is critical...",
      "remediation_hint": "Brief fix action..."
    }}
  ]
}}
"""
        try:
            response_text = self.ai_client._call_llm(prompt)
            
            # Clean DeepSeek <think> tags if present
            if "<think>" in response_text:
                import re
                response_text = re.sub(r"<think>.*?</think>", "", response_text, flags=re.DOTALL).strip()

            # Robust JSON Extraction
            try:
                # 1. Try removing markdown code blocks
                clean_text = response_text
                if "```json" in clean_text:
                    clean_text = clean_text.split("```json")[1].split("```")[0].strip()
                elif "```" in clean_text:
                    clean_text = clean_text.split("```")[1].split("```")[0].strip()
                
                return json.loads(clean_text)
            except (json.JSONDecodeError, IndexError, TypeError):
                # 2. Aggressive Fallback: Find first '{' and last '}'
                try:
                    start_idx = response_text.find("{")
                    end_idx = response_text.rfind("}")
                    if start_idx != -1 and end_idx != -1 and start_idx < end_idx:
                        json_str = response_text[start_idx : end_idx + 1]
                        return json.loads(json_str)
                except Exception:
                    pass
                
                logger.warning(f"Failed to parse AI response. Raw output: {response_text[:200]}...")
                return {"global_strategy": response_text, "recommendations": []}
        except Exception as e:
            logger.error(f"AI prioritization failed: {e}")
            return {"error": str(e), "recommendations": []}
