import json
import logging
from datetime import datetime
from typing import Dict, List, Any, Optional

logger = logging.getLogger(__name__)


def parse_jsonl_file(file_path: str, start_line: int = 0) -> List[Dict[str, Any]]:
    """Parse JSONL file from a specific line number"""
    messages = []
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for i, line in enumerate(f):
                if i < start_line:
                    continue
                if not line.strip():
                    continue
                try:
                    obj = json.loads(line)
                    messages.append(obj)
                except json.JSONDecodeError as e:
                    logger.warning(f"Failed to parse line {i} in {file_path}: {e}")
                    continue
    except Exception as e:
        logger.error(f"Failed to read {file_path}: {e}")
    return messages


def extract_text_content(content: Any) -> str:
    """Extract text from message content"""
    if isinstance(content, str):
        return content

    if isinstance(content, list):
        text_parts = []
        for block in content:
            if isinstance(block, dict):
                if block.get('type') == 'text':
                    text_parts.append(block.get('text', ''))
                elif block.get('type') == 'tool_result':
                    # Include tool results in searchable text
                    result = block.get('content', '')
                    if isinstance(result, str):
                        text_parts.append(result[:500])  # Limit length
        return '\n'.join(text_parts)

    return ''


def extract_tool_calls(content: Any, message_uuid: str, timestamp: str) -> List[Dict[str, Any]]:
    """Extract tool use blocks from message content"""
    tool_calls = []

    if not isinstance(content, list):
        return tool_calls

    for block in content:
        if isinstance(block, dict) and block.get('type') == 'tool_use':
            tool_calls.append({
                'message_uuid': message_uuid,
                'tool_use_id': block.get('id'),
                'tool_name': block.get('name'),
                'tool_input': json.dumps(block.get('input', {})),
                'tool_result': None,
                'is_error': False,
                'timestamp': timestamp
            })

    return tool_calls


def extract_tool_results(content: Any, tool_calls_map: Dict[str, Dict]) -> None:
    """Extract tool results and match them to tool calls"""
    if not isinstance(content, list):
        return

    for block in content:
        if isinstance(block, dict) and block.get('type') == 'tool_result':
            tool_use_id = block.get('tool_use_id')
            if tool_use_id and tool_use_id in tool_calls_map:
                result_content = block.get('content', '')
                if isinstance(result_content, str):
                    tool_calls_map[tool_use_id]['tool_result'] = result_content[:5000]  # Limit length
                tool_calls_map[tool_use_id]['is_error'] = block.get('is_error', False)


def extract_message_data(obj: Dict[str, Any], session_id: str) -> Optional[Dict[str, Any]]:
    """Extract relevant fields from a message object"""
    msg_type = obj.get('type')

    # Only process user and assistant messages
    if msg_type not in ['user', 'assistant']:
        return None

    message = obj.get('message', {})
    role = message.get('role')

    if not role or role not in ['user', 'assistant']:
        return None

    content = message.get('content', '')
    has_tool_use = False
    has_thinking = False

    # Check for tool use and thinking blocks
    if isinstance(content, list):
        for block in content:
            if isinstance(block, dict):
                if block.get('type') == 'tool_use':
                    has_tool_use = True
                elif block.get('type') == 'thinking':
                    has_thinking = True

    # Extract usage metrics
    usage = message.get('usage', {})
    input_tokens = usage.get('input_tokens', 0)
    output_tokens = usage.get('output_tokens', 0)

    return {
        'session_id': session_id,
        'message_uuid': obj.get('uuid'),
        'parent_uuid': obj.get('parentUuid'),
        'role': role,
        'content_text': extract_text_content(content),
        'timestamp': obj.get('timestamp'),
        'prompt_id': obj.get('promptId'),
        'model': message.get('model'),
        'input_tokens': input_tokens,
        'output_tokens': output_tokens,
        'has_tool_use': has_tool_use,
        'has_thinking': has_thinking,
        'stop_reason': message.get('stop_reason')
    }


def calculate_conversation_stats(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
    """Calculate aggregate statistics for a conversation"""
    stats = {
        'message_count': 0,
        'user_message_count': 0,
        'assistant_message_count': 0,
        'total_input_tokens': 0,
        'total_output_tokens': 0,
        'started_at': None,
        'last_updated_at': None,
        'model': None
    }

    for msg in messages:
        if msg.get('role') == 'user':
            stats['user_message_count'] += 1
        elif msg.get('role') == 'assistant':
            stats['assistant_message_count'] += 1
            if not stats['model'] and msg.get('model'):
                stats['model'] = msg.get('model')

        stats['message_count'] += 1
        stats['total_input_tokens'] += msg.get('input_tokens', 0)
        stats['total_output_tokens'] += msg.get('output_tokens', 0)

        timestamp = msg.get('timestamp')
        if timestamp:
            if not stats['started_at'] or timestamp < stats['started_at']:
                stats['started_at'] = timestamp
            if not stats['last_updated_at'] or timestamp > stats['last_updated_at']:
                stats['last_updated_at'] = timestamp

    return stats


def parse_conversation_file(file_path: str, start_line: int = 0) -> Dict[str, Any]:
    """Parse a conversation file and extract all data"""
    raw_messages = parse_jsonl_file(file_path, start_line)

    if not raw_messages:
        return None

    # Extract session metadata from first message
    session_id = None
    project_path = None
    git_branch = None
    slug = None

    for obj in raw_messages:
        if obj.get('sessionId'):
            session_id = obj['sessionId']
        if obj.get('cwd'):
            project_path = obj['cwd']
        if obj.get('gitBranch'):
            git_branch = obj['gitBranch']
        if obj.get('slug'):
            slug = obj['slug']

        if session_id and project_path:
            break

    if not session_id:
        logger.warning(f"No session ID found in {file_path}")
        return None

    # Extract messages
    messages = []
    tool_calls_map = {}  # Map tool_use_id to tool call data

    for obj in raw_messages:
        msg_data = extract_message_data(obj, session_id)
        if msg_data:
            messages.append(msg_data)

            # Extract tool calls from assistant messages
            if msg_data['role'] == 'assistant' and msg_data['has_tool_use']:
                message_content = obj.get('message', {}).get('content', [])
                tool_calls = extract_tool_calls(message_content, msg_data['message_uuid'], msg_data['timestamp'])
                for tc in tool_calls:
                    tool_calls_map[tc['tool_use_id']] = tc

        # Extract tool results from user messages
        if obj.get('type') == 'user':
            message_content = obj.get('message', {}).get('content', [])
            extract_tool_results(message_content, tool_calls_map)

    if not messages:
        return None

    # Calculate stats
    stats = calculate_conversation_stats(messages)

    return {
        'session_id': session_id,
        'project_path': project_path or 'unknown',
        'git_branch': git_branch,
        'slug': slug,
        'file_path': file_path,
        'messages': messages,
        'tool_calls': list(tool_calls_map.values()),
        **stats
    }