import json
import logging
from typing import Any, Dict, Optional, List
from fastapi import APIRouter, HTTPException, Depends, Request
from fastapi.responses import PlainTextResponse
from pydantic import BaseModel
from app.services.ai_fix_service import generate_fix
from app.services.llm_service import LLMConfig
from app.utils.patch_db_utils import save_patch_to_db
from app.database import SessionDep
from app.models import PatchType
from app.auth import get_current_user
from app.models import User
from app.api.scan import SimpleRateLimiter
from app.utils.audit_logger import log_audit_event

_ai_fix_limiter = SimpleRateLimiter(max_calls=10, period_seconds=60)

logger = logging.getLogger(__name__)

router = APIRouter()


class AIFixRequest(BaseModel):
    file_path: str
    file_content: str
    finding: Dict[str, Any]
    llm_config: Optional[Dict[str, Any]] = None  # LLM configuration from frontend
    mode: str = "comment"  # "comment" = @TODO-pq annotation; "codegen" = actual code replacement


class AIFixResponse(BaseModel):
    success: bool
    generated_fix: Optional[str] = None
    explanation: str
    confidence: Optional[float] = None
    warnings: Optional[List[str]] = None
    validation: Optional[Dict[str, Any]] = None
    patch_id: Optional[str] = None


@router.post("/ai-fix/generate", response_model=AIFixResponse)
async def ai_fix_generate(
    body: AIFixRequest,
    request: Request,
    session: SessionDep,
    user: User = Depends(get_current_user)
) -> AIFixResponse:
    """Generate an AI-proposed fix for a scanner finding.

    Simplified approach: LLM automatically detects if vulnerability is quantum-related
    and recommends appropriate fixes (PQC or standard security).
    """
    if not _ai_fix_limiter.check(key=user.id):
        raise HTTPException(status_code=429, detail="Rate limit exceeded. Please wait before generating another fix.")

    try:
        # Parse LLM config from request body
        llm_config = None
        if body.llm_config:
            llm_config = LLMConfig(
                provider=body.llm_config.get("provider", "google"),
                api_key=body.llm_config.get("apiKey") or None,
                base_url=body.llm_config.get("baseUrl"),
                model=body.llm_config.get("model"),
            )

        logger.info(f"Generating fix for {body.finding.get('check_id')} using {llm_config.provider if llm_config else 'no'} provider")

        # Generate fix using the requested mode
        result = generate_fix(
            file_path=body.file_path,
            file_content=body.file_content,
            finding=body.finding,
            llm_config=llm_config,
            mode=body.mode,
        )

        # If fix was successfully generated, save to database
        if result.get("success") and result.get("generated_fix"):
            validation = result.get("validation", {})

            # Generate fingerprint — must match the frontend key in FindingRow:
            #   result?.fingerprint || `${result?.path}-${start_line}`
            finding = body.finding
            start_line = (
                finding.get("start_line")
                or (finding.get("start") or {}).get("line")
                or 0
            )
            # Normalize path separators so fingerprint matches frontend key
            normalized_path = body.file_path.replace("\\", "/")
            fingerprint = (
                finding.get("fingerprint")
                or f"{normalized_path}-{start_line}"
            )

            # Get stats from validation (if available)
            stats = (validation.get("stats") or {}) if validation else {}

            # All patches are AI-generated (comment-based guidance)
            patch_type = PatchType.AI

            # Save patch to database
            patch_id = save_patch_to_db(
                session=session,
                file_path=body.file_path,
                finding_fingerprint=fingerprint,
                patch_diff=result.get("generated_fix"),
                patch_type=patch_type,
                source_data={
                    "check_id": finding.get("check_id"),
                    "severity": finding.get("severity"),
                    "message": finding.get("message"),
                    "start_line": finding.get("start_line"),
                    "end_line": finding.get("end_line"),
                    "validation": validation,
                },
                user_id=user.id,
                lines_added=stats.get("added", 0),
                lines_removed=stats.get("removed", 0),
            )

            if patch_id:
                logger.info(f"Saved AI fix patch {patch_id} for {body.file_path} (valid={validation.get('valid', False)})")
                result["patch_id"] = patch_id

        log_audit_event("ai_fix_generated", user_id=user.id, details={
            "file_path": body.file_path,
            "check_id": body.finding.get("check_id"),
            "patch_id": result.get("patch_id"),
        })
        return AIFixResponse(**result)
    except Exception as e:
        err_str = str(e)
        if "503" in err_str or "UNAVAILABLE" in err_str or "high demand" in err_str.lower() or "ServiceUnavailable" in type(e).__name__:
            raise HTTPException(status_code=503, detail="The AI model is temporarily unavailable due to high demand. Please try again in a moment.")
        logger.exception("AI fix generation failed")
        raise HTTPException(status_code=500, detail=f"AI fix generation failed: {str(e)}")


class AITestRequest(BaseModel):
    llm_config: Dict[str, Any]


class AISummaryRequest(BaseModel):
    findings: List[Dict[str, Any]]
    llm_config: Optional[Dict[str, Any]] = None


class AISummaryResponse(BaseModel):
    success: bool
    summary: Optional[str] = None
    error: Optional[str] = None


class AIReportRequest(BaseModel):
    scan_ids: List[str]
    user_context: Optional[str] = None
    llm_config: Optional[Dict[str, Any]] = None


class AIReportResponse(BaseModel):
    success: bool
    report: Optional[str] = None
    report_id: Optional[str] = None
    error: Optional[str] = None


@router.post("/ai-fix/summary", response_model=AISummaryResponse)
async def ai_summary_generate(
    body: AISummaryRequest,
    user: User = Depends(get_current_user)
) -> AISummaryResponse:
    """Generate AI summary of scan findings."""
    try:
        from app.services.llm_service import generate_completion, LLMConfig
        
        # Parse LLM config
        llm_config = None
        if body.llm_config:
            llm_config = LLMConfig(
                provider=body.llm_config.get("provider", "google"),
                api_key=body.llm_config.get("apiKey") or None,
                base_url=body.llm_config.get("baseUrl"),
                model=body.llm_config.get("model"),
            )
        
        if not llm_config:
            return AISummaryResponse(success=False, error="No LLM configuration provided")
        
        # Prepare findings summary for AI
        findings_summary = []
        for finding in body.findings[:20]:  # Limit to first 20 findings
            findings_summary.append({
                "rule": finding.get("check_id", "unknown"),
                "severity": finding.get("severity", "unknown"),
                "message": finding.get("message", "")[:100],  # Truncate message
                "file": finding.get("path", "unknown"),
                "quantum_vulnerable": finding.get("extra", {}).get("metadata", {}).get("quantum_vulnerable", False)
            })
        
        # Create prompt for AI summary
        prompt = f"""Analyze these security scan findings and provide a concise 1-2 sentence summary focusing on the most critical issues:

Findings ({len(findings_summary)} total):
{chr(10).join([f"- {f['severity'].upper()}: {f['rule']} in {f['file']} - {f['message']}" for f in findings_summary[:10]])}

Provide a plain-English summary highlighting:
1. Most critical security issues found
2. Quantum cryptography vulnerabilities if any
3. Overall risk assessment

Keep response under 150 words and focus on actionable insights."""
        
        summary = generate_completion(prompt, llm_config)
        
        return AISummaryResponse(success=True, summary=summary.strip())
        
    except Exception as e:
        logger.exception("AI summary generation failed")
        return AISummaryResponse(success=False, error=str(e))


@router.post("/ai-fix/test")
async def ai_fix_test(
    body: AITestRequest,
    user: User = Depends(get_current_user),
):
    """Validate LLM configuration by making a minimal test call."""
    from app.services.llm_service import generate_completion, LLMConfig
    try:
        llm_config = LLMConfig(
            provider=body.llm_config.get("provider", "google"),
            api_key=body.llm_config.get("apiKey") or None,
            base_url=body.llm_config.get("baseUrl"),
            model=body.llm_config.get("model"),
        )
        response = generate_completion("Respond with only the word OK.", llm_config)
        return {"success": True, "message": f"Connected — model responded: {response.strip()[:80]}"}
    except Exception as e:
        error = str(e)
        logger.warning(f"AI connection test failed: {error}")
        return {"success": False, "message": error}


@router.post("/ai-fix/generate-report", response_model=AIReportResponse)
async def ai_generate_report(
    body: AIReportRequest,
    session: SessionDep,
    user: User = Depends(get_current_user)
) -> AIReportResponse:
    """Generate comprehensive AI report from multiple scan results and save to DB."""
    try:
        from app.services.llm_service import generate_completion, LLMConfig
        from app.models import GeneratedReport
        import re

        llm_config = None
        if body.llm_config:
            llm_config = LLMConfig(
                provider=body.llm_config.get("provider", "google"),
                api_key=body.llm_config.get("apiKey") or None,
                base_url=body.llm_config.get("baseUrl"),
                model=body.llm_config.get("model"),
            )

        if not llm_config:
            return AIReportResponse(success=False, error="No LLM configuration provided")

        from sqlmodel import select
        from app.models import ScanRecord

        scan_summaries = []
        for scan_id in body.scan_ids:
            try:
                scan = session.exec(select(ScanRecord).where(ScanRecord.id == scan_id)).first()
                if not scan or not scan.scan_data:
                    continue
                scan_data = scan.scan_data
                findings = scan_data.get("enriched_findings") or scan_data.get("data", {}).get("results") or []
                severity_counts: dict = {"critical": 0, "high": 0, "medium": 0, "low": 0}
                for f in findings:
                    raw = (f.get("extra", {}).get("severity") or f.get("severity") or "info").upper()
                    if raw in ("ERROR", "CRITICAL"): severity_counts["critical"] += 1
                    elif raw in ("WARNING", "HIGH"): severity_counts["high"] += 1
                    elif raw in ("INFO", "MEDIUM"): severity_counts["medium"] += 1
                    else: severity_counts["low"] += 1
                yara_matches = scan_data.get("data", {}).get("yara_analysis", {}).get("matches_found", 0)
                quantum_risk = scan_data.get("quantum_risk_score") or scan_data.get("risk_score")
                findings_count = len(findings) or yara_matches
                risk_level = (
                    "CRITICAL" if severity_counts["critical"] > 0 else
                    "HIGH" if severity_counts["high"] > 0 else
                    "MEDIUM" if severity_counts["medium"] > 0 else
                    "LOW" if findings_count > 0 else "CLEAN"
                )
                rules = list({f.get("check_id") or f.get("rule_id") or "unknown" for f in findings[:20]})
                scan_summaries.append({
                    "type": scan.scan_type,
                    "target": scan.target,
                    "timestamp": str(scan.timestamp),
                    "findings_count": findings_count,
                    "severity_counts": severity_counts,
                    "top_rules": rules[:10],
                    "risk_level": risk_level,
                    "quantum_risk_score": quantum_risk,
                })
            except Exception as e:
                logger.warning(f"Failed to fetch scan {scan_id}: {e}")

        if not scan_summaries:
            return AIReportResponse(success=False, error="No valid scan data found")

        context_section = f"\n\nUser Context: {body.user_context}" if body.user_context else ""
        findings_summary = "\n".join([
            f"- {s['type'].upper()} | Target: {s['target']} | "
            f"{s['findings_count']} findings | Risk: {s['risk_level']} | "
            f"Critical: {s['severity_counts']['critical']}, High: {s['severity_counts']['high']}, "
            f"Medium: {s['severity_counts']['medium']}, Low: {s['severity_counts']['low']}"
            + (f" | Quantum Risk Score: {s['quantum_risk_score']}" if s.get('quantum_risk_score') else "")
            + (f"\n  Top rules: {', '.join(s['top_rules'][:5])}" if s.get('top_rules') else "")
            for s in scan_summaries
        ])

        from datetime import datetime, timezone
        report_date = datetime.now(timezone.utc).strftime("%B %d, %Y")
        targets = ", ".join(dict.fromkeys(s["target"] for s in scan_summaries))

        prompt = f"""Generate a comprehensive post-quantum security assessment report in Markdown format.

Scan Results Summary:
{findings_summary}{context_section}

Use these exact values in the report header metadata (do not invent or change them):
- Date: {report_date}
- Assessment Target: {targets}

Create a professional security report with the following structure:

# Post-Quantum Security Assessment Report

## Executive Summary
[2-3 sentences highlighting the most critical quantum cryptography vulnerabilities and overall risk level]

## Risk Overview
[Overall quantum risk assessment with key risk factors]

## Key Findings by Category

### Critical Quantum Vulnerabilities
[Most severe quantum-vulnerable cryptography found]

### Code Security Issues
[Application-level security findings]

### Infrastructure Vulnerabilities
[Network, certificate, and system-level issues]

### Binary Analysis Results
[Findings from executable analysis]

## Post-Quantum Cryptography Recommendations

### Immediate Actions (0-3 months)
[Urgent steps to mitigate quantum risks]

### Short-term Migration (3-12 months)
[Planned transition to quantum-safe algorithms]

### Long-term Strategy (1-3 years)
[Complete post-quantum cryptography implementation]

## Technical Implementation Guide
[Specific technical steps for remediation]

## Compliance and Business Impact
[Regulatory implications and business risk assessment]

Keep the report professional, actionable, and under 1500 words. Focus on quantum cryptography risks and practical migration steps."""

        report_content = generate_completion(prompt, llm_config)
        report_content = report_content.strip()

        # Extract title from first # heading, fallback to timestamp
        title_match = re.search(r'^#\s+(.+)$', report_content, re.MULTILINE)
        title = title_match.group(1).strip() if title_match else f"Security Report {scan_summaries[0]['timestamp'][:10]}"

        # Save to DB
        record = GeneratedReport(
            user_id=user.id,
            title=title,
            user_context=body.user_context,
            content=report_content,
            scan_ids=json.dumps(body.scan_ids),
        )
        session.add(record)
        session.commit()
        session.refresh(record)

        return AIReportResponse(success=True, report=report_content, report_id=record.id)

    except Exception as e:
        logger.exception("AI report generation failed")
        return AIReportResponse(success=False, error=str(e))


@router.get("/ai-fix/reports")
async def list_generated_reports(
    session: SessionDep,
    user: User = Depends(get_current_user),
):
    """List all saved AI-generated reports."""
    from sqlmodel import select
    from app.models import GeneratedReport
    reports = session.exec(
        select(GeneratedReport)
        .where(GeneratedReport.user_id == user.id)
        .order_by(GeneratedReport.created_at.desc())  # type: ignore
    ).all()
    return {"success": True, "data": [
        {
            "id": r.id,
            "title": r.title,
            "user_context": r.user_context,
            "created_at": r.created_at.isoformat(),
            "scan_count": len(json.loads(r.scan_ids)),
        }
        for r in reports
    ]}


@router.delete("/ai-fix/reports/{report_id}")
async def delete_generated_report(
    report_id: str,
    session: SessionDep,
    user: User = Depends(get_current_user),
):
    """Delete a saved AI-generated report."""
    from sqlmodel import select
    from app.models import GeneratedReport
    record = session.exec(select(GeneratedReport).where(
        GeneratedReport.id == report_id,
        GeneratedReport.user_id == user.id,
    )).first()
    if not record:
        raise HTTPException(status_code=404, detail="Report not found")
    session.delete(record)
    session.commit()
    return {"success": True}


@router.get("/ai-fix/reports/{report_id}/export")
async def export_generated_report(
    report_id: str,
    format: str = "md",
    session: SessionDep = None,
    user: User = Depends(get_current_user),
):
    """Download a saved report as markdown or HTML."""
    from sqlmodel import select
    from app.models import GeneratedReport
    from fastapi.responses import HTMLResponse
    record = session.exec(select(GeneratedReport).where(
        GeneratedReport.id == report_id,
        GeneratedReport.user_id == user.id,
    )).first()
    if not record:
        raise HTTPException(status_code=404, detail="Report not found")
    if format == "html":
        html = f"""<!DOCTYPE html><html><head><title>{record.title}</title>
<style>body{{font-family:system-ui;line-height:1.6;margin:40px;max-width:900px}}pre{{background:#f3f4f6;padding:16px;border-radius:8px;white-space:pre-wrap}}</style>
</head><body><pre>{record.content}</pre></body></html>"""
        return HTMLResponse(content=html, headers={"Content-Disposition": f'attachment; filename="report.html"'})
    return PlainTextResponse(content=record.content, headers={"Content-Disposition": f'attachment; filename="report.md"'})
