"""
TaxForge - Document Extractor using Google Gemini API (Free Tier)
"""
import os
import base64
import json
import asyncio
import httpx
from pathlib import Path
from typing import Optional

# Rate limiting settings
MAX_RETRIES = 3
RETRY_DELAY_SECONDS = 5
REQUEST_DELAY_SECONDS = 2  # Delay between requests to avoid rate limiting


# Prompt to extract tax document data
EXTRACTION_PROMPT = """You are a tax document parser. Analyze this tax document image and extract all relevant data.

For W-2 forms, extract:
- employer_name: Company name
- employer_ein: Employer's EIN (XX-XXXXXXX)
- wages: Box 1 - Wages, tips, other compensation
- federal_withheld: Box 2 - Federal income tax withheld
- social_security_wages: Box 3
- social_security_tax: Box 4
- medicare_wages: Box 5
- medicare_tax: Box 6
- state: State abbreviation
- state_wages: State wages
- state_tax_withheld: State tax withheld

For 1099-INT forms, extract:
- payer_name: Financial institution name
- interest_income: Box 1 - Interest income
- federal_withheld: Box 4 - Federal income tax withheld

For 1099-DIV forms, extract:
- payer_name: Institution name
- ordinary_dividends: Box 1a - Total ordinary dividends
- qualified_dividends: Box 1b - Qualified dividends
- capital_gain: Box 2a - Total capital gain distributions
- federal_withheld: Box 4 - Federal income tax withheld

For 1099-B forms, extract:
- payer_name: Broker name
- proceeds: Total proceeds
- cost_basis: Total cost basis
- gain_loss: Net gain or loss

For 1098 forms, extract:
- lender_name: Mortgage company name
- mortgage_interest: Box 1 - Mortgage interest received
- points_paid: Box 6 - Points paid
- property_taxes: Property taxes (if shown)

Respond ONLY with a valid JSON object containing:
{
  "document_type": "W-2" | "1099-INT" | "1099-DIV" | "1099-B" | "1098" | "OTHER",
  "confidence": 0.0-1.0,
  "extracted_data": {
    // relevant fields based on document type
  }
}

Important:
- All monetary values should be numbers (no $ or commas)
- If a field is not visible or unclear, omit it
- Set confidence lower if document is blurry or partially visible
"""


async def extract_from_document(
    file_path: str,
    api_key: Optional[str] = None,
) -> dict:
    """
    Extract tax data from a document using Google Gemini API (Free Tier).
    
    Args:
        file_path: Path to PDF or image file
        api_key: Gemini API key (uses GOOGLE_API_KEY env var if not provided)
        
    Returns:
        Dict with document_type, confidence, and extracted_data
    """
    api_key = api_key or os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
    if not api_key:
        return {
            "error": "No API key provided. Get a free key from aistudio.google.com/apikey",
            "document_type": "UNKNOWN",
            "confidence": 0.0,
            "extracted_data": {}
        }
    
    # Read and encode file
    path = Path(file_path)
    if not path.exists():
        return {
            "error": f"File not found: {file_path}",
            "document_type": "UNKNOWN",
            "confidence": 0.0,
            "extracted_data": {}
        }
    
    file_bytes = path.read_bytes()
    base64_data = base64.standard_b64encode(file_bytes).decode("utf-8")
    
    # Determine media type
    suffix = path.suffix.lower()
    media_type_map = {
        ".pdf": "application/pdf",
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".gif": "image/gif",
        ".webp": "image/webp",
    }
    media_type = media_type_map.get(suffix, "image/png")
    
    # Check file size (Gemini inline limit ~20MB)
    file_size_mb = len(file_bytes) / (1024 * 1024)
    if file_size_mb > 20:
        return {
            "error": f"File too large ({file_size_mb:.1f}MB). Please split into smaller files (<20MB each).",
            "document_type": "UNKNOWN",
            "confidence": 0.0,
            "extracted_data": {}
        }
    
    # Call Gemini API
    # Use gemini-2.0-flash for PDF/document support
    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            response = await client.post(
                f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}",
                headers={
                    "content-type": "application/json",
                },
                json={
                    "contents": [
                        {
                            "parts": [
                                {
                                    "inline_data": {
                                        "mime_type": media_type,
                                        "data": base64_data,
                                    }
                                },
                                {
                                    "text": EXTRACTION_PROMPT,
                                }
                            ]
                        }
                    ],
                    "generationConfig": {
                        "temperature": 0.1,
                        "maxOutputTokens": 4096,
                    }
                },
            )
            
            if response.status_code != 200:
                error_text = response.text
                if "API_KEY_INVALID" in error_text:
                    return {
                        "error": "Invalid API key. Get a free key from aistudio.google.com/apikey",
                        "document_type": "UNKNOWN",
                        "confidence": 0.0,
                        "extracted_data": {}
                    }
                if "INVALID_ARGUMENT" in error_text and "pdf" in error_text.lower():
                    return {
                        "error": "PDF processing failed. Try converting to images (screenshot each page) or use a smaller file.",
                        "document_type": "UNKNOWN",
                        "confidence": 0.0,
                        "extracted_data": {}
                    }
                if "RESOURCE_EXHAUSTED" in error_text or response.status_code == 429:
                    # Rate limited - will be retried by caller
                    return {
                        "error": "RATE_LIMITED",
                        "document_type": "UNKNOWN",
                        "confidence": 0.0,
                        "extracted_data": {}
                    }
                return {
                    "error": f"API error: {response.status_code} - {error_text[:200]}",
                    "document_type": "UNKNOWN",
                    "confidence": 0.0,
                    "extracted_data": {}
                }
            
            result = response.json()
            
            # Extract text from Gemini response
            candidates = result.get("candidates", [])
            if not candidates:
                return {
                    "error": "No response from Gemini",
                    "document_type": "UNKNOWN",
                    "confidence": 0.0,
                    "extracted_data": {}
                }
            
            content = candidates[0].get("content", {}).get("parts", [{}])[0].get("text", "{}")
            
            # Parse JSON from response
            # Handle case where model wraps in markdown
            if "```json" in content:
                content = content.split("```json")[1].split("```")[0]
            elif "```" in content:
                content = content.split("```")[1].split("```")[0]
            
            try:
                parsed = json.loads(content.strip())
                return parsed
            except json.JSONDecodeError:
                return {
                    "error": f"Failed to parse response as JSON: {content[:200]}",
                    "document_type": "UNKNOWN",
                    "confidence": 0.0,
                    "extracted_data": {}
                }
                
    except Exception as e:
        return {
            "error": f"Request failed: {str(e)}",
            "document_type": "UNKNOWN",
            "confidence": 0.0,
            "extracted_data": {}
        }


async def extract_with_retry(
    file_path: str,
    api_key: Optional[str] = None,
) -> dict:
    """
    Extract with automatic retry on rate limiting.
    """
    for attempt in range(MAX_RETRIES):
        result = await extract_from_document(file_path, api_key)
        
        if result.get("error") == "RATE_LIMITED":
            if attempt < MAX_RETRIES - 1:
                # Wait and retry with exponential backoff
                wait_time = RETRY_DELAY_SECONDS * (2 ** attempt)
                await asyncio.sleep(wait_time)
                continue
            else:
                return {
                    "error": "API quota exceeded after retries. Please wait a minute and try again.",
                    "document_type": "UNKNOWN",
                    "confidence": 0.0,
                    "extracted_data": {}
                }
        
        return result
    
    return result


def convert_to_w2(extracted_data: dict) -> dict:
    """Convert extracted data to W-2 format for state."""
    return {
        "employer_name": extracted_data.get("employer_name", "Unknown Employer"),
        "wages": float(extracted_data.get("wages", 0)),
        "federal_withheld": float(extracted_data.get("federal_withheld", 0)),
    }


def convert_to_1099(extracted_data: dict, form_type: str) -> dict:
    """Convert extracted data to 1099 format for state."""
    amount = 0.0
    if form_type == "1099-INT":
        amount = float(extracted_data.get("interest_income", 0))
    elif form_type == "1099-DIV":
        amount = float(extracted_data.get("ordinary_dividends", 0))
    elif form_type == "1099-B":
        amount = float(extracted_data.get("gain_loss", extracted_data.get("proceeds", 0)))
    
    return {
        "payer_name": extracted_data.get("payer_name", extracted_data.get("lender_name", "Unknown Payer")),
        "amount": amount,
        "form_type": form_type,
    }
