import logging
import requests
import json
import re
import time
from datetime import datetime
from typing import Type, TypeVar, get_origin, get_args, Any, Union, Literal
from pydantic import BaseModel
from waveassist.constants import API_BASE_URL

logger = logging.getLogger("waveassist")

T = TypeVar('T', bound=BaseModel)


def call_post_api(path, body) -> tuple:
    url = f"{API_BASE_URL}/{path}"
    headers = { "Content-Type": "application/json" }  # JSON content
    try:
        response = requests.post(url, json=body, headers=headers)  # Sends proper JSON
        response_dict = response.json()

        if str(response_dict.get("success")) == "1":
            return True, response_dict
        else:
            error_message = response_dict.get("message", "Unknown error")
            return False, error_message
    except Exception as e:
        logger.error("API call failed: %s", e)
        return False, str(e)

def call_post_api_with_files(path, body, files=None) -> tuple:
    url = f"{API_BASE_URL}/{path}"
    try:
        response = requests.post(url, data=body, files=files or {})
        response_dict = response.json()
        if str(response_dict.get("success")) == "1":
            return True, response_dict
        else:
            error_message = response_dict.get("message", "Unknown error")
            return False, error_message
    except Exception as e:
        logger.error("API call failed: %s", e)
        return False, str(e)



def call_get_api(path, params) -> tuple:
    url = f"{API_BASE_URL}/{path}"
    headers = { "Content-Type": "application/json" }
    try:
        response = requests.get(url, params=params, headers=headers)
        response_dict = response.json()

        if str(response_dict.get("success")) == "1":
            return True, response_dict.get("data", {})
        else:
            error_message = response_dict.get("message", "Unknown error")
            return False, error_message

    except Exception as e:
        logger.error("API GET call failed: %s", e)
        return False, str(e)




def _get_type_name(annotation: Any) -> str:
    """Get a clean type name from a type annotation."""
    if annotation is None:
        return "null"
    
    origin = get_origin(annotation)
    
    # Handle Optional types (Union[X, None])
    if origin is Union:
        args = get_args(annotation)
        non_none_args = [a for a in args if a is not type(None)]
        if len(non_none_args) == 1:
            return _get_type_name(non_none_args[0])
        return " | ".join(_get_type_name(a) for a in non_none_args)
    
    # Handle Literal (show allowed values)
    if origin is Literal:
        args = get_args(annotation)
        return " | ".join(repr(a) for a in args)

    # Handle List, Dict, etc.
    if origin is list:
        args = get_args(annotation)
        if args:
            return f"list[{_get_type_name(args[0])}]"
        return "list"
    
    if origin is dict:
        args = get_args(annotation)
        if args and len(args) == 2:
            return f"dict[{_get_type_name(args[0])}, {_get_type_name(args[1])}]"
        return "dict"
    
    # Handle Pydantic BaseModel (nested models)
    if isinstance(annotation, type) and issubclass(annotation, BaseModel):
        return annotation.__name__
    
    # Handle basic types
    if hasattr(annotation, '__name__'):
        return annotation.__name__
    
    return str(annotation)


def _generate_template_value(field_annotation: Any, field_description: str | None) -> Any:
    """Generate a template value for a field, handling nested models."""
    # Check if it's a Pydantic model (nested)
    origin = get_origin(field_annotation)
    
    # Handle Optional types
    if origin is Union:
        args = get_args(field_annotation)
        non_none_args = [a for a in args if a is not type(None)]
        if non_none_args:
            return _generate_template_value(non_none_args[0], field_description)
    
    # Handle List of Pydantic models
    if origin is list:
        args = get_args(field_annotation)
        if args and isinstance(args[0], type) and issubclass(args[0], BaseModel):
            return [generate_json_template_dict(args[0])]
        elif args:
            return [f"<{_get_type_name(args[0])}>"]
        return ["<item>"]
    
    # Handle nested Pydantic model
    if isinstance(field_annotation, type) and issubclass(field_annotation, BaseModel):
        return generate_json_template_dict(field_annotation)
    
    # Always include type; append description when present
    type_str = _get_type_name(field_annotation)
    if field_description:
        return f"<{type_str}> {field_description}"
    return f"<{type_str}>"


def generate_json_template_dict(model: Type[BaseModel]) -> dict:
    """Generate a template dictionary showing the structure and descriptions."""
    template = {}
    for name, field in model.model_fields.items():
        template[name] = _generate_template_value(field.annotation, field.description)
    return template


def generate_json_template(model: Type[BaseModel]) -> str:
    """Generate a clean JSON string showing the structure and descriptions."""
    return json.dumps(generate_json_template_dict(model), indent=2)


def _find_balanced_json(content: str, start_char: str, end_char: str) -> str | None:
    """
    Find a balanced JSON object or array starting from the first occurrence of start_char.
    
    Args:
        content: Content to search in
        start_char: Opening character ('{' or '[')
        end_char: Closing character ('}' or ']')
        
    Returns:
        The matched JSON string, or None if not found
    """
    start_idx = content.find(start_char)
    if start_idx == -1:
        return None
    
    depth = 0
    in_string = False
    escape_next = False
    
    for i in range(start_idx, len(content)):
        char = content[i]
        
        if escape_next:
            escape_next = False
            continue
        
        if char == '\\':
            escape_next = True
            continue
        
        if char == '"' and not escape_next:
            in_string = not in_string
            continue
        
        if not in_string:
            if char == start_char:
                depth += 1
            elif char == end_char:
                depth -= 1
                if depth == 0:
                    # Found matching closing brace/bracket
                    return content[start_idx:i + 1]
    
    return None



def extract_json_from_content(content: str) -> Union[dict, list]:
    """
    Extract and parse JSON from content using multiple strategies.
    Includes json_repair as a fallback for malformed JSON.
    
    Args:
        content: Raw content that may contain JSON
        
    Returns:
        Parsed JSON as a dictionary or list
        
    Raises:
        ValueError: If no valid JSON could be extracted
    """
    if not content:
        raise ValueError("Empty content received")
    
    content = content.strip()
    
    # Strategy 1: Try parsing directly (content is pure JSON)
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        pass
    
    # Strategy 2: Extract from ```json code blocks
    json_block_match = re.search(r'```json\s*([\s\S]*?)\s*```', content)
    if json_block_match:
        try:
            return json.loads(json_block_match.group(1).strip())
        except json.JSONDecodeError:
            pass
    
    # Strategy 3: Extract from generic ``` code blocks
    code_block_match = re.search(r'```\s*([\s\S]*?)\s*```', content)
    if code_block_match:
        try:
            return json.loads(code_block_match.group(1).strip())
        except json.JSONDecodeError:
            pass
    
    # Strategy 4: Find JSON object pattern { ... } with balanced matching
    json_object_str = _find_balanced_json(content, '{', '}')
    if json_object_str:
        try:
            return json.loads(json_object_str)
        except json.JSONDecodeError:
            pass
    
    # Strategy 5: Find JSON array pattern [ ... ] with balanced matching
    json_array_str = _find_balanced_json(content, '[', ']')
    if json_array_str:
        try:
            return json.loads(json_array_str)
        except json.JSONDecodeError:
            pass
    
    # Strategy 6: Try json_repair for malformed JSON (fallback)
    try:
        import json_repair
        repaired = json_repair.repair_json(content)
        return json.loads(repaired)
    except (ImportError, Exception):
        # json_repair not available or failed, continue to error
        pass
    
    raise ValueError(f"Could not extract valid JSON from content: {content[:200]}...")


def soft_parse(model_class: Type[T], raw_data: dict) -> T:
    """
    Parse data into a Pydantic model leniently.
    
    - Ignores extra fields not in the model
    - Allows missing optional fields (they become None)
    - Attempts type coercion where possible
    - Handles Pydantic field aliases automatically
    - Falls back to model_construct if validation fails (e.g., missing required fields)
    - Fills in missing fields with None or their default values
    
    This ensures LLM responses are never wasted - we return whatever data was gathered.
    
    Args:
        model_class: Pydantic model class to parse into
        raw_data: Raw dictionary data from LLM
        
    Returns:
        Instance of the model class (may have None for missing required fields)
    """
    # 1. Try validation without filtering first (handles aliases automatically)
    # Pydantic will use aliases if configured with populate_by_name=True
    try:
        return model_class.model_validate(raw_data)
    except Exception:
        pass
    
    # 2. Filter to only valid keys (including aliases)
    valid_keys = set(model_class.model_fields.keys())
    # Also check for alias values
    alias_to_field = {}
    for field_name, field_info in model_class.model_fields.items():
        if hasattr(field_info, 'alias') and field_info.alias:
            alias_to_field[field_info.alias] = field_name
    
    filtered_data = {}
    for k, v in raw_data.items():
        # Include if it's a direct field name or an alias
        if k in valid_keys or k in alias_to_field:
            # Map alias back to field name if needed
            actual_key = alias_to_field.get(k, k)
            filtered_data[actual_key] = v
    
    # 3. Try normal validation with filtered data (handles type coercion, defaults, etc.)
    try:
        return model_class.model_validate(filtered_data)
    except Exception:
        pass
    
    # 4. Safety Fallback: "Just give me what you found"
    # model_construct bypasses validation - returns whatever data is available
    obj = model_class.model_construct(**filtered_data)
    
    # 5. Fill in the gaps: Ensure every field exists
    from pydantic_core import PydanticUndefined
    for field_name in valid_keys:
        if not hasattr(obj, field_name):
            # Set to the field's default if it has one, otherwise None
            field_info = model_class.model_fields[field_name]
            default_value = field_info.default
            
            # Check for default_factory if default is undefined
            if default_value is PydanticUndefined:
                # Check if there's a default_factory to call
                default_factory = getattr(field_info, 'default_factory', None)
                if default_factory is not None and callable(default_factory):
                    default_value = default_factory()
                else:
                    default_value = None
            
            setattr(obj, field_name, default_value)
    
    return obj


def parse_json_response(content: str, response_model: Type[T], model: str) -> T:
    """
    Parse JSON content and validate it against a Pydantic model with soft parsing.
    
    Args:
        content: JSON string to parse (may contain markdown code blocks)
        response_model: Pydantic model class to validate against
        model: Model name (for error messages)
        
    Returns:
        Validated instance of response_model
        
    Raises:
        ValueError: If JSON extraction or parsing fails
    """
    try:
        # Extract JSON using multiple strategies
        parsed_data = extract_json_from_content(content)

        # If the extracted JSON is an array, take the first element
        # (since the model expects an object/dict)
        if isinstance(parsed_data, list):
            if len(parsed_data) == 0:
                raise ValueError(
                    f"Expected JSON object but got empty array. "
                    f"The model '{response_model.__name__}' requires an object, not an array."
                )
            parsed_data = parsed_data[0]
            # Ensure it's a dict after extracting from array
            if not isinstance(parsed_data, dict):
                raise ValueError(
                    f"Expected JSON object but got array with non-object element. "
                    f"The model '{response_model.__name__}' requires an object."
                )

        # Soft parse with lenient validation
        return soft_parse(response_model, parsed_data)
    except ValueError:
        raise
    except Exception as e:
        raise ValueError(
            f"Failed to validate response from model '{model}' against {response_model.__name__}. "
            f"Error: {e}"
        )


def create_json_prompt(prompt: str, response_model: Type[BaseModel]) -> str:
    """
    Create a prompt that requests JSON output matching a Pydantic model structure.
    
    Args:
        prompt: Original user prompt
        response_model: Pydantic model class to generate template from
        
    Returns:
        Enhanced prompt with JSON structure instructions
    """
    template = generate_json_template(response_model)
    return f"""{prompt}
    
    Respond with a JSON object following this structure:
    {template}
    
    Return ONLY the JSON object, no explanations or other text. Return JSON now:"""