"""
Structured data extraction tools
"""

import json
from typing import Dict, Any, List
from bs4 import BeautifulSoup
from mcp.types import TextContent

from ..config import session, logger
from ..schemas import EXTRACTION_SCHEMAS


async def extract_structured_data_tool(args: Dict[str, Any]) -> List[TextContent]:
    """Extract structured data using predefined schemas."""
    url = args.get("url")
    html = args.get("html")
    schema = args.get("schema", "custom")
    custom_schema = args.get("custom_schema", {})
    output_format = args.get("output_format", "json")
    
    try:
        # Get HTML content using efficient API with browser session cookies
        if url:
            from ..utils.api_client import get_api_client
            api_client = get_api_client()

            # Update API client with latest browser session cookies
            api_client.ensure_session_updated()

            result_data = api_client.get(url, timeout=10, headers=headers, debug=debug)

            if not result_data.get("success", False):
                error_msg = result_data.get("error", f"HTTP {result_data.get('status_code', 0)}: Failed to fetch {url}")
                return [TextContent(type="text", text=error_msg)]

            # Extract HTML from API response
            html = result_data.get("content", "")
            if result_data.get("json_data"):
                # If response is JSON, try to extract HTML from it
                json_data = result_data["json_data"]
                if isinstance(json_data, dict):
                    html = json_data.get("html", json_data.get("content", html))
        elif not html:
            return [TextContent(type="text", text="Either 'url' or 'html' parameter is required")]
        
        soup = BeautifulSoup(html, 'lxml')
        
        # Choose extraction schema
        if schema == "custom":
            extraction_schema = custom_schema
        elif schema in EXTRACTION_SCHEMAS:
            extraction_schema = EXTRACTION_SCHEMAS[schema]
        else:
            return [TextContent(type="text", text=f"Unknown schema: {schema}. Available: {list(EXTRACTION_SCHEMAS.keys())}")]
        
        if not extraction_schema:
            return [TextContent(type="text", text="No extraction schema provided")]
        
        # Extract data using schema
        extracted_data = {}
        
        for field_name, selectors in extraction_schema.items():
            if not isinstance(selectors, list):
                selectors = [selectors]
            
            field_data = []
            
            for selector in selectors:
                try:
                    elements = soup.select(selector)
                    for element in elements:
                        if selector.endswith(('img', 'image')):
                            # Handle image sources
                            value = element.get('src') or element.get('data-src')
                        elif element.get('href'):
                            # Handle links
                            value = element.get('href')
                        elif element.get('content'):
                            # Handle meta tags
                            value = element.get('content')
                        else:
                            # Handle text content
                            value = element.get_text(strip=True)
                        
                        if value and value not in field_data:
                            field_data.append(value)
                    
                    if field_data:
                        break  # Found data with this selector, move to next field
                        
                except Exception as e:
                    logger.warning(f"Error with selector {selector}: {e}")
                    continue
            
            # Store extracted data
            if len(field_data) == 1:
                extracted_data[field_name] = field_data[0]
            elif len(field_data) > 1:
                extracted_data[field_name] = field_data
            else:
                extracted_data[field_name] = None
        
        # Format output
        if output_format == "csv":
            import csv
            import io
            
            output = io.StringIO()
            if extracted_data:
                writer = csv.DictWriter(output, fieldnames=extracted_data.keys())
                writer.writeheader()
                
                # Handle case where some fields are arrays
                max_rows = max(
                    len(v) if isinstance(v, list) else 1 
                    for v in extracted_data.values()
                )
                
                for i in range(max_rows):
                    row = {}
                    for key, value in extracted_data.items():
                        if isinstance(value, list):
                            row[key] = value[i] if i < len(value) else ""
                        else:
                            row[key] = value if i == 0 else ""
                    writer.writerow(row)
            
            formatted_data = output.getvalue()
            
        elif output_format == "table":
            # Simple table format
            if extracted_data:
                table_lines = []
                table_lines.append("| Field | Value |")
                table_lines.append("|-------|-------|")
                
                for key, value in extracted_data.items():
                    if isinstance(value, list):
                        value_str = ", ".join(str(v) for v in value[:3])
                        if len(value) > 3:
                            value_str += f" ... (+{len(value)-3} more)"
                    else:
                        value_str = str(value) if value else ""
                    
                    table_lines.append(f"| {key} | {value_str} |")
                
                formatted_data = "\n".join(table_lines)
            else:
                formatted_data = "No data extracted"
        
        else:  # json format
            formatted_data = json.dumps(extracted_data, indent=2, ensure_ascii=False)
        
        result = {
            "url": url or "provided_html",
            "schema": schema,
            "output_format": output_format,
            "fields_extracted": len(extracted_data),
            "extraction_schema": extraction_schema,
            "structured_data": extracted_data,
            "formatted_output": formatted_data
        }
        
        return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))]
        
    except Exception as e:
        return [TextContent(type="text", text=f"Structured data extraction failed: {str(e)}")]